calc_afterglow.py

"""
Project ACSAF: Aerosol & Cloud geometry based Sunset/Sunrise cloud Afterglow Forecaster

This script uses the ECMWF AIFS cloud cover data and CAMS AOD550 data to visualize the cloud cover maps and calculate various parameters related to afterglow.
"""
import xarray as xr
import numpy as np
import pandas as pd
import datetime
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from astral.sun import sun, elevation, azimuth
from astral import LocationInfo
import pytz
import cfgrib
from calc_cloudbase import specific_to_relative_humidity, calc_cloud_base
import logging
logging.basicConfig(
    level=logging.INFO,
    datefmt= '%Y-%m-%d %H:%M:%S',
                    )

run = "00"
run = run.zfill(2)
today = datetime.date.today() #- datetime.timedelta(days=1)
today_str = today.strftime("%Y%m%d")

# Define the latitude and longitude for the region of interest
lat, lon = 51.5, -1.0  #Reading
city = LocationInfo("Reading", "England")
s_tdy = sun(city.observer, date=today)
sunset = s_tdy['sunset']
s_tmr = sun(city.observer, date=today + datetime.timedelta(days=1))
sunrise = s_tmr['sunrise']
sunset_tmr = s_tmr['sunset']

def max_solar_elevation(city, date):
    tz = pytz.timezone(city.timezone)
    observer = city.observer

    # Sample every 5 minutes across the day
    times = [datetime.datetime.combine(date, datetime.datetime.min.time()) + datetime.timedelta(minutes=5*i) for i in range(288)]
    times = [tz.localize(t) for t in times]

    max_angle = max(elevation(observer, t) for t in times)
    return max_angle

max_elev = max_solar_elevation(city, datetime.date.today())
logging.info(f"Maximum solar elevation in {city.name}: {max_elev:.2f}°")

#function to get the azuimuth angle at sunset
def get__sunset_azimuth(city, date):
    tz = pytz.timezone(city.timezone)
    observer = city.observer

    # Get the sunset time
    s = sun(observer, date=date, tzinfo=tz)
    sunset_time = s['sunset']

    # Calculate the azimuth angle at sunset
    azimuth_angle = azimuth(observer, sunset_time)

    return azimuth_angle

sunset_azimuth = get__sunset_azimuth(city, today)
sunset_azimuth_42 = get__sunset_azimuth(city, today + datetime.timedelta(days=1))
logging.info(f"Sunset azimuth angle in {city.name}: {sunset_azimuth:.2f}°")
logging.info(f"sunset time in {city.name}: {sunset}")

def plot_azimuth_line(ax, lon, lat, azimuth_deg, length=5.0):
    """
    Draw a line from (lon, lat) in the direction of azimuth_deg.
    
    Parameters:
    - ax: Matplotlib axis
    - lon, lat: starting point (e.g., city location)
    - azimuth_deg: direction in degrees (0=N, 90=E, etc.)
    - length: length of the line in degrees
    """
    # Convert azimuth to radians
    azimuth_rad = np.deg2rad(azimuth_deg)

    # Estimate the end point of the line
    dx = length * np.sin(azimuth_rad)
    dy = length * np.cos(azimuth_rad)

    end_lon = lon + dx
    end_lat = lat + dy

    # Plot the line
    ax.plot([lon, end_lon], [lat, end_lat], c='orange', lw=2)

def safe_distance(h):
    # Equation to calculate the safe distance in km
    l1 = 2*np.arccos(6371/((6371+h)))*360*6371*2*np.pi
    return l1

def calc_afterglow_time(h, max_elev, l1):
    # Equation to calculate the time of afterglow in km and minutes
    t1 = np.arccos(6371/((6371+h)/360))*24*60/np.sin(max_elev.radians())
    t2 = l1/(6371*2*np.pi*24*60/np.sin(max_elev.radians()))
    z = t1 + t2
    return z

def calc_cloud_hole_size(z):
    # Equation to calculate the cloud hole size in km
    q = (z/1440)*40076

# Define a square box enclosing the region of interest
lat_min, lat_max = lat - 3, lat + 3
lon_min, lon_max = lon - 5, lon + 1

