import xarray as xr
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import numpy as np
import os

def plot_esa_cci_land_cover(esa_cci_1km_path, lat_range=None, lon_range=None, region_name="", output_path=None):
    # Load the data
    df_esa1km = xr.open_dataset(esa_cci_1km_path)

    # Extracting the global data
    data = df_esa1km.lccs_class[0, :, :]

    # Default color map
    default_color_map = {
        0: '#ffff64', 10: '#ffff64', 11: '#ffff00', 12: '#aaf0f0', 20: '#dcf064', 30: '#c8c864', 40: '#006400',
        50: '#00a000', 60: '#00a000', 61: '#aac800', 62: '#003c00', 70: '#003c00', 71: '#005000', 72: '#285000',
        80: '#285000', 81: '#286400', 82: '#788200', 90: '#8ca000', 100: '#be9600', 110: '#966400', 120: '#966400',
        121: '#966400', 122: '#ffb432', 130: '#ffdcd2', 140: '#ffebaf', 150: '#ffc864', 151: '#ffd278', 152: '#ffebaf',
        153: '#00785a', 160: '#009678', 170: '#00dc82', 180: '#c31400', 190: '#fff5d7', 200: '#dcdcdc', 201: '#fff5d7',
        202: '#0046c8', 210: '#ffffff', 220: '#ffffff'
    }

    # New class names
    esa_class_names = {
        0: 'no.data', 10: 'cropland.rainfed', 11: 'cropland.rainfed.herbaceous.cover', 12: 'cropland.rainfed.tree.or.shrub.cover',
        20: 'cropland.irrigated', 30: 'mosaic.cropland', 40: 'mosaic.natural.vegetation', 50: 'tree.broadleaved.evergreen.closed.to.open',
        60: 'tree.broadleaved.deciduous.closed.to.open', 61: 'tree.broadleaved.deciduous.closed', 62: 'tree.broadleaved.deciduous.open',
        70: 'tree.needleleaved.evergreen.closed.to.open', 71: 'tree.needleleaved.evergreen.closed', 72: 'tree.needleleaved.evergreen.open',
        80: 'tree.needleleaved.deciduous.closed.to.open', 81: 'tree.needleleaved.deciduous.closed', 82: 'tree.needleleaved.deciduous.open',
        90: 'tree.mixed', 100: 'mosaic.tree.and.shrub', 110: 'mosaic.herbaceous', 120: 'shrubland', 121: 'shrubland.evergreen',
        122: 'shrubland.deciduous', 130: 'grassland', 140: 'lichens.and.mosses', 150: 'sparse.vegetation', 151: 'sparse.tree',
        152: 'sparse.shrub', 153: 'sparse.herbaceous', 160: 'tree.cover.flooded.fresh.or.brakish.water', 170: 'tree.cover.flooded.saline.water',
        180: 'shrub.or.herbaceous.cover.flooded', 190: 'urban', 200: 'bare.areas', 201: 'bare.areas.consolidated', 202: 'bare.areas.unconsolidated',
        210: 'water', 220: 'snow.and.ice'
    }

    # Check if flag_values and flag_colors attributes are present
    if hasattr(df_esa1km.lccs_class, 'flag_values') and hasattr(df_esa1km.lccs_class, 'flag_colors'):
        colors = df_esa1km.lccs_class.flag_colors.split(' ')
        flag_vals = df_esa1km.lccs_class.flag_values
        flag_labels = df_esa1km.lccs_class.flag_meanings.split(' ')

        # Creating color values with corresponding labels for the classes present in ESA-CCI data
        esa_cci_lulc_codes = {x: i for x, i in zip(flag_labels, flag_vals)}

        # Handling potential mismatch in lengths
        color_map = {}
        for idx, val in enumerate(flag_vals):
            if idx < len(colors):
                color_map[val] = colors[idx]
            else:
                color_map[val] = "#ffffff"  # Assigning white color if out of range
    else:
        # Use default color map if flag_values or flag_colors are not present
        color_map = default_color_map

    # Determine the extent of the plot based on user input
    if lat_range and lon_range:
        data_region = data.sel(lat=slice(lat_range[0], lat_range[1]), lon=slice(lon_range[0], lon_range[1]))
        masked_data_region = np.ma.masked_array(data_region, mask=(data_region == 0))
        extent_region = [lon_range[0], lon_range[1], lat_range[0], lat_range[1]]
        region_suffix = f"_{lat_range[0]}_{lat_range[1]}_{lon_range[0]}_{lon_range[1]}"
    else:
        masked_data_region = np.ma.masked_array(data, mask=(data == 0))  # Assuming 0 is the code for water or no data
        extent_region = [-180, 180, -90, 90]
        region_suffix = ""

    # Filter out masked values and NaN values
    unique_vals = [val for val in np.unique(masked_data_region) if val is not np.ma.masked and not np.isnan(val)]

    # Create the figure with subplots
    fig = plt.figure(figsize=(15, 20))
    gs = fig.add_gridspec(3, 2, height_ratios=[2, 1, 1], width_ratios=[2, 1])

    # Plotting the regional data
    ax1 = fig.add_subplot(gs[0, :])
    cmap = mcolors.ListedColormap([color_map[val] for val in unique_vals if val in color_map])
    norm = mcolors.BoundaryNorm(unique_vals, cmap.N)
    
    ax1.imshow(masked_data_region[::-1], cmap=cmap, norm=norm, extent=extent_region)  # Flipping the data vertically
    ax1.set_title(f'ESA-CCI Land Cover Classification - {region_name}')
    ax1.set_xlabel('Longitude')
    ax1.set_ylabel('Latitude')
    
    # Plotting the global data
    ax2 = fig.add_subplot(gs[1:, 0])
    masked_data_global = np.ma.masked_array(data, mask=(data == 0))  # Full global data
    unique_vals_global = [val for val in np.unique(masked_data_global) if val is not np.ma.masked and not np.isnan(val)]
    
    cmap_global = mcolors.ListedColormap([color_map[val] for val in unique_vals_global if val in color_map])
    norm_global = mcolors.BoundaryNorm(unique_vals_global, cmap_global.N)
    im = ax2.imshow(masked_data_global[::-1], cmap=cmap_global, norm=norm_global, extent=[-180, 180, -90, 90])
    ax2.set_title('ESA-CCI Land Cover Classification - Global')
    ax2.set_xlabel('Longitude')
    ax2.set_ylabel('Latitude')

    # Draw a rectangle to indicate the region
    if lat_range and lon_range:
        rect = plt.Rectangle((lon_range[0], lat_range[0]), lon_range[1] - lon_range[0], lat_range[1] - lat_range[0], 
                             linewidth=2, edgecolor='r', facecolor='none')
        ax2.add_patch(rect)
    
    # Creating a custom legend with multiple columns
    labels_global = [esa_class_names[val] for val in unique_vals_global if val in esa_class_names]
    colors_global = [color_map[val] for val in unique_vals_global if val in color_map]
    handles_global = [plt.Line2D([0, 1], [0, 0], color=color, linewidth=3) for color in colors_global]
    
    # Placing the legend on the right with multiple columns
    ax3 = fig.add_subplot(gs[1:, 1])
    ax3.axis('off')
    ax3.legend(handles_global, labels_global, loc='center', ncol=1)

    plt.tight_layout()

    # Save the combined plot if output_path is provided
    base_name = os.path.splitext(os.path.basename(esa_cci_1km_path))[0]
    file_name = os.path.join(output_path, f"{base_name}{region_suffix}_combined.png")
    # Check if the file exists and delete it
    if os.path.exists(file_name):
        os.remove(file_name)
    plt.savefig(file_name)


    # Save the regional plot separately
    fig_region, ax_region = plt.subplots(figsize=(15, 10))
    ax_region.imshow(masked_data_region[::-1], cmap=cmap, norm=norm, extent=extent_region)
    ax_region.set_title(f'ESA-CCI Land Cover Classification - {region_name}')
    ax_region.set_xlabel('Longitude')
    ax_region.set_ylabel('Latitude')
    plt.tight_layout()
    file_name = os.path.join(output_path, f"{base_name}{region_suffix}_region.png")
    # Check if the file exists and delete it
    if os.path.exists(file_name):
        os.remove(file_name)
    plt.savefig(file_name)


    # Save the global plot separately with the indication rectangle
    fig_global, ax_global = plt.subplots(figsize=(15, 10))
    im_global = ax_global.imshow(masked_data_global[::-1], cmap=cmap_global, norm=norm_global, extent=[-180, 180, -90, 90])
    ax_global.set_title('ESA-CCI Land Cover Classification - Global')
    ax_global.set_xlabel('Longitude')
    ax_global.set_ylabel('Latitude')
    
    if lat_range and lon_range:
        rect = plt.Rectangle((lon_range[0], lat_range[0]), lon_range[1] - lon_range[0], lat_range[1] - lat_range[0], 
                             linewidth=2, edgecolor='r', facecolor='none')
        ax_global.add_patch(rect)
    
    plt.tight_layout()
    file_name = os.path.join(output_path, f"{base_name}_global.png")
    # Check if the file exists and delete it
    if os.path.exists(file_name):
        os.remove(file_name)
    plt.savefig(file_name)

    plt.close(fig)
    plt.close(fig_region)
    plt.close(fig_global)


