# ESA Climate Toolbox imports for accessing and plotting ESA CCI data
from esa_climate_toolbox.core import get_op              # Get predefined operations (e.g., time series, averages)
from esa_climate_toolbox.core import list_ecv_datasets   # List available datasets per ECV (Essential Climate Variable)
from esa_climate_toolbox.core import get_store           # Connect to an ESA data store
from esa_climate_toolbox.core import list_datasets       # List all datasets in a store
from esa_climate_toolbox.core import open_data           # Open datasets with the toolbox
from esa_climate_toolbox.ops import plot                 # High-level plotting functions for CCI data
from esa_climate_toolbox.ops import temporal_aggregation # Converting data into coarser time resolution based

# Used for connecting to remote data sources (e.g., ESA CCI ODP)
from xcube.core.store import new_data_store

# For displaying geographic regions interactively (if desired)
from IPython.display import GeoJSON
import shapely.geometry   # Handling geometric objects like bounding boxes
import folium

# Core data science libraries
import pandas as pd       # For tabular data handling and time manipulation
import numpy as np        # For numerical operations
import xarray as xr       # For multi-dimensional climate data structures

# Mapping and plotting
import matplotlib.pyplot as plt       # Plotting library
from matplotlib.colors import TwoSlopeNorm
import cartopy.crs as ccrs            # Cartographic projections for spatial data
import cartopy.feature as cfeature
import rasterio                       # Used to read GeoTIFF files

# Individual colormap creation
from matplotlib.colors import ListedColormap, BoundaryNorm

# Notebook settings
import warnings
warnings.filterwarnings("ignore") # Suppressing warnings to keep notebook output clean
%matplotlib inline

list_ecv_datasets('Fire')

fire_dataset = 'ESACCI-L4_FIRE-BA-MODIS-2001-2022-fv5.1.zarr'
cci_zarr_store = new_data_store("esa-cci-zarr")
cci_zarr_store.describe_data(fire_dataset)

# open the dataset with the toolbox command open_data from the esa-climate-toolbox
fire_ds, fire_ds_name = open_data(fire_dataset, var_names="burned_area")
fire_ds

# Set bounding box by defining the longitude minimum, latitude minimum, longitude maximum and latitude maximum
bbox_iberianpeninsula = (-10, 35.0, 5, 44.0) # Iberian Peninsula
bbox_west = (-9.5, 39.0, -6.0, 44.0) # North Western part of the Iberian Peninsula 

bbox = bbox_west

# Display bounding box for reference 
geom = shapely.geometry.box(*bbox)
centroid = geom.centroid

m = folium.Map(location=[centroid.y, centroid.x], zoom_start=7)

folium.GeoJson(geom.__geo_interface__).add_to(m)

m

# Building the subset and saving it into a new dataset called fire_sub
subset_spatial_op = get_op('subset_spatial')
fire_sub = subset_spatial_op(ds=fire_ds, region=bbox_west)

import matplotlib as mpl
mpl.rcParams['figure.dpi'] = 150  # here we can increase the resolution of the figures to make them look clearer

# Get the operator plot_map from the toolbox:
plot_map = get_op('plot_map')

# Define the months to plot
dates = ['2017-06-01', '2017-10-01']

# Here is the actual plotting with the toolbox for June
plot_map(
    fire_sub.sel(time=dates[0], method='nearest'), # this is the dataset we want to plot, we chose the date from above via the for loop
    var="burned_area", # here we define the variable we would like to plot
    projection="PlateCarree", # here we define the projection for the plot
    region=bbox_west, # the bbox is used here to define the x and y limits which makes it easy as we already defined them above
    properties=dict(cmap="gist_heat_r", vmin=0, vmax=1e8, xticks=np.arange(bbox[0], bbox[2], 1), yticks=np.arange(bbox[1], bbox[3], 1)), # here you can change the colormap and the min and max of the colorbar
)