ds_18 = xr.open_dataset(f'input/{today_str}{run}0000-18h-oper-fc.grib2', engine = 'cfgrib')
ds_42 = xr.open_dataset(f'input/{today_str}{run}0000-42h-oper-fc.grib2', engine = 'cfgrib')

# t2m at 2m
ds_18_2m = cfgrib.open_dataset(f'input/{today_str}{run}0000-18h-oper-fc.grib2', filter_by_keys={'typeOfLevel': 'heightAboveGround', 'level': 2})
ds_42_2m = cfgrib.open_dataset(f'input/{today_str}{run}0000-42h-oper-fc.grib2', filter_by_keys={'typeOfLevel': 'heightAboveGround', 'level': 2})

def extract_variable(ds, var_name, lat_min, lat_max, lon_min, lon_max, verbose=False):
    var = getattr(ds, var_name)
    var = var.sel(latitude=slice(lat_max, lat_min), longitude=slice(lon_min, lon_max)).squeeze()
    if verbose:
        logging.info(f"{var_name}:")
        logging.info(var.values)
    return var

ds_18_lcc = extract_variable(ds_18, "lcc", lat_min, lat_max, lon_min, lon_max)
ds_18_mcc = extract_variable(ds_18, "mcc", lat_min, lat_max, lon_min, lon_max)
ds_18_hcc = extract_variable(ds_18, "hcc", lat_min, lat_max, lon_min, lon_max)
# tcc (total cloud cover) removed from processing; only use layered covers
ds_42_lcc = extract_variable(ds_42, "lcc", lat_min, lat_max, lon_min, lon_max)
ds_42_mcc = extract_variable(ds_42, "mcc", lat_min, lat_max, lon_min, lon_max)
ds_42_hcc = extract_variable(ds_42, "hcc", lat_min, lat_max, lon_min, lon_max)

cloud_vars_18 = {
    "lcc": ds_18_lcc,
    "mcc": ds_18_mcc,
    "hcc": ds_18_hcc
}

cloud_vars_42 = {
    "lcc": ds_42_lcc,
    "mcc": ds_42_mcc,
    "hcc": ds_42_hcc
}


def plot_cloud_cover_map(data_dict, city, lon, lat, title_prefix, fcst_hr, sunset_azimuth, save_path='output', cmap='gray'):
    """
    Plot cloud cover maps: LCC, MCC, HCC.

    Parameters:
    - data_dict: dict with keys "lcc", "mcc", "hcc" and values as xarray DataArrays
    - city: LocationInfo object (from astral)
    - lon, lat: coordinates of the city
    - title_prefix: string to prefix plot titles
    - save_path: output file path
    - cmap: colormap to use (default: 'gray')
    """
    fig, axs = plt.subplots(2, 2, figsize=(10, 8))
    var_order = [
        ("lcc", "Low Cloud Cover"),
        ("mcc", "Mid Cloud Cover"),
        ("hcc", "High Cloud Cover"),
    ]

    mappables = []
    for ax, (key, label) in zip(axs.flat, var_order):
        if key in data_dict:
            mappable = data_dict[key].plot(ax=ax, cmap=cmap, add_colorbar=False, vmin=0, vmax=100)
            mappables.append(mappable)
            ax.set_title(f"{label}")
            ax.scatter(lon, lat, color='red', marker='x', label=city.name)
            plot_azimuth_line(ax, lon, lat, sunset_azimuth, length=5.0)
        else:
            ax.set_visible(False)
            
    for ax in axs[0]:  # second row
        pos = ax.get_position()
        ax.set_position([pos.x0, pos.y0 - 0.03, pos.width, pos.height])
            
    for ax in axs[1]:  # second row
        pos = ax.get_position()
        ax.set_position([pos.x0, pos.y0 - 0.06, pos.width, pos.height])
                
    if mappables:
        cbar = fig.colorbar(mappables[0], ax=axs, orientation='vertical', fraction=0.03, pad=0.04)
        cbar.set_label("Cloud Cover (%)")

    fig.suptitle(title_prefix, fontsize=14)
    
    # Shared explanation under the title
    fig.text(0.5, 0.92, f"Red × marks {city.name}", ha='center', fontsize=10, color='red')
    fig.text(0.5, 0.90, f"Orange marks sunset azimuth {round(sunset_azimuth,1)}°", ha='center', fontsize=10, color='darkorange')

    plt.savefig(f'{save_path}' + f"/{today_str}{run}0000-{fcst_hr}h-AIFS_cloud_cover_{city.name}.png")
    plt.close()
    
