# ============================================================
# STANDALONE MACHINE LEARNING SECTION
# OBS vs WRF/LASSO vs DP-SCREAM
#
# Goal:
#   Rank which variables are most associated with each model's
#   low-cloud-fraction bias.
#
# Target:
#   cf_bias = model_low_cloud_fraction - observed_low_cloud_fraction
#
# Design choices (based on prior discussion):
#   - Feature sets are CURATED, not auto-selected, to avoid
#     state-vs-bias redundancy and quasi-tautological features
#     (e.g., qc/qr mean are too close to cloud fraction itself).
#   - WRF and DP-SCREAM share a core set of environmental-bias
#     features so their importance bars are directly comparable.
#   - DP-SCREAM gets additional process diagnostics (tendencies,
#     fluxes, TKE) that WRF does not expose.
#   - RH biases (surface + BL) are included since RH is the
#     proximal driver of cloud formation.
# ============================================================

import numpy as np
import pandas as pd
import xarray as xr
import matplotlib.pyplot as plt

from pathlib import Path
from glob import glob
from itertools import compress

from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import LeaveOneGroupOut
from sklearn.metrics import r2_score, mean_squared_error
from sklearn.impute import SimpleImputer
from sklearn.pipeline import Pipeline
from sklearn.inspection import permutation_importance


# ============================================================
# 0. SETTINGS
# ============================================================

BASE_LASSO = "/data/project/ARM_Summer_School_2026/data/modeling/lasso"
BASE_DPSCREAM = "/data/project/ARM_Summer_School_2026/data/modeling/dpscream"

CASES = ["20180709", "20190517", "20190929"]
CASE_LABELS = ["2018-07-09", "2019-05-17", "2019-09-29"]
CASE_COLORS = ["#1f77b4", "#d62728", "#2ca02c"]

ML_CASE_LABELS = CASE_LABELS
ML_CASE_COLORS = CASE_COLORS

UTC_TO_LST = pd.Timedelta(hours=-6)   # SGP is ~ -97.5 deg lon; solar LST ~ UTC-6.5h.
                                       # Using -6 (CST) is the ARM convention.
RANDOM_STATE = 42

OUTDIR = Path("./cloud_model_eval_outputs")
OUTDIR.mkdir(exist_ok=True)

TARGET_COL = "cf_bias"

N_ESTIMATORS = 500
N_REPEATS_PERM = 100


# ============================================================
# 1. SMALL HELPERS
# ============================================================

def safe_display(obj):
    try:
        display(obj)
    except NameError:
        print(obj)


def to_1d(x):
    return np.asarray(x).ravel()


def to_lst(t):
    return pd.DatetimeIndex(pd.to_datetime(t)) + UTC_TO_LST


def safe_compute(da):
    try:
        return da.compute()
    except Exception:
        return da


def qv_to_gkg_values(values):
    """
    Convert qv to g/kg if it looks like kg/kg.
    """
    arr = np.asarray(values, dtype=float)
    finite = arr[np.isfinite(arr)]

    if len(finite) == 0:
        return arr

    med = np.nanmedian(finite)

    if abs(med) < 1:
        return arr * 1000.0
    else:
        return arr


def qv_to_gkg_da(da):
    """
    Same conversion but for xarray DataArray.
    Uses the full-array median rather than a single point.
    """
    if da is None:
        return None

    try:
        med = float(da.median(skipna=True).compute().values)
    except Exception:
        try:
            med = float(np.nanmedian(np.asarray(da.values)))
        except Exception:
            return da

    if not np.isfinite(med):
        return da

    if abs(med) < 1:
        return da * 1000.0
    else:
        return da


def interp_da_to_time(source_da, target_time):
    """
    Interpolate a 1D time DataArray to target_time.
    """
    target_time = pd.to_datetime(target_time)

    if source_da is None:
        return np.full(len(target_time), np.nan)

    source_da = safe_compute(source_da)

    if "time" not in source_da.coords:
        return np.full(len(target_time), np.nan)

    t_src = pd.to_datetime(source_da["time"].values)
    y_src = to_1d(source_da.values)

    if len(t_src) == 0 or len(y_src) == 0:
        return np.full(len(target_time), np.nan)

    return np.interp(
        target_time.astype("int64"),
        t_src.astype("int64"),
        y_src,
        left=np.nan,
        right=np.nan,
    )


def truncate(arr, n):
    arr = to_1d(arr)
    out = np.full(n, np.nan, dtype=float)
    m = min(n, len(arr))
    out[:m] = arr[:m]
    return out


# ============================================================
# 2. LOAD DATASETS
# ============================================================

# -------------------------
# LASSO / WRF / OBS main files
# -------------------------
lasso_files = glob(
    f"{BASE_LASSO}/*/sgplassodiagconfobsmod*/obs_model/"
    "sgplassodiagobsmod*C1.m1*.nc"
)

lasso_files = list(compress(
    lasso_files,
    [("modz" not in f) and ("5C1" not in f) for f in lasso_files]
))

lasso_files = sorted([
    f for f in lasso_files
    if any(case in f for case in CASES)
])

