admin_id = "ITF2"
workdir = "/home/nejk/code/extreme_precip_exposure"

rerun = True  # Set to True to re-download data, False to skip existing files
create_sample_data = True  # Set to True to create sample data, False to skip

import os
import zipfile
import cdsapi
import glob
from pathlib import Path
import xarray as xr
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import geopandas as gpd
import cartopy.crs as ccrs
import cartopy.feature as cfeature
from functools import partial
from tqdm.auto import tqdm
import regionmask

# Set working directory (adjust as needed)
os.chdir(workdir)
print(f"Working directory: {os.getcwd()}")

# Define data directories
data_dir = Path("./data")
cds_dir = data_dir / "c3s_atlas"
hist_dir = cds_dir / "historical"
proj_dir = cds_dir / "projections"
output_dir = data_dir / admin_id / "extreme_precip_exposure"

# Create directories if they don't exist
for directory in [cds_dir, hist_dir, proj_dir, output_dir]:
    directory.mkdir(parents=True, exist_ok=True)

print(f"Data will be saved to: {cds_dir}")
print(f"Output files will be saved to: {output_dir}")

regions_dir = data_dir / 'regions'
nuts_shp = regions_dir / 'NUTS_RG_20M_2024_4326' / 'NUTS_RG_20M_2024_4326.shp'
nuts_gdf = gpd.read_file(nuts_shp)

# Select the region of interest
sel_gdf = nuts_gdf[nuts_gdf['NUTS_ID'] == admin_id]
print(f"Region: {sel_gdf['NUTS_NAME'].values[0]}")
print(f"NUTS ID: {admin_id}")
print(f"Country: {sel_gdf['CNTR_CODE'].values[0]}")
print(f"Bounding box: {sel_gdf.geometry.total_bounds}")

# Rounded bounding box for downloading data (lon_min, lat_min, lon_max, lat_max)
lon_min, lat_min, lon_max, lat_max = sel_gdf.geometry.total_bounds
lon_min = (np.floor(lon_min) -1)
lat_min = (np.floor(lat_min) -1)
lon_max = (np.ceil(lon_max) +1)
lat_max = (np.ceil(lat_max) +1)
bbox = [lon_min, lat_min, lon_max, lat_max]
print(f"Rounded bounding box for data download: {bbox}")

def unzip_cds_data(zipfile_path):
    """
    Unzip data downloaded from CDS and rename to .nc extension.
    
    Parameters:
    -----------
    zipfile_path : Path or str
        Path to the zip file to extract
    """
    zipfile_path = Path(zipfile_path)
    
    with zipfile.ZipFile(zipfile_path, 'r') as zip_ref:
        names = zip_ref.namelist()
        zip_ref.extractall(zipfile_path.parent)
        
        for name in names:
            if name.split(".")[-1] == "nc":
                # Rename the extracted file to match the zip filename
                nc_path = zipfile_path.with_suffix('.nc')
                os.rename(zipfile_path.parent / name, nc_path)
                print(f"  Extracted: {nc_path.name}")
            else:
                # Remove non-NetCDF files
                os.remove(zipfile_path.parent / name)
    
    # Remove the zip file after extraction
    os.remove(zipfile_path)
    print(f"  Removed: {zipfile_path.name}")

# Download historical climate model outputs
print("Downloading historical data (1970-2005)...")
print("=" * 60)

# Create filename
hist_zipfile = hist_dir / f"daily_precip_monmax_1970_2005_{admin_id}.zip"
hist_ncfile = hist_zipfile.with_suffix('.nc')
            
if not rerun and hist_ncfile.exists():
    print(f"  File already exists, skipping: {hist_ncfile.name}")
else:
    dataset = "multi-origin-c3s-atlas"
    request = {
        "variable": ["monthly_maximum_1_day_precipitation"],
        "origin": "cordex_eur_11",
        "domain": "euro_cordex",
        "experiment": "historical",
        "bias_adjustment": "no_bias_adjustment",
        "period": "1970-2005",
        "area": [lat_max, lon_min, lat_min, lon_max], # N, W, S, E
    }
    try:
        client = cdsapi.Client()
        client.retrieve(dataset, request).download(str(hist_zipfile))
        unzip_cds_data(hist_zipfile)
        proj_files.append(hist_ncfile)
    except Exception as e:
        print(f"  Error downloading historical data: {e}")

# Define projection scenarios to download
rcps = ["rcp_2_6","rcp_4_5", "rcp_8_5"]