plot_cloud_cover_map(cloud_vars_18, city, lon, lat,
                     f'{today_str} {run}z +18h EC AIFS cloud cover (today sunset)',
                     '18', sunset_azimuth, save_path='output', cmap='gray')

plot_cloud_cover_map(cloud_vars_42, city, lon, lat,
                     f'{today_str} {run}z +42h EC AIFS cloud cover (tomorrow sunset)',
                     '42', sunset_azimuth, save_path='output', cmap='gray')

RH_18, p_18 = specific_to_relative_humidity(ds_18.q, ds_18.t, ds_18.isobaricInhPa, lat, lon)
cloud_base_lvl_18, z_lcl_18, RH_cb_18 = calc_cloud_base(ds_18_2m["t2m"], ds_18_2m["d2m"], ds_18.t, RH_18, ds_18.isobaricInhPa, lat, lon)

RH_42, p_42 = specific_to_relative_humidity(ds_42.q, ds_42.t, ds_42.isobaricInhPa, lat, lon)
cloud_base_lvl_42, z_lcl_42, RH_cb_42 = calc_cloud_base(ds_42_2m["t2m"], ds_42_2m["d2m"], ds_42.t, RH_42, ds_42.isobaricInhPa, lat, lon)

logging.info('18')
logging.info(cloud_base_lvl_18, z_lcl_18, RH_cb_18)
logging.info('42')
logging.info(cloud_base_lvl_42, z_lcl_42, RH_cb_42)

def azimuth_line_points(lon, lat, azimuth, distance_km, num_points=100):
    """
    Generate a series of points along a given azimuth line (in degrees) starting from the given coordinates.
    
    Parameters:
    - lon: Longitude of the starting point
    - lat: Latitude of the starting point
    - azimuth: Azimuth angle in degrees (0° is north, 90° is east, 180° is south, 270° is west)
    - distance_km: Distance to generate points (in km)
    - num_points: Number of points to generate along the azimuth line
    
    Returns:
    - lons, lats: Arrays of longitudes and latitudes along the azimuth line
    
    Generated using GPT
    
    """
    # Earth radius in km
    earth_radius = 6371.0

    # Convert azimuth to radians
    azimuth_rad = np.deg2rad(azimuth)

    # Generate points along the azimuth line
    dists = np.linspace(0, distance_km, num_points)
    lons = lon + (dists / earth_radius) * np.sin(azimuth_rad) * (180 / np.pi)
    lats = lat + (dists / earth_radius) * np.cos(azimuth_rad) * (180 / np.pi)

    return lons, lats

def extract_cloud_cover_along_azimuth(data_dict, lon, lat, azimuth, distance_km, num_points=20):
    """
    Extract cloud cover data along an azimuth line over a certain distance.
    
    Parameters:
    - data_dict: Dictionary containing cloud cover data ("tcc", "lcc", "mcc", "hcc")
    - lon, lat: Starting coordinates of the city
    - azimuth: Azimuth angle in degrees
    - distance_km: Distance along the azimuth line (in km)
    - num_points: Number of points to sample along the azimuth line
    
    Returns:
    - cloud_cover_data: Dictionary with cloud cover data along the azimuth line
    """
    # Generate azimuth line points
    lons, lats = azimuth_line_points(lon, lat, azimuth, distance_km, num_points)
    
    # Interpolate the cloud cover data along the azimuth line
    cloud_cover_data = np.diag(data_dict.interp(longitude=lons, latitude=lats, method='nearest').values) #diagonal entries are the target coordinates (points along the azimuth line)
    
    return cloud_cover_data