if len(lasso_files) == 0:
    raise FileNotFoundError("No LASSO main obs_model files found.")

print("LASSO / WRF / OBS main files:")
for f in lasso_files:
    print("  ", f)

ds_all = xr.open_mfdataset(
    lasso_files,
    combine="nested",
    concat_dim="date",
    decode_times=True,
    chunks={},
)


# -------------------------
# DP-SCREAM files
# -------------------------
dp_files = glob(f"{BASE_DPSCREAM}/*.nc")

dp_files = sorted([
    f for f in dp_files
    if any(f"{c[:4]}-{c[4:6]}-{c[6:8]}" in f for c in CASES)
])

if len(dp_files) == 0:
    raise FileNotFoundError("No DP-SCREAM files found.")

print("\nDP-SCREAM files:")
for f in dp_files:
    print("  ", f)

ds_dp = xr.open_mfdataset(
    dp_files,
    combine="nested",
    concat_dim="date",
    decode_times=True,
    chunks={},
)


# ============================================================
# 3. DP-SCREAM DERIVED VARIABLES
# ============================================================

# Track which DP variables actually resolved (vs returned None)
dp_loaded_status = {}


def _record(name, value):
    dp_loaded_status[name] = ("OK" if value is not None else "MISSING")
    return value


def dp_domain_mean(var):
    if var not in ds_dp:
        return None
    da = ds_dp[var]
    if "ncol" in da.dims:
        da = da.mean(dim="ncol", skipna=True)
    return da


def dp_low_layer_mean(var, ztop_m=3000.0):
    if var not in ds_dp or "z_mid" not in ds_dp:
        return None
    da = ds_dp[var]
    z = ds_dp["z_mid"]
    da_low = da.where(z <= ztop_m)
    dims_to_mean = [d for d in da_low.dims if d in ["lev", "ncol"]]
    if len(dims_to_mean) > 0:
        da_low = da_low.mean(dim=dims_to_mean, skipna=True)
    return da_low


def dp_low_layer_max(var, ztop_m=3000.0):
    if var not in ds_dp or "z_mid" not in ds_dp:
        return None
    da = ds_dp[var]
    z = ds_dp["z_mid"]
    da_low = da.where(z <= ztop_m)
    if "lev" in da_low.dims:
        da_low = da_low.max(dim="lev", skipna=True)
    if "ncol" in da_low.dims:
        da_low = da_low.mean(dim="ncol", skipna=True)
    return da_low


def dp_low_cf():
    return dp_low_layer_max("cldfrac_tot", ztop_m=3000.0)


def dp_total_cf():
    if "cldfrac_tot" not in ds_dp:
        return None
    da = ds_dp["cldfrac_tot"]
    if "ncol" in da.dims:
        da = da.mean(dim="ncol", skipna=True)
    if "lev" in da.dims:
        da = da.max(dim="lev", skipna=True)
    return da


def dp_lwp_gm2():
    if "LiqWaterPath" not in ds_dp:
        return None
    da = ds_dp["LiqWaterPath"]
    if "ncol" in da.dims:
        da = da.mean(dim="ncol", skipna=True)
    return da * 1000.0


def dp_real_water_fraction(threshold=1e-7, ztop_m=3000.0):
    """
    Fraction of low-level cells with non-trivial liquid+rain water.
    Kept as a "is the cloud real?" diagnostic. NOT replaced by qc_mean,
    which would be too close to LWP / cloud fraction definition.
    """
    needed = ["qc", "qr", "z_mid"]
    if not all(v in ds_dp for v in needed):
        return None
    water = ds_dp["qc"] + ds_dp["qr"]
    low = ds_dp["z_mid"] <= ztop_m
    real = (water.where(low) > threshold).astype(float)
    dims = [d for d in real.dims if d in ["lev", "ncol"]]
    if len(dims) > 0:
        real = real.mean(dim=dims, skipna=True)
    return real


def dp_max_rh_below_3km():
    """
    Max RH below 3 km. RH is the proximal driver of cloud formation.
    """
    if "RelativeHumidity" not in ds_dp or "z_mid" not in ds_dp:
        return None
    rh = ds_dp["RelativeHumidity"].where(ds_dp["z_mid"] <= 3000.0)
    if "lev" in rh.dims:
        rh = rh.max(dim="lev", skipna=True)
    if "ncol" in rh.dims:
        rh = rh.mean(dim="ncol", skipna=True)

    # Convert fraction to percent if needed, using full-array median
    try:
        med = float(rh.median(skipna=True).compute().values)
        if np.isfinite(med) and med <= 2:
            rh = rh * 100.0
    except Exception:
        pass
    return rh


def dp_bl_mean_from_pbl(var, scale=1.0):
    if var not in ds_dp:
        return None
    if "z_mid" not in ds_dp or "pbl_height" not in ds_dp:
        return None
    da = ds_dp[var]
    z = ds_dp["z_mid"]
    pblh = ds_dp["pbl_height"]
    da_bl = da.where(z <= pblh)
    dims = [d for d in da_bl.dims if d in ["lev", "ncol"]]
    if len(dims) > 0:
        da_bl = da_bl.mean(dim=dims, skipna=True)
    return da_bl * scale