# Download projection data
print("Downloading projection data (2006-2100)...")
print("=" * 60)

proj_files = []
download_count = 0
skip_count = 0
error_count = 0

for rcp in rcps:

    # Create filename
    proj_zipfile = proj_dir / f"daily_precip_monmax_{rcp}_2006_2100_{admin_id}.zip"
    proj_ncfile = proj_zipfile.with_suffix('.nc')
                
    # Skip if already downloaded
    if not rerun and proj_ncfile.exists():
        skip_count += 1
        proj_files.append(proj_ncfile)
        continue                  
            
    dataset = "multi-origin-c3s-atlas"
    request = {
        "variable": ["monthly_maximum_1_day_precipitation"],
        "origin": "cordex_eur_11",
        "experiment": rcp,
        "bias_adjustment": "no_bias_adjustment",
        "period": "2006-2100",
        "area": [lat_max, lon_min, lat_min, lon_max], # N, W, S, E
    }
                
    try:
        client = cdsapi.Client()
        client.retrieve(dataset, request).download(str(proj_zipfile))
        unzip_cds_data(proj_zipfile)
        proj_files.append(proj_ncfile)
        download_count += 1
    except Exception as e:
        error_count += 1
        print(f"  Error downloading projection data: {e}")
    

print("=" * 60)
print(f"Projection data download summary:")
print(f"  New downloads: {download_count}")
print(f"  Already existed: {skip_count}")
print(f"  Errors/unavailable: {error_count}")
print(f"  Total projection files: {len(proj_files)}")

# Use regionmask to create a 2D mask for the NUTS region from sel_gdf geometry and with grid from the historical dataset
nuts_mask = regionmask.mask_geopandas(sel_gdf, xr.open_dataset(hist_dir / f"daily_precip_monmax_1970_2005_{admin_id}.nc"))
# reset index to np.nan/1
nuts_mask = xr.where(nuts_mask >= 1, 1, np.nan)
nuts_mask.plot()

# Number of grid points in the mask:
print(f"\nNumber of grid points in the mask: {np.sum(~np.isnan(nuts_mask.values))}")

# Read historical NetCDF file
hist_ncfile = hist_dir / f"daily_precip_monmax_1970_2005_{admin_id}.nc" 

if hist_ncfile.exists():
    # Open historical data
    precip_hist = xr.open_dataset(hist_ncfile)
    print("\nHistorical dataset:")
    print(precip_hist)

    # print all multi model ensemble members by gcm_model and rcm_model and ensemble member combinations and associate with member_id 
    mm_ensemble = precip_hist.groupby("member_id").groups.keys()
    print("\nMulti-model ensemble member:")
    i = 0
    for combo in mm_ensemble:
        print(combo) #  the member_id is defined as {gcm_institution}_{gcm_model}_{gcm_variant}_{rcm_institution}_{rcm_model}_{rcm_variant}
        #print(f"{precip_hist['gcm_institution'].values[i]}_{precip_hist['gcm_model'].values[i]}_{precip_hist['gcm_variant'].values[i]}_{precip_hist['rcm_institution'].values[i]}_{precip_hist['rcm_model'].values[i]}_{precip_hist['rcm_variant'].values[i]}")
        i += 1
else:
    print("Historical data file not found. Please download the data first.")

n_mmens = len(mm_ensemble)
print(f"\n {"-"*80} \n   * {n_mmens} multi-model ensemble members found in the historical dataset. \n {"-"*80} \n")

# Extract data for NUTS region using regionmask
precip_hist_nuts = precip_hist.where(nuts_mask == 1, drop=True)
print(f"\nData for {admin_id}:")
print(precip_hist_nuts)

#precip_hist_nuts.isel(time=0).isel(member=0)['rx1day'].plot()
#precip_hist_nuts.isel(member=0)['rx1day'].max(dim=['lon','lat']).plot()