def plot_cloud_cover_along_azimuth(cloud_cover_data, azimuth, distance_km, fcst_hr, threshold, cloud_lvl_used, save_path='output'):
    """
    Plot cloud cover data along the azimuth line.
    
    Parameters:
    - cloud_cover_data: Array containing cloud cover data along the azimuth line
    - azimuth: Azimuth angle in degrees
    - distance_km: Distance along the azimuth line (in km)
    - save_path: Path to save the plot
    """
    fig, ax = plt.subplots(figsize=(10, 6))
    
    # Plot the cloud cover data
    ax.plot(np.linspace(0, distance_km, len(cloud_cover_data)), cloud_cover_data, label=f"Cloud Cover ({cloud_lvl_used})")
    
    # Check for the first distance where cloud cover falls below the threshold
    below_threshold_index = np.argmax(cloud_cover_data <= threshold)
    
    if below_threshold_index > 0:  # Ensure that the threshold is met somewhere in the data
        distance_below_threshold = np.linspace(0, distance_km, len(cloud_cover_data))[below_threshold_index]

        logging.info(f"The cloud cover falls below {threshold}% at {distance_below_threshold} km.")
    else:
        logging.info(f"Cloud cover does not fall below {threshold}%.")
        distance_below_threshold = np.nan
        
    # Plot the threshold line
    ax.axhline(y=threshold, color='r', linestyle='--', label=f'{threshold}% threshold')
    
    if ~np.isnan(distance_below_threshold):
        ax.axvline(x=distance_below_threshold, color='r', linestyle='--', label=f'{threshold}% threshold')
    ax.legend()
    
    ax.set_ylim(0,100)
    ax.set_xlabel('Distance along Azimuth (km)')
    ax.set_ylabel('Cloud Cover (%)')
    ax.set_title(f'{today_str} {run}z +{fcst_hr}h EC AIFS cloud cover Along Azimuth path ')
    #set subtitle
    ax.text(0.5, 1.02, f"Azimuth {azimuth}°", fontsize=10)
    ax.legend()
    
    plt.savefig(save_path + f"/{today_str}{run}0000_{fcst_hr}h_AIFS_cloud_cover_azimuth_{city.name}.png")
    plt.close()
    return distance_below_threshold
    
cloud_cover_data = extract_cloud_cover_along_azimuth(ds_18_tcc, lon, lat, sunset_azimuth, 500, num_points=20)
cloud_cover_data_42 = extract_cloud_cover_along_azimuth(ds_42_tcc, lon, lat, sunset_azimuth_42, 500, num_points=20)

def get_cloud_extent(data_dict, lon, lat, azimuth, cloud_base_lvl: float, fcst_hr, distance_km = 500, num_points=20, threshold=60.0):
    if cloud_base_lvl > 0.0 and cloud_base_lvl <= 2000.0:
        data = data_dict['lcc']
        cloud_lvl_used = 'lcc'
    if cloud_base_lvl > 2000.0 and cloud_base_lvl <= 6000.0:
        data = data_dict['mcc']
        cloud_lvl_used = 'mcc'
    if cloud_base_lvl > 6000.0:
        data = data_dict['hcc']
        cloud_lvl_used = 'hcc'
    elif np.all(np.isnan(cloud_base_lvl)): 
        logging.info(f"Cloud base level {cloud_base_lvl} is invalid. There is no cloud cover. Afterglow not probable.")
        logging.info("Trying to use tcc for cloud cover data")
        data = data_dict['tcc']
        cloud_lvl_used = 'tcc'
        cloud_present = False
        try:
            cloud_cover_data = extract_cloud_cover_along_azimuth(data, lon, lat, azimuth, distance_km, num_points)
            distance_below_threshold = plot_cloud_cover_along_azimuth(cloud_cover_data, azimuth, distance_km, fcst_hr, threshold, cloud_lvl_used, save_path='output')
        except:
            logging.error(f"Cloud cover data is not available for forecast hour {fcst_hr}.")
            cloud_present = False
    # Check if data is associated with a value or a NaN
    if np.all(np.isnan(data)):
        logging.error(f"Cloud cover NaN for forecast hour {fcst_hr}. There is no stratiform cloud cover. Afterglow not probable.")
        cloud_present = False
        return cloud_present
    else:
        cloud_cover_data = extract_cloud_cover_along_azimuth(data, lon, lat, azimuth, distance_km, num_points)
        distance_below_threshold = plot_cloud_cover_along_azimuth(cloud_cover_data, azimuth, distance_km, fcst_hr, threshold, cloud_lvl_used, save_path='output')
    distance_below_threshold = distance_below_threshold * 1000 # convert to meters
    return distance_below_threshold