def dp_low_layer_mean_rh():
    """
    Boundary-layer mean RH, in percent.
    """
    if "RelativeHumidity" not in ds_dp:
        return None
    raw = dp_bl_mean_from_pbl("RelativeHumidity", scale=1.0)
    if raw is None:
        return None
    try:
        med = float(raw.median(skipna=True).compute().values)
        if np.isfinite(med) and med <= 2:
            raw = raw * 100.0
    except Exception:
        pass
    return raw


def dp_surface_rh():
    """
    Approximate surface RH from T_2m and qv_2m + ps if needed.
    If a direct surface RH variable exists, use it; otherwise return None
    and we'll let the RH bias features only cover BL.
    """
    for candidate in ["RH_2m", "rh_2m", "RelativeHumidity_2m"]:
        if candidate in ds_dp:
            rh = ds_dp[candidate]
            if "ncol" in rh.dims:
                rh = rh.mean(dim="ncol", skipna=True)
            try:
                med = float(rh.median(skipna=True).compute().values)
                if np.isfinite(med) and med <= 2:
                    rh = rh * 100.0
            except Exception:
                pass
            return rh
    return None


# Resolve all DP variables and record status
dp_low_cf_da          = _record("dp_low_cf",          dp_low_cf())
dp_total_cf_da        = _record("dp_total_cf",        dp_total_cf())
dp_lwp_da             = _record("dp_lwp",             dp_lwp_gm2())

dp_T2m_da             = _record("T_2m",               dp_domain_mean("T_2m"))
dp_qv2m_da            = _record("qv_2m",              qv_to_gkg_da(dp_domain_mean("qv_2m")))
dp_pblh_da            = _record("pbl_height",         dp_domain_mean("pbl_height"))

dp_T_bl_da = None
for candidate in ["T_mid", "T"]:
    if candidate in ds_dp:
        dp_T_bl_da = dp_bl_mean_from_pbl(candidate, scale=1.0)
        break
_record("T_bl (from T_mid/T)", dp_T_bl_da)

dp_qv_bl_da = None
for candidate in ["qv", "qv_mid"]:
    if candidate in ds_dp:
        dp_qv_bl_da = qv_to_gkg_da(dp_bl_mean_from_pbl(candidate, scale=1.0))
        break
_record("qv_bl (from qv/qv_mid)", dp_qv_bl_da)

dp_rh_surface_da      = _record("RH_surface",         dp_surface_rh())
dp_rh_bl_da           = _record("RH_BL",              dp_low_layer_mean_rh())

real_water_fraction_da = _record("real_water_fraction", dp_real_water_fraction())
max_RH_below_3km_da    = _record("max_RH_below_3km",    dp_max_rh_below_3km())

sensible_heat_flux_da = _record("surf_sens_flux",     dp_domain_mean("surf_sens_flux"))
latent_heat_flux_da   = _record("surf_lat_flux",      dp_domain_mean("surface_upward_latent_heat_flux"))

shoc_T_tend_da        = _record("shoc_T_mid_tend",    dp_low_layer_mean("shoc_T_mid_tend"))
p3_T_tend_da          = _record("p3_T_mid_tend",      dp_low_layer_mean("p3_T_mid_tend"))
rad_T_tend_da         = _record("rrtmgp_T_mid_tend",  dp_low_layer_mean("rrtmgp_T_mid_tend"))

shoc_qv_tend_da       = _record("shoc_qv_tend",       dp_low_layer_mean("shoc_qv_tend"))
p3_qv_tend_da         = _record("p3_qv_tend",         dp_low_layer_mean("p3_qv_tend"))

tke_da                = _record("tke",                dp_low_layer_mean("tke"))
eddy_diff_mom_da      = _record("eddy_diff_mom",      dp_low_layer_mean("eddy_diff_mom"))

# Print which DP variables actually loaded - critical for diagnosing the feature set
print("\n" + "=" * 70)
print("DP-SCREAM variable load status:")
print("=" * 70)
for k, v in dp_loaded_status.items():
    print(f"  {k:35s} {v}")


# ============================================================
# 4. LASSO / WRF / OBS HELPERS
# ============================================================

def get_lasso_series(var, case_index, source_type):
    """source_type: 0 = OBS, 1 = WRF/LASSO"""
    if var not in ds_all:
        return None
    return ds_all[var].sel(source_type=source_type).isel(date=case_index)


def extract_lasso_to_obs_time(var, case_index, source_type, obs_time, qv=False):
    da = get_lasso_series(var, case_index, source_type)
    if da is None:
        return np.full(len(obs_time), np.nan)
    vals = interp_da_to_time(da, obs_time)
    if qv:
        vals = qv_to_gkg_values(vals)
    return vals


def extract_dp_to_obs_time(dp_da, case_index, obs_time):
    if dp_da is None:
        return np.full(len(obs_time), np.nan)
    try:
        da = dp_da.isel(date=case_index)
    except Exception:
        return np.full(len(obs_time), np.nan)
    return interp_da_to_time(da, obs_time)


