#!pip install cdsapi #Uncomment this part if cdsapi is not installed

#!pip install cdo  #Uncomment this part if cdo is not installed

#!pip install cartopy #Uncomment this part if cartopy is not installed

import cdsapi
URL = 'https://cds.climate.copernicus.eu/api'
KEY = '#####' 
client = cdsapi.Client(url=URL,key=KEY)

from cdo import *
cdo = Cdo()
import xarray as xr
import os
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import cartopy

# And some stuff for fancy plotting
from cartopy.mpl.geoaxes import GeoAxes
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
from matplotlib.patches import ConnectionPatch
from mpl_toolkits.axes_grid1 import make_axes_locatable
from mpl_toolkits.axes_grid1.inset_locator import inset_axes, mark_inset

HOME = os.path.expanduser("~")
SCRATCH = os.environ.get("SCRATCH", "/tmp")
WRKDIR = SCRATCH #change this if needed

print(f"Working directory: {WRKDIR}")

DATADIR=os.path.join(WRKDIR, "CERRA")
if not os.path.isdir(DATADIR):
    os.makedirs(DATADIR)
fCERRA1 = os.path.join(DATADIR,"WS_"+"period1"+".grb")
fCERRA2 = os.path.join(DATADIR,"WS_"+"period2"+".grb")
fCERRA = os.path.join(DATADIR, "WS_total_period"+".grb")

DATADIR=os.path.join(WRKDIR,"ERA")
if not os.path.isdir(DATADIR):
    os.makedirs(DATADIR)
fERA1 = os.path.join(DATADIR,"WS_"+"period1"+".nc")
fERA2 = os.path.join(DATADIR,"WS_"+"period2"+".nc")
fERA = os.path.join(DATADIR, "WS_total_period.nc")

print(f"Will be saving data in files: \n{fCERRA}\n{fERA}")

def download_cerra(year, month, days, file):
    dataset = "reanalysis-cerra-single-levels"
    request = {
        "variable": ["10m_wind_speed"],
        "level_type": "surface_or_atmosphere",
        "data_type": ["reanalysis"],
        "product_type": "forecast",
        "year": [year],
        "month": [month],
        "day": days,
        "time": [
            "00:00", "03:00", "06:00", 
            "09:00", "12:00", "15:00",
            "18:00", "21:00",
        ],
        "leadtime_hour": ["1"],
        "data_format": "grib"
    }
    client.retrieve(dataset, request, file)

def download_era(year, month, days, file):
    dataset = "reanalysis-era5-single-levels"
    request = {
        "product_type": "reanalysis",
        "variable": [ "10m_u_component_of_wind","10m_v_component_of_wind"],
        "year": [f"{year}"],
        "month": [f"{month}"],
        "day": days, 
        "time": [
            "00:00", "03:00", "06:00",
            "09:00", "12:00", "15:00",
            "18:00", "21:00"
        ],
        "area": [80, -60, 20, 75],  
        "format": "netcdf"
    }
    client.retrieve(dataset, request, file)

year = 2023
month1, month2 = 10, 11
days1, days2 = ["30", "31"], ["01", "02", "03", "04", "05"]

download_cerra(year, month1, days1, fCERRA1)
download_cerra(year, month2, days2, fCERRA2)

download_era(year, month1, days1, fERA1)
download_era(year, month2, days2, fERA2)

def mergetime(f1, f2, fout):
    cdo.mergetime(input=f"{f1} {f2}", output=fout)

mergetime(fCERRA1, fCERRA2, fCERRA)
mergetime(fERA1, fERA2, fERA)

def remap_to_cerra_grid(source_file, cerra_grid_file, output_nc, fmt="nc4", zip_level=4):
    options = f"-f {fmt} -z zip_{zip_level}" if fmt in ("nc4", "nc4c") else f"-f {fmt}"
    cdo.remapbil(cerra_grid_file, input=source_file, output=output_nc, options=options)

fERA_CERRA_grid = os.path.join(WRKDIR,"ERA","WS_on_CERRA_grid.nc")

remap_to_cerra_grid(fERA, fCERRA, fERA_CERRA_grid)