distance_below_threshold_18 = get_cloud_extent(cloud_vars_18, lon, lat, sunset_azimuth, cloud_base_lvl_18, '18')
distance_below_threshold_42 = get_cloud_extent(cloud_vars_42, lon, lat, sunset_azimuth, cloud_base_lvl_42, '42')

def geom_condition(cloud_base_height, cloud_extent, LCL):
    """
    Calculate the geometric condition for afterglow based on cloud base height and cloud extent.
    
    Parameters:
    - cloud_base_height: Cloud base height (hPa)
    - cloud_extent: Cloud extent (m)
    
    Returns:
    - geom_condition: Geometric condition for afterglow
    """
    # Cosntant
    R = 6371*10**3  # Radius of the Earth in meters
    lf_ma = 2*np.sqrt(2*R*cloud_base_height) 
    try:
        logging.info(f"lf_ma: {round(lf_ma)} m")
        geom_condition_LCL_used = False
    except ValueError as e:
        logging.error(f"Error: {e}")
        logging.info(f"lf_ma is {lf_ma}")
        logging.info('We will assume lf_ma using LCL')
        lf_ma = 2*np.sqrt(2*R*LCL)
        logging.info(f"lf_ma: {round(lf_ma)} m")
        geom_condition_LCL_used = True
        geom_condition = False
    # Compare with cloud_extent
    if lf_ma > cloud_extent:
        geom_condition = True
    else:
        geom_condition = False
    logging.info("cloud geometry condition:", geom_condition)
    return geom_condition, geom_condition_LCL_used, lf_ma

geom_cond_18, geom_condition_LCL_used_18, lf_ma_18 = geom_condition(cloud_base_lvl_18, distance_below_threshold_18, z_lcl_18)
geom_cond_42, geom_condition_LCL_used_42, lf_ma_42 = geom_condition(cloud_base_lvl_42, distance_below_threshold_42, z_lcl_42)

def get_elevation_afterglow(cloud_base_lvl, distance_below_threshold, lf_ma, lcl):
    """
    Calculate the elevation angle for afterglow based on cloud base height and distance below threshold.
    
    Parameters:
    - cloud_base_lvl: Cloud base level (m)
    - distance_below_threshold: Distance below threshold (m)
    - lf_ma: Distance to the cloud base (m)
    - lcl: Lifted condensation level (m)
    
    returns:
    - theta: Elevation angle (radians)

    """
    #Constants
    R = 6371*10**3  # Radius of the Earth in meters
    
    if distance_below_threshold == False or np.isnan(distance_below_threshold):
        logging.info("Cloud cover and elevation estimated using tcc")
        theta = np.arctan((cloud_base_lvl-(lf_ma**2/(2*R)))/lf_ma)
        logging.info(f"Elevation angle: {np.rad2deg(theta)}°")
        return theta
    elif np.isnan(cloud_base_lvl):
        logging.info("Cloud cover and elevation estimated using tcc and LCL")
        theta = np.arctan((lcl-((distance_below_threshold**2)/(2*R)))/distance_below_threshold)
        logging.info(f"Elevation angle: {np.rad2deg(theta)}°")
        return theta
    else:
        theta = np.arctan((cloud_base_lvl-(distance_below_threshold**2/(2*R)))/distance_below_threshold)
        logging.info(f"Elevation angle: {np.rad2deg(theta)}°")
        return theta

theta_18 = get_elevation_afterglow(cloud_base_lvl_18, distance_below_threshold_18, lf_ma_18, z_lcl_18)
theta_42 = get_elevation_afterglow(cloud_base_lvl_42, distance_below_threshold_42, lf_ma_42, z_lcl_42)