# ============================================================
# 5. BUILD COMPARISON TABLE
# ============================================================

def build_ml_compare_df():
    """
    Build one dataframe with both:
      model = LASSO
      model = DP-SCREAM
    Time grid = OBS low cloud fraction time.
    """
    if "low_cloud_fraction_arscl" not in ds_all:
        raise ValueError("low_cloud_fraction_arscl missing from ds_all.")

    rows = []
    ncase = min(3, ds_all.sizes["date"], ds_dp.sizes["date"])

    for i in range(ncase):
        # OBS cloud fraction defines time grid
        obs_cf_da = get_lasso_series("low_cloud_fraction_arscl", i, 0)
        obs_time = pd.to_datetime(obs_cf_da["time"].values)
        obs_time_lst = to_lst(obs_time)
        n = len(obs_time)

        obs_low_cf = truncate(safe_compute(obs_cf_da).values, n)

        # OBS environmental
        obs_T2m  = extract_lasso_to_obs_time("temperature_surface", i, 0, obs_time)
        obs_qv2m = extract_lasso_to_obs_time("water_vapor_mixing_ratio_surface", i, 0, obs_time, qv=True)
        obs_rh2m = extract_lasso_to_obs_time("rh_surface", i, 0, obs_time)

        obs_T_bl  = extract_lasso_to_obs_time("temperature_boundary_layer", i, 0, obs_time)
        obs_qv_bl = extract_lasso_to_obs_time("water_vapor_mixing_ratio_boundary_layer", i, 0, obs_time, qv=True)
        obs_rh_bl = extract_lasso_to_obs_time("rh_boundary_layer", i, 0, obs_time)

        obs_lcl = extract_lasso_to_obs_time("lcl", i, 0, obs_time)
        obs_cbh = extract_lasso_to_obs_time("cloud_base_height", i, 0, obs_time)
        obs_lwp = extract_lasso_to_obs_time("lwp", i, 0, obs_time)

        # WRF/LASSO environmental
        wrf_low_cf = extract_lasso_to_obs_time("low_cloud_fraction_arscl", i, 1, obs_time)
        wrf_lwp    = extract_lasso_to_obs_time("lwp", i, 1, obs_time)

        wrf_T2m  = extract_lasso_to_obs_time("temperature_surface", i, 1, obs_time)
        wrf_qv2m = extract_lasso_to_obs_time("water_vapor_mixing_ratio_surface", i, 1, obs_time, qv=True)
        wrf_rh2m = extract_lasso_to_obs_time("rh_surface", i, 1, obs_time)

        wrf_T_bl  = extract_lasso_to_obs_time("temperature_boundary_layer", i, 1, obs_time)
        wrf_qv_bl = extract_lasso_to_obs_time("water_vapor_mixing_ratio_boundary_layer", i, 1, obs_time, qv=True)
        wrf_rh_bl = extract_lasso_to_obs_time("rh_boundary_layer", i, 1, obs_time)

        wrf_lcl = extract_lasso_to_obs_time("lcl", i, 1, obs_time)
        wrf_cbh = extract_lasso_to_obs_time("cloud_base_height", i, 1, obs_time)

        # DP-SCREAM
        dp_low_cf_v = extract_dp_to_obs_time(dp_low_cf_da, i, obs_time)
        dp_lwp_v    = extract_dp_to_obs_time(dp_lwp_da, i, obs_time)

        dp_T2m_v  = extract_dp_to_obs_time(dp_T2m_da, i, obs_time)
        dp_qv2m_v = extract_dp_to_obs_time(dp_qv2m_da, i, obs_time)
        dp_rh2m_v = extract_dp_to_obs_time(dp_rh_surface_da, i, obs_time)
        dp_pblh_v = extract_dp_to_obs_time(dp_pblh_da, i, obs_time)

        dp_T_bl_v  = extract_dp_to_obs_time(dp_T_bl_da, i, obs_time)
        dp_qv_bl_v = extract_dp_to_obs_time(dp_qv_bl_da, i, obs_time)
        dp_rh_bl_v = extract_dp_to_obs_time(dp_rh_bl_da, i, obs_time)

        dp_real_water = extract_dp_to_obs_time(real_water_fraction_da, i, obs_time)
        dp_max_rh     = extract_dp_to_obs_time(max_RH_below_3km_da, i, obs_time)

        dp_shf = extract_dp_to_obs_time(sensible_heat_flux_da, i, obs_time)
        dp_lhf = extract_dp_to_obs_time(latent_heat_flux_da, i, obs_time)

        dp_shoc_T = extract_dp_to_obs_time(shoc_T_tend_da, i, obs_time) * 86400.0
        dp_p3_T   = extract_dp_to_obs_time(p3_T_tend_da, i, obs_time) * 86400.0
        dp_rad_T  = extract_dp_to_obs_time(rad_T_tend_da, i, obs_time) * 86400.0

        dp_shoc_qv = extract_dp_to_obs_time(shoc_qv_tend_da, i, obs_time) * 86400.0 * 1000.0
        dp_p3_qv   = extract_dp_to_obs_time(p3_qv_tend_da, i, obs_time) * 86400.0 * 1000.0

        dp_tke  = extract_dp_to_obs_time(tke_da, i, obs_time)
        dp_eddy = extract_dp_to_obs_time(eddy_diff_mom_da, i, obs_time)

        for j in range(n):
            base = {
                "case_index": i,
                "case": CASE_LABELS[i],
                "time_utc": obs_time[j],
                "time_lst": obs_time_lst[j],
                "hour_lst": obs_time_lst[j].hour + obs_time_lst[j].minute / 60.0,
                "obs_low_cf": obs_low_cf[j],
            }

            # -------------------------
            # WRF/LASSO row
            # -------------------------
            wrf_cf_bias = wrf_low_cf[j] - obs_low_cf[j]

            rows.append({
                **base,
                "model": "LASSO",
                "model_low_cf": wrf_low_cf[j],
                "model_lwp": wrf_lwp[j],
                "cf_bias": wrf_cf_bias,
                "lwp_bias": wrf_lwp[j] - obs_lwp[j],

                # Environmental biases (shared core)
                "T2m_bias":      wrf_T2m[j]  - obs_T2m[j],
                "qv2m_bias_gkg": wrf_qv2m[j] - obs_qv2m[j],
                "rh2m_bias":     wrf_rh2m[j] - obs_rh2m[j],
                "T_bl_bias":     wrf_T_bl[j]  - obs_T_bl[j],
                "qv_bl_bias_gkg": wrf_qv_bl[j] - obs_qv_bl[j],
                "rh_bl_bias":    wrf_rh_bl[j] - obs_rh_bl[j],

                # WRF-only
                "lcl_bias": wrf_lcl[j] - obs_lcl[j],
                "cbh_bias": wrf_cbh[j] - obs_cbh[j],

                # DP-only -> NaN
                "pblh": np.nan, "pblh_minus_lcl": np.nan, "max_RH_below_3km": np.nan,
                "real_water_fraction": np.nan,
                "sensible_heat_flux": np.nan, "latent_heat_flux": np.nan,
                "shoc_T_tend_Kday": np.nan, "p3_T_tend_Kday": np.nan, "rad_T_tend_Kday": np.nan,
                "shoc_qv_tend_gkgday": np.nan, "p3_qv_tend_gkgday": np.nan,
                "tke": np.nan, "eddy_diff_mom": np.nan,
            })

            # -------------------------
            # DP-SCREAM row
            # -------------------------
            dp_cf_bias = dp_low_cf_v[j] - obs_low_cf[j]

            rows.append({
                **base,
                "model": "DP-SCREAM",
                "model_low_cf": dp_low_cf_v[j],
                "model_lwp": dp_lwp_v[j],
                "cf_bias": dp_cf_bias,
                "lwp_bias": dp_lwp_v[j] - obs_lwp[j],

                # Environmental biases (shared core)
                "T2m_bias":      dp_T2m_v[j]  - obs_T2m[j],
                "qv2m_bias_gkg": dp_qv2m_v[j] - obs_qv2m[j],
                "rh2m_bias":     dp_rh2m_v[j] - obs_rh2m[j],
                "T_bl_bias":     dp_T_bl_v[j]  - obs_T_bl[j],
                "qv_bl_bias_gkg": dp_qv_bl_v[j] - obs_qv_bl[j],
                "rh_bl_bias":    dp_rh_bl_v[j] - obs_rh_bl[j],

                # WRF-only -> NaN
                "lcl_bias": np.nan,
                "cbh_bias": np.nan,

                # DP-only process diagnostics
                "pblh": dp_pblh_v[j],
                "pblh_minus_lcl": dp_pblh_v[j] - obs_lcl[j],
                "max_RH_below_3km": dp_max_rh[j],
                "real_water_fraction": dp_real_water[j],
                "sensible_heat_flux": dp_shf[j],
                "latent_heat_flux": dp_lhf[j],
                "shoc_T_tend_Kday": dp_shoc_T[j],
                "p3_T_tend_Kday": dp_p3_T[j],
                "rad_T_tend_Kday": dp_rad_T[j],
                "shoc_qv_tend_gkgday": dp_shoc_qv[j],
                "p3_qv_tend_gkgday": dp_p3_qv[j],
                "tke": dp_tke[j],
                "eddy_diff_mom": dp_eddy[j],
            })

    return pd.DataFrame(rows)