# File 1: all data
# - convert xarray dataset precip_hist_nuts to a pandas dataframe
# - cleanup to save storage space and make it easier to work with for further analysis
# - save all lon/lat points in the NUTS region to a csv file with columns for time (format: YYYY-MM), member/member_id (but also rcm and gcm details, i.e. gcm_institution, gcm_model, gcm_variant, rcm_institution, rcm_model, rcm_variant), lon, lat, and rx1day
precip_hist_nuts_df = precip_hist_nuts.to_dataframe().reset_index()
precip_hist_nuts_df['time'] = precip_hist_nuts_df['time'].dt.strftime('%Y-%m')
precip_hist_nuts_df.drop(columns=['lon_bnds', 'lat_bnds', 'bnds', 'time_bnds', 'crs'], inplace=True) # remove some unnecessary columns
precip_hist_nuts_df.dropna(inplace=True) # drop all rows with missing values (i.e. points outside the NUTS region)
precip_hist_nuts_df['admin_id'] = admin_id
precip_hist_nuts_df.to_csv(f"{output_dir}/{admin_id}_precip_hist_all_points.csv", index=False)
print(f"Written file: {output_dir}/{admin_id}_precip_hist_all_points.csv")

# File 2: annual maximum and NUTS maximum of rx1day for each time step and member/member_id 
precip_hist_nuts_amax = precip_hist_nuts.groupby('time.year').max()
precip_hist_nuts_amax_df = precip_hist_nuts_amax.to_dataframe().reset_index()
precip_hist_nuts_amax_df.drop(columns=['lon_bnds', 'lat_bnds', 'bnds', 'time_bnds', 'crs'], inplace=True) # remove some unnecessary columns
precip_hist_nuts_amax_df['admin_id'] = admin_id
precip_hist_nuts_amax_df.to_csv(f"{output_dir}/{admin_id}_precip_hist_amax_all_points.csv", index=False)
print(f"Written file: {output_dir}/{admin_id}_precip_hist_amax_all_points.csv")

if create_sample_data:
    
    # Sample grid point; the one closest to the centroid of the NUTS region:
    nuts_centroid = (
        sel_gdf.to_crs(epsg=3035)   # projected CRS for accurate centroid
        .geometry.centroid
        .to_crs(sel_gdf.crs)        # back to original CRS (lon/lat)
        .values[0]
    )
    lon = precip_hist_nuts['lon'].values
    lat = precip_hist_nuts['lat'].values
    lon_idx = np.argmin(np.abs(lon - nuts_centroid.x))
    lat_idx = np.argmin(np.abs(lat - nuts_centroid.y))

    # Save monthly data for one member and one grid point (the one closest to the centroid of the NUTS region)
    imember = 0
    precip_hist_nuts_df[(precip_hist_nuts_df['member'] == imember) & (precip_hist_nuts_df['lon'] == lon[lon_idx]) & (precip_hist_nuts_df['lat'] == lat[lat_idx])].to_csv(f"{output_dir}/{admin_id}_precip_hist_sample_model_gridpoint.csv", index=False)
    print(f"Written file: {output_dir}/{admin_id}_precip_hist_sample_model_gridpoint.csv")

    # Save just annual max of one member but for all grid points in NUTS region
    imember = 0
    precip_hist_nuts_amax_df[precip_hist_nuts_amax_df['member'] == imember].to_csv(f"{output_dir}/{admin_id}_precip_hist_amax_sample_model.csv", index=False)
    print(f"Written file: {output_dir}/{admin_id}_precip_hist_amax_sample_model.csv")

    # Save just annual max one grid point (the one closest to the centroid of the NUTS region) but for all models
    precip_hist_nuts_amax_df[(precip_hist_nuts_amax_df['lon'] == lon[lon_idx]) & (precip_hist_nuts_amax_df['lat'] == lat[lat_idx])].to_csv(f"{output_dir}/{admin_id}_precip_hist_amax_sample_gridpoint.csv", index=False)
    print(f"Written file: {output_dir}/{admin_id}_precip_hist_amax_sample_gridpoint.csv")

# Load projection data for each RCP
projection_data = {}

for rcp in rcps:
    proj_ncfile = proj_dir / f"daily_precip_monmax_{rcp}_2006_2100_{admin_id}.nc"
   
    if proj_ncfile.exists():
        try:
            # Load dataset
            ds = xr.open_dataset(proj_ncfile)
            projection_data[rcp] = ds
        except Exception as e:
            print(f"Error loading {proj_ncfile}: {e}")

print(f"\nSuccessfully loaded {len(projection_data)} model combinations")

# Show example
if len(projection_data) > 0:
    example_key = list(projection_data.keys())[0]
    print(f"\nExample dataset ({example_key}):")
    print(projection_data[example_key])

precip_proj_nuts = {}

for rcp in rcps:
    precip_proj_nuts[rcp] = projection_data[rcp].where(nuts_mask == 1, drop=True)