cerra_ds = xr.open_dataset(fCERRA)
var = "si10"
cerra_ds = cerra_ds.assign_coords(longitude=(((cerra_ds.longitude + 180) % 360) - 180))
cerra_daily_max = cerra_ds[var].resample(time="D").max()
cerra_daily_mean = cerra_ds["si10"].resample(time="D").mean()

era_ds = xr.open_dataset(fERA_CERRA_grid)
era_ds = era_ds.rename({"valid_time": "time"}) if "valid_time" in era_ds.dims else era_ds
wind_speed_data = np.sqrt(era_ds["u10"]**2 + era_ds["v10"]**2)
wind_speed = xr.DataArray(wind_speed_data, coords=era_ds.coords, dims=era_ds.dims)
era_ds["wind_speed"] = wind_speed
era_ds = era_ds.assign_coords(longitude=cerra_ds.longitude)
era_ds = era_ds.assign_coords(lattiude=cerra_ds.latitude)
era_daily_max = era_ds["wind_speed"].resample(time="D").max()
era_daily_mean = era_ds["wind_speed"].resample(time="D").max()

lats = cerra_daily.latitude
lons = cerra_daily.longitude

proj = ccrs.LambertConformal(
    central_latitude=50.0,
    central_longitude=8.0,
    standard_parallels=(50.0,),
    false_easting=2937000.04508007,
    false_northing=2937000.00945589,
    globe=ccrs.Globe(semimajor_axis=6371229.0, semiminor_axis=6371229.0),
)

time = cerra_daily_max["time"]
nrows = len(time)

fig, axs = plt.subplots(nrows, 3, figsize=(7.5, 23),subplot_kw={"projection": proj})
fig.subplots_adjust(left=0.06,right=0.98,top=0.93,bottom=0.10,wspace=0.0,hspace=0.12)

axs[0, 0].set_title("CERRA", fontsize=14, pad=10)
axs[0, 1].set_title("ERA5", fontsize=14, pad=10)
axs[0, 2].set_title("CERRA − ERA5", fontsize=14, pad=10)

m_speed = None
m_diff = None

for i in range(nrows):
    print(f"Working on row {i}...")
    date = pd.to_datetime(time[i].values).strftime("%Y-%m-%d")
    axs[i, 0].text(-0.08, 0.5, date,transform=axs[i, 0].transAxes,rotation=90, va="center", ha="right", fontsize=11)
    cer = cerra_daily_max.isel(time=i)
    era = era_daily_max.isel(time=i)
    panels = [(cer, 0, 25, "viridis"),(era, 0, 25, "viridis"),(cer - era, -5, 5, "RdBu_r"),]
    for j, (data, vmin, vmax, cmap) in enumerate(panels):
        ax = axs[i, j]
        pm = ax.pcolormesh(lons, lats, data, transform=ccrs.PlateCarree(), shading="auto", cmap=cmap, vmin=vmin, vmax=vmax)
        ax.coastlines(resolution="50m", linewidth=0.6)
        ax.set_xticks([])
        ax.set_yticks([])
        if j in (0, 1) and m_speed is None:
            m_speed = pm
        if j == 2 and m_diff is None:
            m_diff = pm
            
cbar1 = fig.colorbar(m_speed, ax=axs[:, :2], orientation="horizontal", shrink=0.95, pad=0.03)
cbar1.set_label("Wind speed (m/s)")

cbar2 = fig.colorbar(m_diff, ax=axs[:, 2], orientation="horizontal", shrink=0.95, pad=0.03)
cbar2.set_label("CERRA − ERA5 (m/s)")

fig.suptitle("Evolution of Storm: Daily Maximum 10m Wind Speed", fontsize=16)
plt.show()

DATADIR=os.path.join(WRKDIR,"OBSERVATIONS")
if not os.path.isdir(DATADIR):
    os.makedirs(DATADIR)
fEOBS_ZIP = os.path.join(DATADIR, "WS_total.zip")
fEOBS_TOTAL = os.path.join(DATADIR,"WS_"+"total"+".nc")
fEOBS_SUB = os.path.join(DATADIR, "WS_subsection.nc")
fEOBS_CERRA_GRID = os.path.join(DATADIR, "WS_subsection_on_CERRA.nc")