ml_compare_df = build_ml_compare_df()

print("\n" + "=" * 70)
print(f"Built ml_compare_df: {len(ml_compare_df)} rows")
print("=" * 70)
print(ml_compare_df["model"].value_counts())
ml_compare_df.to_csv(OUTDIR / "ml_model_comparison_bias_table.csv", index=False)


# ============================================================
# 6. FEATURE NAME LOOKUP
# ============================================================

feature_long_names = {
    # Shared environmental bias
    "T2m_bias":        "2-m temperature bias [K]",
    "qv2m_bias_gkg":   "2-m humidity bias [g kg$^{-1}$]",
    "rh2m_bias":       "2-m relative humidity bias [%]",
    "T_bl_bias":       "BL temperature bias [K]",
    "qv_bl_bias_gkg":  "BL humidity bias [g kg$^{-1}$]",
    "rh_bl_bias":      "BL relative humidity bias [%]",

    # WRF-only
    "lcl_bias":        "LCL bias [m]",
    "cbh_bias":        "Cloud-base-height bias [m]",

    # DP-only
    "pblh":            "PBL height [m]",
    "pblh_minus_lcl":  "PBL height − observed LCL [m]",
    "max_RH_below_3km":"Max RH below 3 km [%]",
    "real_water_fraction": "Real-water cloud occurrence",
    "sensible_heat_flux": "Sensible heat flux [W m$^{-2}$]",
    "latent_heat_flux":   "Latent heat flux [W m$^{-2}$]",
    "shoc_T_tend_Kday":   "SHOC T tendency [K day$^{-1}$]",
    "p3_T_tend_Kday":     "P3 T tendency [K day$^{-1}$]",
    "rad_T_tend_Kday":    "Radiation T tendency [K day$^{-1}$]",
    "shoc_qv_tend_gkgday":"SHOC qv tendency [g kg$^{-1}$ day$^{-1}$]",
    "p3_qv_tend_gkgday":  "P3 qv tendency [g kg$^{-1}$ day$^{-1}$]",
    "tke":             "TKE [m$^2$ s$^{-2}$]",
    "eddy_diff_mom":   "Eddy diff. for momentum [m$^2$ s$^{-1}$]",
}


