Drought duration

admin_id = "EL64"

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

# Define data directories
data_dir = Path("../data")
ecde_dir = data_dir / "ecde"
reanalysis_dir = ecde_dir / "reanalysis"
proj_dir = ecde_dir / "projections"
output_dir = data_dir / "EL64" / "drought_hazard"

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

print(f"Data will be saved to: {ecde_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}")

def unzip_ecde_data(zipfile_path):
    """
    Unzip ECDE 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}")

# Define filename and years for reanalysis data
reanalysis_zipfile = reanalysis_dir / "drought_duration_nuts2_reanalysis_1940_2023.zip"
reanalysis_ncfile = reanalysis_zipfile.with_suffix('.nc')

# Download reanalysis data
print("Downloading reanalysis drought duration data (1940-2023)...")
print("=" * 60)

# Skip if already downloaded
if reanalysis_ncfile.exists():
    print(f"Reanalysis data already exists, skipping download")
else:
    print(f"Downloading reanalysis data (1940-2023)...")
    
    dataset = "sis-ecde-climate-indicators"
    request = {
        "variable": ["duration_of_meteorological_droughts"],
        "origin": "reanalysis",
        "temporal_aggregation": ["yearly"],
        "spatial_aggregation": "regional_layer",
        "regional_layer": ["nuts_level_2"],
    }
    
    try:
        client = cdsapi.Client()
        client.retrieve(dataset, request).download(str(reanalysis_zipfile))
        unzip_ecde_data(reanalysis_zipfile)
        print(f"  ✓ Download complete")
    except Exception as e:
        print(f"  ✗ Error downloading reanalysis data: {e}")

print("=" * 60)
print(f"Reanalysis data download complete!")

# Define model combinations
# Note: Not all combinations are available in the dataset

gcms = ["ec_earth", "hadgem2_es", "ipsl_cm5a_mr", "mpi_esm_lr", "noresm1_m"]
rcms = ["cclm4_8_17", "hirham5", "racmo22e", "rca4", "wrf381p"]
rcps = ["rcp4_5", "rcp8_5"]
enss = ["r12i1p1", "r1i1p1", "r3i1p1"]

print(f"Model combinations to attempt:")
print(f"  GCMs: {len(gcms)}")
print(f"  RCMs: {len(rcms)}")
print(f"  RCPs: {len(rcps)}")
print(f"  Ensemble members: {len(enss)}")
print(f"  Maximum possible downloads: {len(gcms) * len(rcms) * len(rcps) * len(enss)}")
print(f"\nNote: Not all combinations are available in the dataset; we will check for their availability.")

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

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

# Calculate total iterations for progress bar
total_iterations = len(gcms) * len(rcms) * len(rcps) * len(enss)

# Create progress bar
with tqdm(total=total_iterations, desc="Processing model combinations") as pbar:
    for gcm in gcms:
        for rcm in rcms:
            for rcp in rcps:
                for ens in enss:
                    
                    years = [str(y) for y in range(1950, 2101)]
                    
                    # Create filename
                    proj_zipfile = proj_dir / f"drought_duration_nuts2_{gcm}_{rcm}_{rcp}_{ens}_1950_2100.zip"
                    proj_ncfile = proj_zipfile.with_suffix('.nc')
                    
                    # Skip if already downloaded
                    if proj_ncfile.exists():
                        skip_count += 1
                        proj_files.append(proj_ncfile)
                        pbar.update(1)
                        continue
                    
                    model_info = f"{gcm}/{rcm}/{rcp}/{ens}"
                    pbar.set_postfix_str(f"Downloading: {model_info}")
                    
                    dataset = "sis-ecde-climate-indicators"
                    request = {
                        "variable": ["duration_of_meteorological_droughts"],
                        "origin": "projections",
                        "gcm": gcm,
                        "rcm": rcm,
                        "experiment": rcp,
                        "ensemble_member": ens,
                        "temporal_aggregation": ["yearly"],
                        "spatial_aggregation": "regional_layer",
                        "regional_layer": ["nuts_level_2"],
                    }
                    
                    try:
                        client = cdsapi.Client()
                        client.retrieve(dataset, request).download(str(proj_zipfile))
                        unzip_ecde_data(proj_zipfile)
                        proj_files.append(proj_ncfile)
                        download_count += 1
                    except Exception as e:
                        error_count += 1
                        # Silently count errors without printing each one
                    
                    pbar.update(1)

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)}")

# Read reanalysis NetCDF file
reanalysis_ncfile = reanalysis_dir / "drought_duration_nuts2_reanalysis_1940_2023.nc"

if reanalysis_ncfile.exists():
    # Open reanalysis data
    drought_dur_reanalysis = xr.open_dataset(reanalysis_ncfile)
    print("\nReanalysis dataset:")
    print(drought_dur_reanalysis)
    
    # Extract data for EL64
    drought_dur_reanalysis_el64 = drought_dur_reanalysis.sel(nuts=admin_id)
    print(f"\nData for {admin_id}:")
    print(drought_dur_reanalysis_el64)
else:
    print("Reanalysis data file not found. Please download the data first.")

# Find all projection NetCDF files
proj_nc_files = sorted(proj_dir.glob("drought_duration_nuts2_*_1950_2100.nc"))
print(f"Found {len(proj_nc_files)} projection data files")

print(f"\nAvailable model combinations:")
for i, file in enumerate(proj_nc_files[:5]):
    print(f"  {i+1}. {file.stem}")
if len(proj_nc_files) > 5:
    print(f"  ... and {len(proj_nc_files) - 5} more")

# Load projection data for each model file
projection_data = {}

for file in proj_nc_files:
    # Extract model info from filename
    # Format: drought_duration_nuts2_{gcm}_{rcm}_{rcp}_{ens}_1950_2100.nc
    parts = file.stem.split('_')
    
    # Find where RCP starts (it always starts with 'rcp')
    rcp_idx = next((i for i, p in enumerate(parts) if p.startswith('rcp')), None)
    
    if rcp_idx:
        gcm = '_'.join(parts[3:rcp_idx-1])  # Everything between 'nuts2' and RCM
        rcm = parts[rcp_idx-1]  # Part just before RCP
        rcp = f"{parts[rcp_idx]}_{parts[rcp_idx+1]}"  # e.g., rcp4_5 or rcp8_5
        ens = parts[rcp_idx+2]  # Ensemble member after RCP (before year range)
        
        model_key = f"{gcm}_{rcm}_{rcp}_{ens}"
        
        try:
            # Load dataset and extract EL64 data
            ds = xr.open_dataset(file)
            ds_el64 = ds.sel(nuts=admin_id)
            projection_data[model_key] = ds_el64
        except Exception as e:
            print(f"Error loading {model_key}: {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])

# Save reanalysis data to CSV
if 'drought_dur_reanalysis_el64' in locals():
    reanalysis_df = drought_dur_reanalysis_el64['dmd'].to_dataframe().reset_index()
    reanalysis_df['scenario'] = 'reanalysis'
    reanalysis_df['model'] = 'reanalysis'
    reanalysis_csv = output_dir / f"drought_duration_reanalysis_{admin_id}.csv"
    reanalysis_df.to_csv(reanalysis_csv, index=False)
    print(f"Saved reanalysis data to: {reanalysis_csv}")
    print(f"  Shape: {reanalysis_df.shape}")
    print(f"  Columns: {list(reanalysis_df.columns)}")

# Save projection data to CSV
if len(projection_data) > 0:
    all_proj_data = []
    
    for model_key, ds in projection_data.items():
        # Parse model information from model_key
        # Format: {gcm}_{rcm}_{rcp}_{ens}
        parts = model_key.split('_')
        
        # Find where RCP starts (it always has format rcp4 or rcp8)
        rcp_idx = next(i for i, p in enumerate(parts) if p.startswith('rcp'))
        
        gcm = '_'.join(parts[:rcp_idx-1])  # Everything before RCM
        rcm = parts[rcp_idx-1]  # Part just before RCP
        rcp = f"{parts[rcp_idx]}_{parts[rcp_idx+1]}"  # e.g., rcp4_5 or rcp8_5
        ens = '_'.join(parts[rcp_idx+2:])  # Everything after RCP (e.g., r1i1p1)
        
        # Convert to dataframe
        df = ds['dmd'].to_dataframe().reset_index()
        df['scenario'] = rcp.upper()  # Use RCP4_5 or RCP8_5
        df['gcm'] = gcm
        df['rcm'] = rcm
        df['ensemble'] = ens
        df['model'] = model_key
        
        all_proj_data.append(df)
    
    # Combine all projection data
    proj_df = pd.concat(all_proj_data, ignore_index=True)
    
    # Save to CSV
    proj_csv = output_dir / f"drought_duration_projections_{admin_id}.csv"
    proj_df.to_csv(proj_csv, index=False)
    print(f"\nSaved projection data to: {proj_csv}")
    print(f"  Shape: {proj_df.shape}")
    print(f"  Columns: {list(proj_df.columns)}")
    print(f"  Unique models: {proj_df['model'].nunique()}")
    print(f"  Scenarios: {proj_df['scenario'].unique()}")

# Create a combined CSV with both reanalysis and projections
if 'reanalysis_df' in locals() and 'proj_df' in locals():
    # Align columns
    common_cols = ['time', 'nuts', 'dmd', 'scenario', 'model']
    
    reanalysis_df_clean = reanalysis_df[common_cols].copy()
    proj_df_clean = proj_df[['time', 'nuts', 'dmd', 'scenario', 'model']].copy()
    
    combined_df = pd.concat([reanalysis_df_clean, proj_df_clean], ignore_index=True)
    combined_df = combined_df.sort_values('time').reset_index(drop=True)
    
    combined_csv = output_dir / f"drought_duration_all_{admin_id}.csv"
    combined_df.to_csv(combined_csv, index=False)
    print(f"\nSaved combined data to: {combined_csv}")
    print(f"  Shape: {combined_df.shape}")
    print(f"  Time range: {combined_df['time'].min()} to {combined_df['time'].max()}")