print(f"Will be saving data in files: \n{fEOBS_ZIP}\n{fEOBS_TOTAL}\n{fEOBS_SUB}\n{fEOBS_CERRA_GRID}")

def download_observations(period, file):
    dataset = "insitu-gridded-observations-europe"
    request = {
        "product_type": "ensemble_mean",
        "variable": ["wind_speed"],
        "grid_resolution": "0_1deg",
        "period": period,
        "version": ["31_0e"]
    }
    client.retrieve(dataset, request, file)

download_observations("2011_2024", fEOBS_ZIP)

import zipfile

nc_paths = []

with zipfile.ZipFile(fEOBS_ZIP, "r") as z:
    for name in z.namelist():
        if name.endswith(".nc"):
            z.extract(name, DATADIR)
            nc_paths.append(os.path.join(DATADIR, name))

os.rename(nc_paths[0], fEOBS_TOTAL)

cdo.seldate(f"2023-10-30,2023-11-5", input=fEOBS_TOTAL, output=fEOBS_SUB)
remap_to_cerra_grid(fEOBS_SUB,fCERRA,fEOBS_CERRA_GRID)

obsar_raw = xr.open_dataset(fEOBS_CERRA_GRID)
obs_daily_mean = obsar_raw["fg"]

cer_mean_ev = cerra_daily_mean.mean(dim="time", skipna=True)
era_mean_ev = era_daily_mean.mean(dim="time", skipna=True)
obs_mean_ev = obs_daily_mean.mean(dim="time", skipna=True)

mask = np.isfinite(obs_mean_ev)
cer_mean_m = cer_mean_ev.where(mask)
era_mean_m = era_mean_ev.where(mask)

bias_cer = (cer_mean_ev - obs_mean_ev).where(mask)
bias_era = (era_mean_ev - obs_mean_ev).where(mask)

fig, axs = plt.subplots(2, 3, figsize=(14, 8),subplot_kw={"projection": proj},constrained_layout=True)
axs[0, 0].set_title("Mean: E-OBS", fontsize=13, pad=10)
axs[0, 1].set_title("Mean: CERRA", fontsize=13, pad=10)
axs[0, 2].set_title("Mean: ERA5", fontsize=13, pad=10)
axs[1, 0].set_title("", fontsize=13, pad=10)
axs[1, 1].set_title("Event-Mean Bias: CERRA - E-OBS", fontsize=13, pad=10)
axs[1, 2].set_title("Event-Mean Bias: ERA5 - E-OBS", fontsize=13, pad=10)

vmax_speed = 15
vmax_bias  = 4

panels = {
    (0, 0): (obs_mean_ev,    0,  vmax_speed, "viridis"),
    (0, 1): (cer_mean_m,  0,  vmax_speed, "viridis"),
    (0, 2): (era_mean_m,  0,  vmax_speed, "viridis"),
    (1, 0): (None,        None, None,      None),      
    (1, 1): (bias_cer,   -vmax_bias, vmax_bias, "RdBu_r"),
    (1, 2): (bias_era,   -vmax_bias, vmax_bias, "RdBu_r"),
}

m_speed = None
m_bias  = None

for (r, c), (data, vmin, vmax, cmap) in panels.items():
    ax = axs[r, c]
    if data is None: ax.axis("off"); continue
    pm = ax.pcolormesh(lons, lats, data, transform=ccrs.PlateCarree(), shading="auto", cmap=cmap, vmin=vmin, vmax=vmax)
    ax.coastlines(resolution="50m", linewidth=0.6)
    ax.set_xticks([])
    ax.set_yticks([])
    if r == 0 and m_speed is None: m_speed = pm
    if r == 1 and c in (1, 2) and m_bias is None: m_bias = pm

cbar1 = fig.colorbar(m_speed, ax=axs[0, :], orientation="vertical", shrink=0.95, pad=0.03)
cbar1.set_label("Event-Mean Wind Speed (m/s)")
cbar2 = fig.colorbar(m_bias, ax=axs[1, 1:], orientation="vertical", shrink=0.95, pad=0.03)
cbar2.set_label("Bias (model − E-OBS) (m/s)")

fig.suptitle("Event-Mean Wind Speed and Bias (30-10-2023 to 5-11-2023)", fontsize=15)
plt.show()