# ============================================================
# 7. CURATED FEATURE SETS
# ============================================================
# Design notes:
#  - SHARED features are identical between WRF and DP-SCREAM so the
#    importance ranking can be directly compared.
#  - WRF gets LCL/CBH biases additionally (DP has no LCL/CBH variables).
#  - DP gets process diagnostics that WRF doesn't expose.
#  - Excluded on purpose:
#      * model_T2m, model_qv2m, model_T_bl, model_qv_bl (state-vs-bias
#        redundancy with the *_bias features)
#      * qc_mean_gkg, qr_mean_gkg, qc_occurrence (mechanically too close
#        to cloud-fraction definition; would be tautological)
#      * model_low_cf, model_lwp (part of the target)

SHARED_FEATURES = [
    "T2m_bias",
    "qv2m_bias_gkg",
    "rh2m_bias",
    "T_bl_bias",
    "qv_bl_bias_gkg",
    "rh_bl_bias",
]

WRF_EXTRA_FEATURES = [
    "lcl_bias",
    "cbh_bias",
]

DP_EXTRA_FEATURES = [
    "pblh",
    "pblh_minus_lcl",
    "max_RH_below_3km",
    "real_water_fraction",
    "sensible_heat_flux",
    "latent_heat_flux",
    "shoc_T_tend_Kday",
    "p3_T_tend_Kday",
    "rad_T_tend_Kday",
    "shoc_qv_tend_gkgday",
    "p3_qv_tend_gkgday",
    "tke",
    "eddy_diff_mom",
]

wrf_feature_candidates    = SHARED_FEATURES + WRF_EXTRA_FEATURES
scream_feature_candidates = SHARED_FEATURES + DP_EXTRA_FEATURES


# ============================================================
# 8. TRAIN ONE MODEL
# ============================================================