# - convert xarray dataset precip_proj_nuts to a pandas dataframe, but include all rcp scenario as a column in the dataframe
# - cleanup to save storage space and make it easier to work with for further analysis
# - save all lon/lat points in the NUTS region to a csv file with columns for time (format: YYYY-MM), member/member_id (but also rcm and gcm details, i.e. gcm_institution, gcm_model, gcm_variant, rcm_institution, rcm_model, rcm_variant), lon, lat, and rx1day
precip_proj_nuts_df = pd.DataFrame()
for rcp in rcps:
    dummy = precip_proj_nuts[rcp].to_dataframe().reset_index()
    dummy['rcp'] = rcp
    precip_proj_nuts_df = pd.concat([precip_proj_nuts_df, dummy], ignore_index=True)
precip_proj_nuts_df['time'] = precip_proj_nuts_df['time'].dt.strftime('%Y-%m')
precip_proj_nuts_df.drop(columns=['lon_bnds', 'lat_bnds', 'bnds', 'time_bnds', 'crs'], inplace=True) # remove some unnecessary columns
precip_proj_nuts_df.dropna(inplace=True) # drop all rows with missing values (i.e. points outside the NUTS region)
precip_proj_nuts_df['admin_id'] = admin_id
precip_proj_nuts_df.to_csv(f"{output_dir}/{admin_id}_precip_proj_all_points.csv", index=False)
print(f"Written file: {output_dir}/{admin_id}_precip_proj_all_points.csv")

# File 2: annual maximum
precip_proj_nuts_amax_df = pd.DataFrame()
for rcp in rcps:
    dummy = precip_proj_nuts[rcp].groupby('time.year').max().to_dataframe().reset_index()
    dummy['rcp'] = rcp
    precip_proj_nuts_amax_df = pd.concat([precip_proj_nuts_amax_df, dummy], ignore_index=True)
precip_proj_nuts_amax_df.drop(columns=['lon_bnds', 'lat_bnds', 'bnds', 'time_bnds', 'crs'], inplace=True) # remove some unnecessary columns
precip_proj_nuts_amax_df['admin_id'] = admin_id
precip_proj_nuts_amax_df.to_csv(f"{output_dir}/{admin_id}_precip_proj_amax_all_points.csv", index=False)
print(f"Written file: {output_dir}/{admin_id}_precip_proj_amax_all_points.csv")

if create_sample_data:
    
    # Sample grid point; the one closest to the centroid of the NUTS region:
    nuts_centroid = (
        sel_gdf.to_crs(epsg=3035)   # projected CRS for accurate centroid
        .geometry.centroid
        .to_crs(sel_gdf.crs)        # back to original CRS (lon/lat)
        .values[0]
    )
    lon = precip_proj_nuts[rcps[0]]['lon'].values
    lat = precip_proj_nuts[rcps[0]]['lat'].values
    lon_idx = np.argmin(np.abs(lon - nuts_centroid.x))
    lat_idx = np.argmin(np.abs(lat - nuts_centroid.y))

    # Save monthly data for one member and one grid point (the one closest to the centroid of the NUTS region)
    imember = 0
    precip_proj_nuts_df[(precip_proj_nuts_df['member'] == imember) & (precip_proj_nuts_df['lon'] == lon[lon_idx]) & (precip_proj_nuts_df['lat'] == lat[lat_idx])].to_csv(f"{output_dir}/{admin_id}_precip_proj_sample_model_gridpoint.csv", index=False)
    print(f"Written file: {output_dir}/{admin_id}_precip_proj_sample_model_gridpoint.csv")

    # Save just annual max of one member but for all grid points in NUTS region
    imember = 0
    precip_proj_nuts_amax_df[precip_proj_nuts_amax_df['member'] == imember].to_csv(f"{output_dir}/{admin_id}_precip_proj_amax_sample_model.csv", index=False)
    print(f"Written file: {output_dir}/{admin_id}_precip_proj_amax_sample_model.csv")

    # Save just annual max one grid point (the one closest to the centroid of the NUTS region) but for all models
    precip_proj_nuts_amax_df[(precip_proj_nuts_amax_df['lon'] == lon[lon_idx]) & (precip_proj_nuts_amax_df['lat'] == lat[lat_idx])].to_csv(f"{output_dir}/{admin_id}_precip_proj_amax_sample_gridpoint.csv", index=False)
    print(f"Written file: {output_dir}/{admin_id}_precip_proj_amax_sample_gridpoint.csv")