# and next we plot October - the dates are taken from the dates we defined above
plot_map(
     fire_sub.sel(time=dates[1], method='nearest'), # this is the dataset we want to plot, we chose the date from above via the for loop
     var="burned_area", # here we define the variable we would like to plot
     projection="PlateCarree", # here we define the projection for the plot
     region=bbox_west, # the bbox is used here to define the x and y limits which makes it easy as we already defined them above
     properties=dict(cmap="gist_heat_r", vmin=0, vmax=1e8, xticks=np.arange(bbox[0], bbox[2], 1), yticks=np.arange(bbox[1], bbox[3], 1)), # here you can change the colormap and the min and max of the colorbar
)

soilmoisture_dataset = 'ESACCI-SOILMOISTURE-L3S-SSMV-COMBINED-1978-2021-fv07.1.zarr'
cci_zarr_store = new_data_store('esa-cci-zarr')
cci_zarr_store.describe_data(soilmoisture_dataset)

# open the dataset with the toolbox command open_data from the esa-climate-toolbox
sm_ds, sm_ds_name = open_data(soilmoisture_dataset, var_names="sm")

# now we can directly build a subset of the dataset and will use the same a for the fire dataset
subset_spatial_op = get_op('subset_spatial')
sm_sub = subset_spatial_op(ds=sm_ds, region=bbox)
sm_sub

# get the function from the esa-climate-toolbox
from esa_climate_toolbox.ops import temporal_aggregation

# Create monthly means
sm_monthly_mean = temporal_aggregation(sm_sub, 
                                method='mean', # here we can define what to do with the data, e.g. 'sum', 'mean'
                                period='M') # 'M' = month, other inputs are e.g. 'Y' = year, 'QS-DEC' = quarterly 
display(sm_monthly_mean)

# get the climatology operation
climatology_op = get_op('climatology')
sm_monthly_clim = climatology_op(ds=sm_monthly_mean)
sm_monthly_clim

# Define the month to plot
date = '2017-10-01'
climatology_time = f"1978-{date[5:7]}-01"

# calculate the difference of 2017 to the mean soil moisture:
sm_plotting = (sm_monthly_mean["sm"].sel(time=date, method='nearest') - \
       sm_monthly_clim["sm"].sel(time=climatology_time, method='nearest')).compute()

# Do the plotting
plot_map(
    sm_plotting.to_dataset(name="sm"), 
    var="sm", 
    title='Soil Moisture Anomaly October 2017 vs 1978–2021 mean',
    projection="PlateCarree", 
    region=bbox, 
    properties=dict(cmap='BrBG', 
                    vmin=-0.1, 
                    vmax=0.1,
                    cbar_kwargs={"label": "Soil Moisture Anomaly (m³/m³)"},
                    xticks=np.arange(bbox[0], bbox[2], 1), 
                    yticks=np.arange(bbox[1], bbox[3], 1)
                   ), 
)
plt.show()

# Type your code here:
lst_dataset = 'ESACCI-LST-L3S-LST-IRCDR_-0.01deg_1MONTHLY_DAY-1995-2020-fv3.00.zarr'

# Solution to Exercise 1

# load the dataset
lst_ds, lst_name = open_data(lst_dataset, var_names="lst")
# get the spatial subset
subset_spatial_op = get_op('subset_spatial')
lst_sub = subset_spatial_op(ds=lst_ds, region=bbox_west)
# get the climatology operation
climatology_op = get_op('climatology')
lst_clim = climatology_op(ds=lst_sub)

# Define the two months to plot
dates = ['2017-09-01', '2017-10-01']
norm = TwoSlopeNorm(vmin=-5, vcenter=0, vmax=10)
# Compute the anomaly for June and plot
climatology_time = f"1995-{dates[0][5:7]}-01"
data_june = (lst_sub["lst"].sel(time=dates[0], method='nearest') -
             lst_clim["lst"].sel(time=climatology_time, method='nearest')).compute().to_dataset(name="lst")
plot_map(
    data_june,
    var="lst",
    title=str(dates[0]),
    projection="PlateCarree",
    region=bbox,
    properties=dict(cmap="bwr", 
                    norm=norm,
                    cbar_kwargs={"label": 'Land Surface Temperature Anomaly (K)'},
                    xticks=np.arange(bbox[0], bbox[2], 1), 
                    yticks=np.arange(bbox[1], bbox[3], 1)
    )
)
plt.show()