def get_afterglow_time(lat, today, distance_below_threshold, lf_ma, cloud_base_lvl, z_lvl):
    
    day_of_year = today.timetuple().tm_yday
    
    R = 6371  # Earth's radius in km
    T = 1440  # minutes in a day

    if np.isnan(cloud_base_lvl):
        cloud_base_lvl = z_lvl
    if lf_ma >= distance_below_threshold:
        # Convert latitude to radians
        phi = np.deg2rad(lat)

        # Approximate solar declination angle (in degrees then radians)
        decl_deg = 23.44 * np.sin(np.deg2rad((360 / 365) * (day_of_year - 81)))
        delta = np.deg2rad(decl_deg)

        # Linear speed of sunrays at given latitude and time of year based on https://doi.org/10.1016/j.asr.2023.08.036
        speed = (2 * np.pi * R * np.cos(phi)) / T * np.cos(delta) #In km/minutes
        
        total_afterglow_time = np.sqrt((R*cloud_base_lvl))/speed #minutes
        
        # Make a straight line between origin and point (total_afterglow_time, lf_ma)
        
        t1 =  -distance_below_threshold * (total_afterglow_time/lf_ma) # rearrange purple y=mx
        t2 =  total_afterglow_time + (-distance_below_threshold)*(2*total_afterglow_time/lf_ma) # rearrange purple y=mx
        
        overhead_afterglow_time = np.abs(t1-t2)
        
        actual_afterglow_time = total_afterglow_time - overhead_afterglow_time
        
        actual_afterglow_time = actual_afterglow_time + ((cloud_base_lvl/np.tan(np.deg2rad(5)))/(21*1000/60)) # Accounting for the clouds in visual contact assuming 5 deg elevation
        
        logging.info(f"Total Afterglow time: {total_afterglow_time} seconds")
        logging.info(f"Overhead Afterglow time: {overhead_afterglow_time} seconds")
        logging.info(f"Actual Afterglow time: {actual_afterglow_time} seconds")
    else:
        logging.info(f"Sun ray extent lf_max is less than cloud extent. No afterglow is possible.")
        actual_afterglow_time = 0
    return actual_afterglow_time

actual_afterglow_time_18 = get_afterglow_time(lat, today, distance_below_threshold_18, lf_ma_18, cloud_base_lvl_18, z_lcl_18)
actual_afterglow_time_42 = get_afterglow_time(lat, today, distance_below_threshold_42, lf_ma_42, cloud_base_lvl_42, z_lcl_42)

# One method is to use Equivalent cloud heigh  = cloud base level - equivalent surface height as highlighted in the paper
# But uncertainty is high and the exact mechanism is not clear
# Here we use a simplier and quicker method, but yet to be verified
# The AOD value directly controls the value of the afterglow liklihood index in the weighted equation

# Incorporate AOD
from calc_aod import calc_aod
dust_aod550, total_aod550, dust_aod550_ratio = calc_aod(run, today_str, city.name) #Array of shape (2,) first is 18h , second is 42h