def point_timeseries_with_locator(lat0, lon0, cerra_ds, era_ds, obs_ds, cerra_var, era_var, obs_var, time_step):
    pad = 120
    
    cer = cerra_ds[cerra_var]
    era = era_ds[era_var]
    obs = obs_ds[obs_var]
    
    lat2d = cerra_ds["latitude"]
    lon2d = cerra_ds["longitude"]
    dist2 = (lat2d - lat0) ** 2 + (lon2d - lon0) ** 2
    j, i = np.unravel_index(dist2.argmin().item(), dist2.shape)
    lat_pt = float(lat2d.isel(y=j, x=i))
    lon_pt = float(lon2d.isel(y=j, x=i))

    print(f"Selected grid cell: y={j}, x={i}")
    print(f"Nearest grid lat/lon: {lat_pt:.4f}, {lon_pt:.4f}")

    cer_ts = cer.isel(y=j, x=i)
    era_ts = era.isel(y=j, x=i)
    obs_ts = obs.isel(y=j, x=i)
        
    #proj = ccrs.PlateCarree()
    fig, axs = plt.subplots(1, 2, figsize=(14, 5), gridspec_kw={"width_ratios": [2.2, 1.3]})

    ax = axs[0]
    cer_ts.plot(ax=ax, label="CERRA", linewidth=1, color="tab:blue", ls="--")
    (cer_ts.resample(time="D").mean()).plot(ax=ax, label="CERRA daily mean", linewidth=2, color="tab:blue", ls="-", marker="o")
    era_ts.plot(ax=ax, label="ERA5", linewidth=1, linestyle="--", color="tab:orange")
    (era_ts.resample(time="D").mean()).plot(ax=ax, label="ERA daily mean", linewidth=2, color="tab:orange", ls="-", marker="o")
    obs_ts.plot(ax=ax, label="OBS", linewidth=2, marker="o", color="tab:green")
    ax.set_title(f"Wind Speed at Selected Grid Cell (Target: {lat0:.3f}N, {lon0:.3f}W)")
    ax.set_xlabel("")
    ax.set_ylabel("Wind speed (m/s)")
    ax.grid(True, alpha=0.3)
    ax.legend(frameon=False, loc="upper left")

    axs[1].remove()
    axm = fig.add_subplot(1, 2, 2, projection=ccrs.PlateCarree())
    bg = cer.isel(time=time_step)
    bg = bg.where(np.isfinite(obs.isel(time=time_step)))
    y0 = max(j - pad, 0)
    y1 = min(j + pad, bg.sizes["y"] - 1)
    x0 = max(i - pad, 0)
    x1 = min(i + pad, bg.sizes["x"] - 1)
    bgz  = bg.isel(y=slice(y0, y1), x=slice(x0, x1))
    latz = lat2d.isel(y=slice(y0, y1), x=slice(x0, x1))
    lonz = lon2d.isel(y=slice(y0, y1), x=slice(x0, x1))
    sc = axm.scatter(lons.values.ravel(), lats.values.ravel(), c=bg.values.ravel(), s=2)
    axm.scatter([lon0], [lat0], marker="s", s=70, facecolors="None", edgecolors="white")
    axm.coastlines(resolution="50m", linewidth=0.8)
    axm.set_title("Selected Grid Cell")
    axm.set_extent([float(lonz.min()), float(lonz.max()), float(latz.min()), float(latz.max())],crs=ccrs.PlateCarree())
    axm.gridlines(draw_labels=False, linewidth=0.3, alpha=0.5)
    cbar = fig.colorbar(sc, ax=axm, orientation="vertical", shrink=0.8, pad=0.02)
    date = pd.to_datetime(time[time_step].values).strftime("%Y-%m-%d")
    cbar.set_label(f"CERRA Wind Speed at {date} (m/s)")
    plt.show()

lat0 = 43.22
lon0 = -7.25
point_timeseries_with_locator(lat0, lon0, cerra_ds, era_ds, obsar_raw, "si10", "wind_speed", "fg", 2)

import subprocess

subprocess.run([
    "python3", "-m", "nbconvert",
    "--to", "html",
    "Draft2_Ciarán.ipynb"
])