if __name__ == "__main__":
    # Define the regions and their parameters
    regions = [
        {
            "lat_range": [12.5, 14.5],
            "lon_range": [13, 15.5],
            "region_name": "Lake Chad",
            "output_path": "/gpfs/projects/bsc32/cerise/outputs/bsc032013/plots/ESA_CCI_1KM"
        },
        {
            "lat_range": [22, 24.5],
            "lon_range": [29, 32],
            "region_name": "Toschka Lakes",
            "output_path": "/gpfs/projects/bsc32/cerise/outputs/bsc032013/plots/ESA_CCI_1KM"
        },
        {
            "lat_range": [36, 44],
            "lon_range": [-10, 5],
            "region_name": "Iberian Peninsula",
            "output_path": "/gpfs/projects/bsc32/cerise/outputs/bsc032013/plots/ESA_CCI_1KM"
        }
    ]

    # Loop through each region and each year
    for region in regions:
        for year in range(1994, 2020):
            esa_cci_1km_path = f"/gpfs/projects/bsc32/cerise/outputs/bsc032013/0.0089_deg/lccsclass_{year}_0.0089deg.nc"
            region_name_with_year = f"{region['region_name']} {year}"
            plot_esa_cci_land_cover(
                esa_cci_1km_path=esa_cci_1km_path, 
                lat_range=region["lat_range"], 
                lon_range=region["lon_range"], 
                region_name=region_name_with_year, 
                output_path=region["output_path"]
            )