# Weighted likelihod index
def weighted_likelihood_index(geom_condition, aod, dust_aod_ratio, cloud_base_lvl, max_RH, theta):
    """
    Calculate the weighted likelihood index based on AOD, afterglow time, cloud base level, geometric condition, and cloud extent.
    
    Parameters:
    - geom_condition: Geometric condition for afterglow (True/False)
    - aod: AOD value
    - cloud_base_lvl: Cloud base level (m)
    - max_RH: Maximum relative humidity (%)
    - theta: Elevation angle (radians)
    
    Returns:
    - likelihood_index: Weighted likelihood index for afterglow
    """
    if np.isnan(cloud_base_lvl):
        cloud_base_score = 0
    if cloud_base_lvl <= 2000.0:
        cloud_base_score = 0.2
    elif cloud_base_lvl <= 6000.0:
        cloud_base_score = 0.7
    else:
        cloud_base_score = 1
    
    if aod >= 0 and aod <= 0.3:
        aod_score = 1
    elif aod > 0.3 and aod <= 0.5:
        aod_score = 0.8
    elif aod > 0.5 and aod <= 0.7:
        aod_score = 0.2
    elif aod > 0.7:
        aod_score = 0
    
    if dust_aod_ratio >= 0 and dust_aod_ratio <= 0.3:
        dust_ratio_score = 0.2
    elif dust_aod_ratio > 0.3 and dust_aod_ratio <= 0.4:
        dust_ratio_score = 0.8
    elif dust_aod_ratio > 0.5:
        dust_ratio_score = 1
        
    if theta >= 0 and theta <= np.deg2rad(5):
        theta_score = 0.4
    elif theta > np.deg2rad(5):
        theta_score = 1
    elif theta < 0:
        theta_score = -1
    
    if max_RH <= 70: 
        rh_score = -0.5
    if max_RH > 70 and max_RH <= 80:
        rh_score = 0.2
    if max_RH > 80 and max_RH <= 85:
        rh_score = 0.5
    if max_RH > 85 and max_RH <= 95:
        rh_score = 1
    if max_RH > 95:
        rh_score = 0.2
    
    if geom_condition == True:
        # Constants
        geom_condition_weight = 0.4
        aod_weight = 0.1
        dust_aod_ratio_weight = 0.1
        max_RH_weight = 0.05
        cloud_base_lvl_weight = 0.3
        theta_weight = 0.05
        
    elif geom_condition == False:
        # Constants
        geom_condition_weight = 0.7
        aod_weight = 0.05
        dust_aod_ratio_weight = 0.05
        max_RH_weight = 0.05
        cloud_base_lvl_weight = 0.1
        theta_weight = 0.05
        
    # Normalise cloud_base_lvl to 0 to 1
    cloud_base_score = np.clip(cloud_base_score/9000, 0, 1)
    
    # Binary flag for geom_condition
    geom_flag = 1 if geom_condition else 0

    # Compute weighted index
    likelihood_index = (
        (geom_flag ** 3) * geom_condition_weight +
        (aod_score ** 1)* aod_weight +
        (dust_ratio_score ** 0.8 )* dust_aod_ratio_weight +
        (rh_score ** 1) * max_RH_weight +
        (cloud_base_score ** 2) * cloud_base_lvl_weight +
        (theta_score ** 1 )* theta_weight
)   
    
    likelihood_index = np.clip(likelihood_index, 0, 1)  # Ensure it's between 0 and 1
    
    # Scale to 0-100 and round to whole number
    likelihood_index = round(likelihood_index * 100)
    return likelihood_index

likelihood_index_18 = weighted_likelihood_index(geom_cond_18, total_aod550[0], dust_aod550_ratio[0], cloud_base_lvl_18, RH_cb_18, theta_18)
likelihood_index_42 = weighted_likelihood_index(geom_cond_42, total_aod550[1], dust_aod550_ratio[1], cloud_base_lvl_42, RH_cb_42, theta_42)
logging.info(likelihood_index_18)
logging.info(likelihood_index_42)

def possible_colours(cloud_base_lvl, total_aod_550):
    """
    Determine the possible colours for the afterglow based on cloud base level, inferred cloud cover and AOD_550.
    Cloud cover is inferred through RH of the cloud base level. 
    
    Parameters:
    """
    if np.isnan(cloud_base_lvl):
        color = None
    if cloud_base_lvl <= 2000.0:
        if total_aod_550 <= 0.2:
            color = ('orange-red',)
        else:
            color = ('dirty-orange',)
    elif cloud_base_lvl > 2000.0 and cloud_base_lvl <= 6000.0:
        color = ('orange-yellow', 'orange-red', 'dark-red') 
    elif cloud_base_lvl > 6000.0:
        color = ('golden-yellow', 'orange-red', 'crimson', 'magenta', 'dark-red')
    
    return color

possible_colors_18 = possible_colours(cloud_base_lvl_18, total_aod550[0])
possible_colors_42 = possible_colours(cloud_base_lvl_42, total_aod550[1])
logging.info(f"Possible colors for afterglow 18: {possible_colors_18}")
logging.info(f"Possible colors for afterglow 42: {possible_colors_42}")

# Color fill based on index using the 'plasma' colormap
def color_fill(index):
    # Normalize the index to a value between 0 and 1 for the colormap
    norm = plt.Normalize(vmin=0, vmax=100)
    cmap = plt.get_cmap('magma')  # 
    return cmap(norm(index))