# Compute the anomaly for October and plot
climatology_time = f"1995-{dates[1][5:7]}-01"
data_october = (lst_sub["lst"].sel(time=dates[1], method='nearest') -
                lst_clim["lst"].sel(time=climatology_time, method='nearest')).compute().to_dataset(name="lst")
plot_map(
    data_october,
    var="lst",
    title=str(dates[1]),
    projection="PlateCarree",
    region=bbox,
    properties=dict(cmap="bwr", 
                    norm=norm,
                    cbar_kwargs={"label": 'Land Surface Temperature Anomaly (K)'},
                    xticks=np.arange(bbox[0], bbox[2], 1), 
                    yticks=np.arange(bbox[1], bbox[3], 1)
    )
)
plt.show()

# list_ecv_datasets('Biomass')

# Type your code here ...
biomass_dataset_100m = 'ESACCI-BIOMASS-L4-AGB-MERGED-100m-2010-2020-fv4.0.zarr'
bbox_small = (-9.0, 39.5, -7.0, 41.0) # smaller area

# Solution to Exercise 2

biomass_dataset_100m = 'ESACCI-BIOMASS-L4-AGB-MERGED-100m-2010-2020-fv4.0.zarr'
biomass_100m, biomass_name = open_data(biomass_dataset_100m,var_names="agb")
bbox_small = (-9.0, 39.5, -7.0, 41.0) # smaller area

subset_spatial_op = get_op('subset_spatial')

biomass_sub = subset_spatial_op(ds=biomass_100m, region=bbox_small)


# using the toolbox operator "diff" to calcualte the difference in agb:
from esa_climate_toolbox.ops import diff
biomass_diff = diff(
    ds = biomass_sub,
    ds2 = biomass_sub.sel(time='2017-1-1',method='nearest')
)

#
plot_map(
    biomass_diff.sel(time='2018-01-01', method='nearest'), 
    var="agb",
    projection="PlateCarree", 
    region=bbox_small, 
    properties=dict(cmap="PiYG", vmin=-100, vmax=100, xticks=np.arange(bbox_small[0], bbox_small[2], 0.2), yticks=np.arange(bbox_small[1], bbox_small[3], 0.2))
)
plt.show()

# Type your code here: 
co_dataset = 'esacci.PREC.mon.L3S.CO.multi-sensor.multi-platform.IASI_MOPITT_MERGED_LATMOS-180x360_1M.1-0.r1'

# Solution to Exercise 3

from matplotlib.colors import LogNorm

co_dataset = 'esacci.PREC.mon.L3S.CO.multi-sensor.multi-platform.IASI_MOPITT_MERGED_LATMOS-180x360_1M.1-0.r1'

co_ds, co_name = open_data(co_dataset,var_names="co_total_column_day")



plot_map(
    co_ds.sel(time='2017-10-16', method='nearest'), 
    var="co_total_column_day",
    projection="PlateCarree", 
    region=bbox_iberianpeninsula,
    properties=dict(cmap="RdYlBu_r", vmin=1e18, vmax=2e18, 
                    xticks=np.arange(bbox_iberianpeninsula[0], bbox_iberianpeninsula[2], 1), 
                    yticks=np.arange(bbox_iberianpeninsula[1], bbox_iberianpeninsula[3], 1))
)

# Here you see the first steps to open the dataset... note that the dataset is rather large and this takes a while

# fire_pixel_dataset = "esacci.FIRE.mon.L3S.BA.MODIS.Terra.MODIS_TERRA.v5-1.pixel"

# cci_store = new_data_store('esa-cci')
# #cci_store.describe_data(fire_pixel_dataset)

# places = ["AREA_3"]

# # now we use the data file name and the places we defined to open the data tree:
# fire_dt = cci_store.open_data(
#     fire_pixel_dataset,
#     place_names=places
# )
# fire_dt