def train_one_model_ml(
    df,
    model_name,
    feature_candidates,
    target_col="cf_bias",
    save_prefix=None,
    save=True,
):
    if save_prefix is None:
        save_prefix = model_name.replace("-", "").replace("/", "_").replace(" ", "_")

    sub = df[df["model"] == model_name].copy()

    # Drop features with no usable data
    features = [
        c for c in feature_candidates
        if c in sub.columns and sub[c].notna().sum() >= 3
    ]

    use = sub[
        features + [target_col, "case_index", "case", "time_lst", "model"]
    ].replace([np.inf, -np.inf], np.nan)
    use = use[np.isfinite(use[target_col])].copy()

    print("\n" + "=" * 95)
    print(f"Training ML for: {model_name}")
    print("=" * 95)
    print(f"Target: {target_col}")
    print(f"Rows:   {len(use)}")
    print("Features:")
    for c in features:
        print(f"  {c:24s} non-NaN = {use[c].notna().sum()} / {len(use)}")

    if len(use) < 8 or len(features) == 0:
        print(f"Not enough data for {model_name}.")
        return None

    X = use[features]
    y = use[target_col].values
    groups = use["case_index"].values

    # Equal weighting per case so one long case doesn't dominate
    case_counts = use["case_index"].value_counts().to_dict()
    sample_weight = use["case_index"].map(lambda c: 1.0 / case_counts[c]).values
    sample_weight = sample_weight / sample_weight.mean()  # normalize

    rf = Pipeline([
        ("imputer", SimpleImputer(strategy="median")),
        ("model", RandomForestRegressor(
            n_estimators=N_ESTIMATORS,
            random_state=RANDOM_STATE,
            max_depth=4,
            min_samples_leaf=3,
        )),
    ])

    # --- Leave-one-case-out CV with skill score ---
    logo = LeaveOneGroupOut()
    cv_rows = []
    for train_idx, test_idx in logo.split(X, y, groups):
        rf.fit(
            X.iloc[train_idx], y[train_idx],
            model__sample_weight=sample_weight[train_idx],
        )
        pred = rf.predict(X.iloc[test_idx])
        train_mean = np.average(y[train_idx], weights=sample_weight[train_idx])
        baseline_pred = np.full_like(y[test_idx], train_mean, dtype=float)

        model_rmse = np.sqrt(mean_squared_error(y[test_idx], pred))
        baseline_rmse = np.sqrt(mean_squared_error(y[test_idx], baseline_pred))
        skill = 1.0 - (model_rmse**2 / baseline_rmse**2) if baseline_rmse > 0 else np.nan

        cv_rows.append({
            "model": model_name,
            "test_case": use.iloc[test_idx]["case"].iloc[0],
            "model_rmse": model_rmse,
            "baseline_rmse": baseline_rmse,
            "skill_score": skill,   # >0 means RF beats predicting the mean
            "model_r2": r2_score(y[test_idx], pred),
        })
    cv_df = pd.DataFrame(cv_rows)
    print("\nLeave-one-case-out validation:")
    safe_display(cv_df)

    # --- Final model on all data ---
    rf.fit(X, y, model__sample_weight=sample_weight)
    rf_model = rf.named_steps["model"]

    importance_df = pd.DataFrame({
        "feature": features,
        "importance": rf_model.feature_importances_,
    }).sort_values("importance", ascending=False)
    importance_df["feature_long"] = importance_df["feature"].map(
        feature_long_names).fillna(importance_df["feature"])

    print("\nRF feature importance:")
    safe_display(importance_df)

    # --- Permutation importance ---
    perm = permutation_importance(
        rf, X, y,
        n_repeats=N_REPEATS_PERM,
        random_state=RANDOM_STATE,
        sample_weight=sample_weight,
    )
    perm_df = pd.DataFrame({
        "feature": features,
        "importance_mean": perm.importances_mean,
        "importance_std": perm.importances_std,
    }).sort_values("importance_mean", ascending=False)
    perm_df["feature_long"] = perm_df["feature"].map(
        feature_long_names).fillna(perm_df["feature"])

    print("\nPermutation importance:")
    safe_display(perm_df)

    # --- Save ---
    if save:
        cv_df.to_csv(OUTDIR / f"{save_prefix}_cv_scores.csv", index=False)
        importance_df.to_csv(OUTDIR / f"{save_prefix}_rf_feature_importance.csv", index=False)
        perm_df.to_csv(OUTDIR / f"{save_prefix}_permutation_importance.csv", index=False)

    # --- Plots ---
    for kind, plot_df_full, xlabel, fname in [
        ("rf",   importance_df, "Random Forest feature importance", "rf_feature_importance"),
        ("perm", perm_df,       "Permutation importance",            "permutation_importance"),
    ]:
        fig, ax = plt.subplots(figsize=(10, max(4, 0.4 * len(plot_df_full))))
        plot_df = plot_df_full.iloc[::-1]
        if kind == "perm":
            ax.barh(plot_df["feature_long"], plot_df["importance_mean"],
                    xerr=plot_df["importance_std"])
        else:
            ax.barh(plot_df["feature_long"], plot_df["importance"])
        ax.set_xlabel(xlabel)
        ax.set_title(f"{model_name}: predictors of low-cloud-fraction bias")
        ax.grid(axis="x", alpha=0.25)
        plt.tight_layout()
        if save:
            fig.savefig(OUTDIR / f"{save_prefix}_{fname}.png",
                        dpi=300, bbox_inches="tight")
        plt.show()

    # --- Predicted vs actual ---
    pred = rf.predict(X)
    fig, ax = plt.subplots(figsize=(6, 5))
    for k, case in enumerate(ML_CASE_LABELS):
        idx = use["case"] == case
        if idx.sum() == 0:
            continue
        ax.scatter(y[idx], pred[idx], color=ML_CASE_COLORS[k],
                   s=70, edgecolor="k", linewidth=0.5, label=case)
    lo = np.nanmin([np.nanmin(y), np.nanmin(pred)])
    hi = np.nanmax([np.nanmax(y), np.nanmax(pred)])
    ax.plot([lo, hi], [lo, hi], "k--", lw=1)
    ax.axhline(0, color="gray", lw=0.8)
    ax.axvline(0, color="gray", lw=0.8)
    ax.set_xlabel("Actual low-cloud-fraction bias")
    ax.set_ylabel("Predicted low-cloud-fraction bias")
    ax.set_title(f"{model_name}: predicted vs actual")
    ax.grid(alpha=0.25)
    ax.legend(frameon=False)
    plt.tight_layout()
    if save:
        fig.savefig(OUTDIR / f"{save_prefix}_predicted_vs_actual.png",
                    dpi=300, bbox_inches="tight")
    plt.show()

    return {
        "model_name": model_name,
        "data": use,
        "features": features,
        "rf": rf,
        "cv": cv_df,
        "importance": importance_df,
        "permutation": perm_df,
    }


# ============================================================
# 9. TRAIN BOTH MODELS
# ============================================================