def create_dashboard(index_today, index_tomorrow, city, latitude, longitude, azimuth, afterglow_length_18, afterglow_length_42, possible_colors_18, possible_colors_42):
    fig, ax = plt.subplots(1, 2, figsize=(12, 6))  # Adjust size as needed

    # Set background colour for each subplot
    for a in ax:
        a.set_facecolor('black')  # Each 'a' is an individual axis
        a.set_xticks([])
        a.set_yticks([])
        a.spines['top'].set_visible(False)
        a.spines['right'].set_visible(False)
        a.spines['left'].set_visible(False)
        a.spines['bottom'].set_visible(False)

    # Also set the figure background
    fig.patch.set_facecolor('black')
    
    plt.suptitle("ACSAF: Aerosol and Cloud geometry based Sunset cloud Afterglow Forecaster Dashboard", fontsize=16, weight='bold', color='white')
        
    # Box settings for today and tomorrow
    # Larger, square boxes with index numbers only in bold
    ax[0].text(0.1, 0.5, f"{index_today}", fontsize=40, fontweight='bold', ha='center', va='center', 
            color='white', bbox=dict(facecolor=color_fill(index_today), edgecolor='black', boxstyle='round,pad=2'))

    ax[0].text(1.0, 0.5, f"{index_tomorrow}", fontsize=40, fontweight='bold', ha='right', va='center',
            color='white', bbox=dict(facecolor=color_fill(index_tomorrow), edgecolor='black', boxstyle='round,pad=2'))

    # Add text outside the box for "Today" and "Tomorrow"
    ax[0].text(0.1, 0.1, "Today", fontsize=18, ha='center', va='center', color='white', fontweight='bold')
    ax[0].text(0.9, 0.1, "Tomorrow", fontsize=18, ha='center', va='center', color='white', fontweight='bold')

    # Title for the left subplot
    ax[0].axis('off')  # Turn off axis for this subplot


    sunset_tz = pytz.timezone("Europe/London")  # Desired timezone for Reading (UK)

    # Convert the sunset time from UTC to local time
    sunset_local = sunset.astimezone(sunset_tz)
    
    # Get sunset time of tomorrow
    sunset_local_42 = sunset_tmr.astimezone(sunset_tz)
    
    # Info text on the right
    info_text = (
        f"Today: {today.strftime('%Y-%m-%d')}\n"
        f"City: {city.name}\n"
        f"Latitude: {latitude:.2f}°\n"
        f"Longitude: {longitude:.2f}°\n"
        f"Azimuth: {round(azimuth)}°\n"
        f"Sunset Time: {sunset_local.strftime('%H:%M:%S')} \n"
        f"Afterglow Length: {round(afterglow_length_18)} seconds after sunset\n"
        f"Colors: {', '.join(possible_colors_18) }\n"
    )
    ax[1].text(0.7, 0.6, info_text, fontsize=15, ha='center', va='center', color='black', 
           bbox=dict(facecolor='lightgray', alpha=0.8))
    
    # Info text on the right
    info_text = (
        f"Tomorrow: {(today + datetime.timedelta(days=1)).strftime('%Y-%m-%d')}\n"
        f"Sunset Time: {sunset_local_42.strftime('%H:%M:%S')} \n"
        f"Afterglow Length: {round(afterglow_length_42)} seconds after sunset\n"
        f"Colors: {', '.join(possible_colors_42)}\n"
    )
    ax[1].text(0.7, 0.2, info_text, fontsize=15, ha='center', va='center', color='black', 
           bbox=dict(facecolor='lightgray', alpha=0.8))

    ax[1].axis('off')  # Turn off axis for this subplot
        
    # Title
    fig.text(0.01, 0.01, "Index ranges 0-100. Higher indicates more favourable cloud afterglow conditions for occurence and intense colors. \n"
             "Based on daily 00z ECMWF AIFS and CAMS forecast. Generate on daily morning. See supplementary figures for details.",
         ha='left', va='bottom', color='white', fontsize=8)
    
    plt.savefig(f'output/{today_str}{run}0000_afterglow_dashboard_{city.name}.png', dpi=400)

create_dashboard(
    likelihood_index_18, likelihood_index_42, city, lat, lon, sunset_azimuth, actual_afterglow_time_18, actual_afterglow_time_42, possible_colors_18, possible_colors_42 
)

# Work on the case with ecloud extent too large
# The case where cloud cover increases midway along the azimuth path?