wrf_ml_result = train_one_model_ml(
    df=ml_compare_df,
    model_name="LASSO",
    feature_candidates=wrf_feature_candidates,
    target_col=TARGET_COL,
    save_prefix="ml_wrf_lasso",
    save=True,
)

scream_ml_result = train_one_model_ml(
    df=ml_compare_df,
    model_name="DP-SCREAM",
    feature_candidates=scream_feature_candidates,
    target_col=TARGET_COL,
    save_prefix="ml_dp_scream",
    save=True,
)


# ============================================================
# 10. SIDE-BY-SIDE COMPARISON PLOT
# ============================================================
# Two-panel plot:
#   Left  panel: shared features only (directly comparable across models)
#   Right panel: each model's full feature set, sorted independently

def compare_two_models_sidebyside(
    result_a, result_b,
    label_a="WRF/LASSO", label_b="DP-SCREAM",
    save=True,
):
    if result_a is None or result_b is None:
        print("Missing model result.")
        return None

    a_perm = result_a["permutation"].set_index("feature")
    b_perm = result_b["permutation"].set_index("feature")

    # ---- Panel 1: SHARED features only ----
    shared = [f for f in SHARED_FEATURES
              if f in a_perm.index and f in b_perm.index]

    shared_df = pd.DataFrame({
        "feature": shared,
        "feature_long": [feature_long_names.get(f, f) for f in shared],
        label_a: [a_perm.loc[f, "importance_mean"] for f in shared],
        label_a + "_std": [a_perm.loc[f, "importance_std"] for f in shared],
        label_b: [b_perm.loc[f, "importance_mean"] for f in shared],
        label_b + "_std": [b_perm.loc[f, "importance_std"] for f in shared],
    })
    # Sort by max across the two models
    shared_df["max_imp"] = shared_df[[label_a, label_b]].max(axis=1)
    shared_df = shared_df.sort_values("max_imp", ascending=True)

    fig, axes = plt.subplots(1, 2, figsize=(14, 6),
                             gridspec_kw={"width_ratios": [1, 1.2]})

    # ---- Left: shared features grouped bars ----
    ax = axes[0]
    y_pos = np.arange(len(shared_df))
    height = 0.4
    ax.barh(y_pos - height/2, shared_df[label_a], height,
            xerr=shared_df[label_a + "_std"], label=label_a, color="#1f77b4")
    ax.barh(y_pos + height/2, shared_df[label_b], height,
            xerr=shared_df[label_b + "_std"], label=label_b, color="#d62728")
    ax.set_yticks(y_pos)
    ax.set_yticklabels(shared_df["feature_long"])
    ax.set_xlabel("Permutation importance")
    ax.set_title("Shared environmental-bias features\n(directly comparable)")
    ax.grid(axis="x", alpha=0.25)
    ax.legend(frameon=False, loc="lower right")
    ax.axvline(0, color="k", lw=0.6)

    # ---- Right: each model's top features, sorted independently ----
    ax = axes[1]

    a_top = result_a["permutation"].head(10).iloc[::-1].copy()
    b_top = result_b["permutation"].head(10).iloc[::-1].copy()

    # Stack them: model A on bottom, model B on top, separated by a gap
    gap = 1
    a_y = np.arange(len(a_top))
    b_y = np.arange(len(b_top)) + len(a_top) + gap

    ax.barh(a_y, a_top["importance_mean"], xerr=a_top["importance_std"],
            color="#1f77b4", label=label_a)
    ax.barh(b_y, b_top["importance_mean"], xerr=b_top["importance_std"],
            color="#d62728", label=label_b)

    yticks = list(a_y) + list(b_y)
    yticklabels = list(a_top["feature_long"]) + list(b_top["feature_long"])
    ax.set_yticks(yticks)
    ax.set_yticklabels(yticklabels)

    # Visual separator
    ax.axhline(len(a_top) + gap/2 - 0.5, color="gray", lw=0.8, ls="--")
    ax.set_xlabel("Permutation importance")
    ax.set_title("Top features per model\n(sorted independently)")
    ax.grid(axis="x", alpha=0.25)
    ax.legend(frameon=False, loc="lower right")
    ax.axvline(0, color="k", lw=0.6)

    plt.tight_layout()

    if save:
        fig.savefig(OUTDIR / "ml_compare_wrf_lasso_vs_dp_scream.png",
                    dpi=300, bbox_inches="tight")
        shared_df.to_csv(OUTDIR / "ml_compare_shared_features.csv", index=False)

    plt.show()

    print("\nShared-feature comparison:")
    safe_display(shared_df.sort_values("max_imp", ascending=False))

    return shared_df


shared_compare = compare_two_models_sidebyside(
    wrf_ml_result, scream_ml_result,
    label_a="WRF/LASSO", label_b="DP-SCREAM",
    save=True,
)


# ============================================================
# 11. SUMMARY
# ============================================================

print("\n" + "=" * 70)
print("Output files saved in:", OUTDIR.resolve())
print("=" * 70)
for f in sorted(OUTDIR.glob("*")):
    print(" ", f.name)