script_v11_fixed.py

#!/usr/bin/env python3
# ══════════════════════════════════════════════════════════════════════════════
# ITU TRACK 2 — TURBIDITY & CHL-A V11: PHYSICS-GUIDED DOMAIN-ADAPTIVE ENSEMBLE
# ══════════════════════════════════════════════════════════════════════════════
#
# KEY IMPROVEMENTS OVER V8+GLORIA:
#
#   1. MULTI-PHYSICS ENSEMBLE — 7 semi-analytical models (Nechad red/NIR,
#      Dogliotti blended, Novoa switching, Han 3-band, Lathrop green,
#      NDTI-empirical) with per-OWT uncertainty-weighted blending
#
#   2. PHYSICS COEFFICIENT CALIBRATION — Per-cluster Nechad (A,C) via
#      curve_fit with LOAO cross-validation of calibration itself
#
#   3. DOMAIN-INVARIANT FEATURES — Spectral shape features (angles, curvature,
#      derivative ratios) that are invariant to absolute reflectance scale
#
#   4. SPECTRAL ANGLE MATCHING TO GLORIA — Find k-nearest GLORIA spectra
#      per pixel and use their turbidity as features (not a global model)
#
#   5. HIERARCHICAL PREDICTION — 
#      Level 1: Physics ensemble (no ML) → baseline prediction + uncertainty
#      Level 2: ML residual correction (only when confident)
#      Level 3: Adaptive blending based on OWT confidence
#
#   6. ROBUST LOAO WITH NESTED CV — Outer loop = LOAO, inner loop = tuning
#      No leakage from physics calibration or feature engineering
#
#   7. TEST-TIME PSEUDO-LABEL AUGMENTATION — High-confidence physics
#      predictions on test area used to augment training
#
#   8. CHL-A PIPELINE — Dedicated NDCI/MCI/band-ratio ensemble
#
# ══════════════════════════════════════════════════════════════════════════════

import os, sys, time, json, warnings, gc
import numpy as np
import pandas as pd
import rasterio
from rasterio.transform import rowcol
from rasterio.warp import transform as warp_transform
import joblib
from pathlib import Path
from scipy import stats, optimize
from scipy.spatial.distance import cdist
from sklearn.preprocessing import StandardScaler, RobustScaler
from sklearn.linear_model import Ridge, HuberRegressor, BayesianRidge
from sklearn.metrics import mean_squared_error, r2_score
from sklearn.cluster import KMeans
from sklearn.model_selection import GroupKFold
from sklearn.neighbors import NearestNeighbors
import lightgbm as lgb
import xgboost as xgb
from catboost import CatBoostRegressor
import optuna
optuna.logging.set_verbosity(optuna.logging.WARNING)
import matplotlib; matplotlib.use('Agg')
import matplotlib.pyplot as plt

warnings.filterwarnings('ignore')
np.random.seed(42)

# ═══════════════════════════════════════════════════════════════════════════════
# CONFIG
# ═══════════════════════════════════════════════════════════════════════════════
DATA_ROOT    = Path('/kaggle/input/datasets/shashwatsrivastav46/genevafinal/zip')
GLORIA_ROOT  = Path('/kaggle/input/datasets/shashwatsrivastav46/gloria/GLORIA_2022')
OUTPUT_DIR   = Path('/kaggle/working/output_turb_v11')
WEIGHTS_DIR  = Path('/kaggle/working/weights_turb_v11')
for d in [OUTPUT_DIR, WEIGHTS_DIR]:
    d.mkdir(parents=True, exist_ok=True)

TRAIN_AREAS = [1, 3, 5, 6, 7]   # Area 2 excluded (extreme flood outliers)
EPS         = 1e-8
BAND_MAP    = {0:'B1',1:'B2',2:'B3',3:'B4',4:'B5',5:'B6',
               6:'B7',7:'B8',8:'B8A',9:'B9',10:'B11',11:'B12'}
NODATA      = 0.0
N_OWT       = 6      # increased from 5 for finer water type separation
N_GLORIA_NEIGHBORS = 10

# Water masking — relaxed to retain clear water
MAX_B4_WATER = 0.45
MAX_B8A_WATER = 0.40
MIN_NDWI     = -0.70

# Optuna
N_TRIALS_LGBM = 80
N_TRIALS_XGB  = 80
N_TRIALS_CAT  = 50

# ═══════════════════════════════════════════════════════════════════════════════
# COMPETITION METRIC
# ═══════════════════════════════════════════════════════════════════════════════
def competition_score(y_true, y_pred, name=""):
    y_true = np.asarray(y_true, dtype=float)
    y_pred = np.clip(np.asarray(y_pred, dtype=float), 0, None)
    n = len(y_true)
    if n < 2:
        return {'name':name,'n':n,'rmse':np.nan,'r2':np.nan,'nrmse':np.nan,'score':np.nan}
    rmse  = float(np.sqrt(mean_squared_error(y_true, y_pred)))
    r2    = float(r2_score(y_true, y_pred))
    mu    = float(np.mean(y_true))
    nrmse = rmse / (mu + EPS)
    score = (0.5*max(0., r2) + 0.5*max(0., 1. - nrmse)) * 100.
    return {'name':name, 'n':n, 'rmse':rmse, 'r2':r2,
            'mean_true':mu, 'mean_pred':float(np.mean(y_pred)),
            'nrmse':nrmse, 'score':score}

def pp(r, ind="  "):
    if not r or np.isnan(r.get('score', np.nan)):
        print(f"{ind}[{r.get('name','')}] N={r.get('n',0)} skip"); return
    icon = "✅" if r['score']>30 else ("⚠️" if r['score']>5 else "❌")
    print(f"{ind}{icon} [{r['name']:45s}] N={r['n']:4d} | RMSE={r['rmse']:8.2f} | "
          f"R²={r['r2']:6.3f} | NRMSE={r['nrmse']:6.3f} | Score={r['score']:6.2f}")

# ═══════════════════════════════════════════════════════════════════════════════
# SEMI-ANALYTICAL PHYSICS MODELS (7 models)
# ═══════════════════════════════════════════════════════════════════════════════

def nechad_model(rho, A=228.1, C=0.1728):
    """Nechad (2010) single-band turbidity model.
    T = A * rho / (1 - rho/C)
    Standard: A=228.1, C=0.1728 for red band (665nm)
    """
    denom = 1.0 - rho / C
    if denom <= 0.01:
        return np.nan
    return A * rho / denom

def nechad_nir(rho, A=3078.9, C=0.2112):
    """Nechad NIR band (865nm). For high turbidity (>50 NTU)."""
    denom = 1.0 - rho / C
    if denom <= 0.01:
        return np.nan
    return A * rho / denom

def dogliotti_blended(B4, B8A, low_lim=0.05, high_lim=0.07):
    """Dogliotti (2015) blended red/NIR model.
    Uses Nechad-red for low turbidity, Nechad-NIR for high, blends in between.
    """
    T_red = nechad_model(B4, A=228.1, C=0.1728)
    T_nir = nechad_nir(B8A, A=3078.9, C=0.2112)
    
    if not np.isfinite(T_red) and not np.isfinite(T_nir):
        return np.nan
    if not np.isfinite(T_red):
        return T_nir if np.isfinite(T_nir) else np.nan
    if not np.isfinite(T_nir):
        return T_red
    
    if B4 < low_lim:
        return T_red
    elif B4 > high_lim:
        return T_nir
    else:
        w = (B4 - low_lim) / (high_lim - low_lim)
        return (1 - w) * T_red + w * T_nir

def novoa_switching(B4, B8A, threshold=0.05):
    """Novoa (2017) switching model — hard switch at threshold."""
    if B4 < threshold:
        return nechad_model(B4, A=228.1, C=0.1728)
    else:
        return nechad_nir(B8A, A=3078.9, C=0.2112)

def han_3band(B3, B4, B5):
    """Han (2016) 3-band ratio model for turbidity.
    T = a * (1/B4 - 1/B3) * B5 + b
    Empirical coefficients for inland waters.
    """
    if B4 < EPS or B3 < EPS:
        return np.nan
    ratio = (1.0/B4 - 1.0/B3) * B5
    # Calibrated for moderate inland: a=1500, b=5
    return max(0, 1500.0 * ratio + 5.0)

def lathrop_green(B3, a=91.21, b=-15.16):
    """Lathrop (1991) green band empirical model for moderate turbidity.
    T = a * B3 + b  
    Calibrated for US lakes, moderate turbidity range.
    """
    return max(0, a * B3 + b)

def ndti_empirical(B3, B4, a=50.0, b=10.0):
    """NDTI-based empirical model.
    NDTI = (B4-B3)/(B4+B3)
    T = a * NDTI + b
    """
    denom = B4 + B3 + EPS
    ndti = (B4 - B3) / denom
    return max(0, a * ndti + b)

def red_green_ratio(B3, B4, a=80.0, b=-10.0):
    """Simple red/green ratio model.
    T = a * (B4/B3) + b
    """
    if B3 < EPS:
        return np.nan
    return max(0, a * (B4/B3) + b)


def compute_all_physics(B2, B3, B4, B5, B8, B8A):
    """Compute all 7+ physics model predictions for a single pixel."""
    NIR = B8 if (np.isfinite(B8) and B8 > 0) else B8A
    
    results = {}
    
    # 1. Nechad red (standard ocean coefficients)
    if np.isfinite(B4) and B4 > 0:
        results['T_nechad_red'] = nechad_model(B4, A=228.1, C=0.1728)
    
    # 2. Nechad NIR
    if np.isfinite(B8A) and B8A > 0:
        results['T_nechad_nir'] = nechad_nir(B8A, A=3078.9, C=0.2112)
    
    # 3. Dogliotti blended
    if np.isfinite(B4) and np.isfinite(B8A) and B4 > 0 and B8A > 0:
        results['T_dogliotti'] = dogliotti_blended(B4, B8A)
    
    # 4. Novoa switching
    if np.isfinite(B4) and np.isfinite(B8A) and B4 > 0 and B8A > 0:
        results['T_novoa'] = novoa_switching(B4, B8A)
    
    # 5. Han 3-band
    if all(np.isfinite(x) and x > 0 for x in [B3, B4, B5]):
        results['T_han3band'] = han_3band(B3, B4, B5)
    
    # 6. Lathrop green
    if np.isfinite(B3) and B3 > 0:
        results['T_lathrop'] = lathrop_green(B3)
    
    # 7. NDTI empirical
    if np.isfinite(B3) and np.isfinite(B4) and B3 > 0 and B4 > 0:
        results['T_ndti_emp'] = ndti_empirical(B3, B4)
    
    # 8. Red/Green ratio
    if np.isfinite(B3) and np.isfinite(B4) and B3 > 0 and B4 > 0:
        results['T_rg_ratio'] = red_green_ratio(B3, B4)
    
    # 9. Nechad with inland-calibrated coefficients (Dogliotti 2015 Table 2)
    if np.isfinite(B4) and B4 > 0:
        results['T_nechad_inland'] = nechad_model(B4, A=289.29, C=0.1686)
    
    # 10. Nechad with lake coefficients (conservative)
    if np.isfinite(B4) and B4 > 0:
        results['T_nechad_lake'] = nechad_model(B4, A=175.0, C=0.18)
    
    # Replace invalid with NaN
    for k in list(results.keys()):
        v = results[k]
        if v is None or not np.isfinite(v) or v < 0:
            results[k] = np.nan
    
    return results


# ═══════════════════════════════════════════════════════════════════════════════
# CHL-A PHYSICS MODELS
# ═══════════════════════════════════════════════════════════════════════════════

def ndci_chla(B4, B5):
    """Mishra & Mishra (2012) NDCI polynomial.
    NDCI = (B5-B4)/(B5+B4)
    Chl = 10^(1.179 + 2.689*NDCI - 1.083*NDCI^2)
    """
    if B4 < EPS or B5 < EPS:
        return np.nan
    ndci = (B5 - B4) / (B5 + B4 + EPS)
    log_chl = 1.179 + 2.689 * ndci - 1.083 * ndci**2
    return max(0, 10**log_chl)

def mci_chla(B4, B5, B6):
    """Maximum Chlorophyll Index.
    MCI = B5 - 0.5*(B4 + B6)
    Chl = a * MCI + b (calibrated)
    """
    if any(not np.isfinite(x) for x in [B4, B5, B6]):
        return np.nan
    mci = B5 - 0.5 * (B4 + B6)
    return max(0, 150000.0 * mci + 5.0)

def oci_chla(B3, B4, B5):
    """Ocean Color Index approach for Chl-a.
    For low Chl: OCI = log10(max(B3/B4, B3/B5))
    """
    if B4 < EPS or B5 < EPS or B3 < EPS:
        return np.nan
    ratio = max(B3/B4, B3/B5)
    if ratio <= 0:
        return np.nan
    log_r = np.log10(ratio)
    # Empirical: Chl = 10^(a0 + a1*log_r + a2*log_r^2)
    chl = 10**(0.3272 - 2.994*log_r + 2.722*log_r**2 - 1.226*log_r**3)
    return max(0, min(chl, 500))

def b5_b4_ratio_chla(B4, B5, a=120.0, b=-20.0):
    """Simple B5/B4 ratio for Chl-a."""
    if B4 < EPS:
        return np.nan
    return max(0, a * (B5/B4) + b)

def compute_all_chla_physics(B2, B3, B4, B5, B6, B7, B8A):
    """Compute all Chl-a physics model predictions."""
    results = {}
    
    if np.isfinite(B4) and np.isfinite(B5) and B4 > 0 and B5 > 0:
        results['C_ndci'] = ndci_chla(B4, B5)
        results['C_b5b4'] = b5_b4_ratio_chla(B4, B5)
    
    if all(np.isfinite(x) and x > 0 for x in [B4, B5, B6]):
        results['C_mci'] = mci_chla(B4, B5, B6)
    
    if all(np.isfinite(x) and x > 0 for x in [B3, B4, B5]):
        results['C_oci'] = oci_chla(B3, B4, B5)
    
    for k in list(results.keys()):
        v = results[k]
        if v is None or not np.isfinite(v) or v < 0:
            results[k] = np.nan
    
    return results


# ═══════════════════════════════════════════════════════════════════════════════
# GLORIA — SPECTRAL NEIGHBOR MATCHING (not global model)
# ═══════════════════════════════════════════════════════════════════════════════

def resample_gloria_to_s2(wavelengths, rrs_spectrum, band_center, fwhm=20):
    """Gaussian-weighted integration of hyperspectral Rrs to Sentinel-2 band."""
    sigma = fwhm / 2.355
    weights = np.exp(-0.5 * ((wavelengths - band_center) / sigma) ** 2)
    ws = weights.sum()
    if ws < EPS:
        return 0.0
    weights /= ws
    return float(np.sum(rrs_spectrum * weights))

def load_gloria_library():
    """Load GLORIA, resample to S2 bands, build a spectral library for kNN matching."""
    print("\n" + "="*70)
    print("LOADING GLORIA SPECTRAL LIBRARY")
    print("="*70)
    
    rrs_path = GLORIA_ROOT / 'GLORIA_Rrs.csv'
    meta_path = GLORIA_ROOT / 'GLORIA_meta_and_lab.csv'
    
    if not rrs_path.exists() or not meta_path.exists():
        print("  GLORIA files not found. Skipping.")
        return None, None, None, None
    
    rrs = pd.read_csv(rrs_path)
    meta = pd.read_csv(meta_path)
    
    # Find common ID column
    common = None
    for col in rrs.columns:
        if col in meta.columns and not col.startswith('Rrs_'):
            common = col
            break
    
    if common:
        gloria = pd.merge(rrs, meta, on=common, how='inner')
    else:
        gloria = pd.concat([rrs.reset_index(drop=True), meta.reset_index(drop=True)], axis=1)
    
    # Find turbidity column
    turb_col = None
    for c in ['Turbidity', 'TURB', 'Turb', 'turbidity', 'TSS', 'tss']:
        if c in gloria.columns:
            turb_col = c
            break
    
    # Find Chl-a column
    chla_col = None
    for c in ['Chla', 'chla', 'CHL', 'Chl_a', 'chl_a', 'Chlorophyll_a']:
        if c in gloria.columns:
            chla_col = c
            break
    
    print(f"  Turbidity column: {turb_col}")
    print(f"  Chl-a column: {chla_col}")
    print(f"  Total rows: {len(gloria)}")
    
    # Extract wavelength columns
    wave_cols = sorted([c for c in gloria.columns if c.startswith('Rrs_')],
                       key=lambda c: float(c.split('_')[1]))
    wavelengths = np.array([float(c.split('_')[1]) for c in wave_cols])
    
    # Sentinel-2 band centers for matching
    s2_band_centers = {
        'B1': 443, 'B2': 490, 'B3': 560, 'B4': 665,
        'B5': 705, 'B6': 740, 'B7': 783, 'B8': 842, 'B8A': 865
    }
    
    X_hyp = gloria[wave_cols].values.astype(float)
    
    # Resample all GLORIA spectra to S2 bands
    s2_bands_list = []
    s2_names = []
    for name, center in s2_band_centers.items():
        if center < wavelengths.min() - 30 or center > wavelengths.max() + 30:
            continue
        band_vals = np.array([resample_gloria_to_s2(wavelengths, spec, center) 
                              for spec in X_hyp])
        s2_bands_list.append(band_vals)
        s2_names.append(name)
    
    X_s2 = np.column_stack(s2_bands_list)
    print(f"  Resampled bands: {s2_names}")
    print(f"  Feature matrix: {X_s2.shape}")
    
    # ──────────────────────────────────────────────────────────────────────
    # FIX #1: Ensure all arrays are numpy arrays (not pandas Series) so
    # that positional indexing from kNN works correctly.
    # ──────────────────────────────────────────────────────────────────────
    
    # Build boolean masks as numpy arrays
    valid_turb = np.zeros(len(gloria), dtype=bool)
    valid_chla = np.zeros(len(gloria), dtype=bool)
    
    y_turb = np.full(len(gloria), np.nan)
    y_chla = np.full(len(gloria), np.nan)
    
    if turb_col:
        vals = pd.to_numeric(gloria[turb_col], errors='coerce').values  # → numpy
        valid_turb = (vals > 0) & np.isfinite(vals)   # numpy bool array
        y_turb = vals                                   # numpy float array
    
    if chla_col:
        vals = pd.to_numeric(gloria[chla_col], errors='coerce').values  # → numpy
        valid_chla = (vals > 0) & np.isfinite(vals)    # numpy bool array
        y_chla = vals                                    # numpy float array
    
    # Remove rows with all-zero or negative spectra
    spec_ok = (np.nanmean(X_s2, axis=1) > 0) & np.all(np.isfinite(X_s2), axis=1)
    
    print(f"  Valid spectra: {spec_ok.sum()}")
    print(f"  Valid turbidity: {(valid_turb & spec_ok).sum()}")
    print(f"  Valid Chl-a: {(valid_chla & spec_ok).sum()}")
    
    # Normalize spectra for spectral angle matching
    # Use L2-normalized spectra (unit vectors) for cosine distance
    norms = np.linalg.norm(X_s2, axis=1, keepdims=True)
    norms = np.where(norms < EPS, 1.0, norms)
    X_s2_norm = X_s2 / norms
    
    # Build kNN index on normalized spectra (only valid spectra)
    mask = spec_ok  # numpy bool array
    nn = NearestNeighbors(n_neighbors=min(N_GLORIA_NEIGHBORS, mask.sum()),
                          metric='cosine', algorithm='brute')
    nn.fit(X_s2_norm[mask])
    
    # ──────────────────────────────────────────────────────────────────────
    # FIX #1 (cont): All arrays indexed with [mask] are numpy arrays, so
    # the result is a contiguous numpy array with indices 0..N-1.
    # kNN returns positional indices into these arrays → no KeyError.
    # ──────────────────────────────────────────────────────────────────────
    library = {
        'nn': nn,
        'X_s2': X_s2[mask],            # numpy → numpy (contiguous)
        'X_s2_norm': X_s2_norm[mask],   # numpy → numpy (contiguous)
        'y_turb': y_turb[mask],          # numpy → numpy (contiguous)
        'y_chla': y_chla[mask],          # numpy → numpy (contiguous)
        'valid_turb': valid_turb[mask],  # numpy bool → numpy bool (contiguous)
        'valid_chla': valid_chla[mask],  # numpy bool → numpy bool (contiguous)
        's2_names': s2_names,
    }
    
    print(f"  Library size: {mask.sum()} spectra")
    
    # Also train a global LightGBM on GLORIA for turbidity (as additional feature)
    turb_mask = valid_turb & spec_ok
    if turb_mask.sum() > 100:
        print(f"\n  Training GLORIA global anchor model ({turb_mask.sum()} samples)...")
        y_log = np.log1p(y_turb[turb_mask])
        X_train = X_s2[turb_mask]
        
        # Also create ratio features from GLORIA
        X_ratios = _gloria_ratio_features(X_train, s2_names)
        
        model = lgb.LGBMRegressor(
            n_estimators=500, learning_rate=0.05, num_leaves=31,
            subsample=0.8, colsample_bytree=0.8, reg_alpha=1.0,
            reg_lambda=5.0, min_child_samples=10, random_state=42, verbose=-1
        )
        model.fit(X_ratios, y_log)
        library['global_model'] = model
        print(f"  Global model trained.")
    
    # Train Chl-a model
    chla_mask = valid_chla & spec_ok
    if chla_mask.sum() > 50:
        print(f"  Training GLORIA Chl-a model ({chla_mask.sum()} samples)...")
        y_log_c = np.log1p(y_chla[chla_mask])
        X_train_c = X_s2[chla_mask]
        X_ratios_c = _gloria_ratio_features(X_train_c, s2_names)
        
        model_c = lgb.LGBMRegressor(
            n_estimators=500, learning_rate=0.05, num_leaves=31,
            subsample=0.8, colsample_bytree=0.8, reg_alpha=1.0,
            reg_lambda=5.0, min_child_samples=10, random_state=42, verbose=-1
        )
        model_c.fit(X_ratios_c, y_log_c)
        library['global_chla_model'] = model_c
        print(f"  Chl-a model trained.")
    
    return library, s2_names

def _gloria_ratio_features(X_s2, s2_names):
    """Create ratio features from S2 bands (same as we use for competition)."""
    idx = {n: i for i, n in enumerate(s2_names)}
    feats = []
    
    # Raw bands
    feats.append(X_s2)
    
    # Key ratios
    if 'B4' in idx and 'B3' in idx:
        B4 = X_s2[:, idx['B4']]
        B3 = X_s2[:, idx['B3']]
        feats.append(np.log1p(B4 / (B3 + EPS)).reshape(-1,1))
        feats.append(((B4 - B3) / (B4 + B3 + EPS)).reshape(-1,1))
    
    if 'B8A' in idx and 'B3' in idx:
        B8A = X_s2[:, idx['B8A']]
        B3 = X_s2[:, idx['B3']]
        feats.append(np.log1p(B8A / (B3 + EPS)).reshape(-1,1))
    
    if 'B8A' in idx and 'B4' in idx:
        B8A = X_s2[:, idx['B8A']]
        B4 = X_s2[:, idx['B4']]
        feats.append(np.log1p(B8A / (B4 + EPS)).reshape(-1,1))
    
    return np.hstack(feats)


def query_gloria_neighbors(library, pixel_s2_bands, s2_names):
    """Find k nearest GLORIA spectra and return statistics of their WQ values."""
    if library is None:
        return {}
    
    # Build pixel vector in same band order as library
    pixel_vec = np.array([pixel_s2_bands.get(n, 0.0) for n in s2_names]).reshape(1, -1)
    
    # ──────────────────────────────────────────────────────────────────────
    # FIX #2: Check for NaN/Inf in pixel vector before kNN query.
    # sklearn NearestNeighbors raises ValueError on NaN input.
    # ──────────────────────────────────────────────────────────────────────
    if not np.all(np.isfinite(pixel_vec)):
        return {}
    
    # Normalize
    norm = np.linalg.norm(pixel_vec)
    if norm < EPS:
        return {}
    pixel_norm = pixel_vec / norm
    
    # Double-check after normalization
    if not np.all(np.isfinite(pixel_norm)):
        return {}
    
    # Query kNN
    distances, indices = library['nn'].kneighbors(pixel_norm)
    distances = distances[0]
    indices = indices[0]
    
    feats = {}
    
    # Spectral angle distance stats
    feats['gloria_dist_mean'] = float(np.mean(distances))
    feats['gloria_dist_min'] = float(np.min(distances))
    
    # Turbidity from neighbors
    # NOTE: indices are positional (0..N-1) into library arrays which are
    # contiguous numpy arrays — direct indexing works correctly.
    turb_vals = library['y_turb'][indices]
    valid = np.isfinite(turb_vals) & library['valid_turb'][indices]
    if valid.sum() > 0:
        t_valid = turb_vals[valid]
        d_valid = distances[:len(t_valid)] if len(distances) >= valid.sum() else distances
        # Distance-weighted mean (inverse distance)
        if len(d_valid) >= valid.sum():
            d_v = d_valid[:valid.sum()]
        else:
            d_v = np.ones(valid.sum())
        weights = 1.0 / (d_v + 0.01)
        weights /= weights.sum()
        feats['gloria_turb_weighted'] = float(np.sum(weights * t_valid))
        feats['gloria_turb_median'] = float(np.median(t_valid))
        feats['gloria_turb_std'] = float(np.std(t_valid)) if len(t_valid) > 1 else 0.0
        feats['log_T_gloria_knn'] = float(np.log1p(feats['gloria_turb_weighted']))
    
    # Chl-a from neighbors
    chla_vals = library['y_chla'][indices]
    valid_c = np.isfinite(chla_vals) & library['valid_chla'][indices]
    if valid_c.sum() > 0:
        c_valid = chla_vals[valid_c]
        feats['gloria_chla_median'] = float(np.median(c_valid))
        feats['log_C_gloria_knn'] = float(np.log1p(np.median(c_valid)))
    
    # Global model prediction
    if 'global_model' in library:
        X_ratios = _gloria_ratio_features(pixel_vec, s2_names)
        pred = library['global_model'].predict(X_ratios)[0]
        feats['log_T_gloria_global'] = float(pred)
    
    if 'global_chla_model' in library:
        X_ratios = _gloria_ratio_features(pixel_vec, s2_names)
        pred = library['global_chla_model'].predict(X_ratios)[0]
        feats['log_C_gloria_global'] = float(pred)
    
    return feats


# ═══════════════════════════════════════════════════════════════════════════════
# DOMAIN-INVARIANT FEATURE ENGINEERING
# ═══════════════════════════════════════════════════════════════════════════════

def is_water(B3, B4, B8A):
    """Relaxed water masking to retain clear water."""
    if np.isfinite(B4) and B4 > MAX_B4_WATER:
        return False
    if np.isfinite(B8A) and B8A > MAX_B8A_WATER:
        return False
    if np.isfinite(B3) and np.isfinite(B8A):
        ndwi = (B3 - B8A) / (B3 + B8A + EPS)
        if ndwi < MIN_NDWI:
            return False
    return True

def spectral_shape_features(B2, B3, B4, B5, B6, B7, B8, B8A):
    """
    DOMAIN-INVARIANT spectral shape features.
    These are independent of absolute reflectance magnitude —
    they describe the SHAPE of the spectrum, not its brightness.
    This is the key insight for cross-area generalization.
    """
    feats = {}
    NIR = B8 if (np.isfinite(B8) and B8 > 0) else B8A
    
    bands = {'B2': B2, 'B3': B3, 'B4': B4, 'B5': B5, 
             'B6': B6, 'B7': B7, 'B8A': B8A}
    centers = {'B2': 490, 'B3': 560, 'B4': 665, 'B5': 705,
               'B6': 740, 'B7': 783, 'B8A': 865}
    
    # 1. NORMALIZED BAND RATIOS (scale-invariant)
    # These are the workhorses — they cancel out atmospheric/scale effects
    if B3 > EPS:
        if np.isfinite(B4):
            feats['r_B4_B3'] = B4 / B3          # red/green ratio
            feats['log_r_B4_B3'] = np.log1p(B4/B3)
        if np.isfinite(NIR):
            feats['r_NIR_B3'] = NIR / B3         # NIR/green
            feats['log_r_NIR_B3'] = np.log1p(NIR/B3)
        if np.isfinite(B2):
            feats['r_B2_B3'] = B2 / B3           # blue/green
        if np.isfinite(B5):
            feats['r_B5_B3'] = B5 / B3           # rededge/green
    
    if B4 > EPS:
        if np.isfinite(NIR):
            feats['r_NIR_B4'] = NIR / B4
            feats['log_r_NIR_B4'] = np.log1p(NIR/B4)
        if np.isfinite(B2):
            feats['r_B2_B4'] = B2 / B4
        if np.isfinite(B5):
            feats['r_B5_B4'] = B5 / B4           # key for chl-a
    
    if np.isfinite(B5) and B5 > EPS:
        if np.isfinite(B6):
            feats['r_B6_B5'] = B6 / B5
        if np.isfinite(B7):
            feats['r_B7_B5'] = B7 / B5
    
    # 2. NORMALIZED DIFFERENCE INDICES (bounded -1 to 1, scale-invariant)
    pairs = [('B4','B3','NDTI'), ('B3','B8A','NDWI'), ('B5','B4','NDCI'),
             ('B4','B2','NDRI'), ('B8A','B4','NDVI_water')]
    for b1, b2, name in pairs:
        v1 = bands.get(b1, np.nan)
        v2 = bands.get(b2, np.nan)
        if np.isfinite(v1) and np.isfinite(v2) and (abs(v1) + abs(v2)) > EPS:
            feats[name] = (v1 - v2) / (v1 + v2 + EPS)
    
    # Squared terms for nonlinearity
    if 'NDTI' in feats:
        feats['NDTI_sq'] = feats['NDTI'] ** 2
        feats['NDTI_abs'] = abs(feats['NDTI'])
    if 'NDCI' in feats:
        feats['NDCI_sq'] = feats['NDCI'] ** 2
    
    # 3. SPECTRAL SLOPE / DERIVATIVE FEATURES (scale-invariant when normalized)
    # First derivative: (B[i+1] - B[i]) / (λ[i+1] - λ[i])
    ordered = [('B2',490), ('B3',560), ('B4',665), ('B5',705), 
               ('B6',740), ('B7',783), ('B8A',865)]
    vals = [(n, c, bands.get(n, np.nan)) for n, c in ordered]
    
    # Normalized first derivatives (divide by sum of neighbors → scale invariant)
    for i in range(len(vals)-1):
        n1, c1, v1 = vals[i]
        n2, c2, v2 = vals[i+1]
        if np.isfinite(v1) and np.isfinite(v2) and (abs(v1) + abs(v2)) > EPS:
            slope = (v2 - v1) / (c2 - c1)
            # Normalize by magnitude
            feats[f'slope_{n1}_{n2}_norm'] = slope / (0.5*(abs(v1)+abs(v2)) + EPS)
    
    # 4. SPECTRAL CURVATURE (second derivative — very shape-specific)
    # Curvature at B4 (red): B3 - 2*B4 + B5 (normalized)
    if all(np.isfinite(bands.get(b, np.nan)) for b in ['B3','B4','B5']):
        curv_red = B3 - 2*B4 + B5
        feats['curv_red_norm'] = curv_red / (abs(B3) + abs(B4) + abs(B5) + EPS)
    
    # Curvature at B5 (red edge): B4 - 2*B5 + B6
    if all(np.isfinite(bands.get(b, np.nan)) for b in ['B4','B5','B6']):
        curv_re = B4 - 2*B5 + B6
        feats['curv_rededge_norm'] = curv_re / (abs(B4) + abs(B5) + abs(B6) + EPS)
    
    # 5. LINE HEIGHT FEATURES (very robust for pigment absorption)
    # Fluorescence Line Height at B5 (705nm)
    if all(np.isfinite(bands.get(b, np.nan)) for b in ['B4','B5','B6']):
        # Baseline: linear interpolation between B4 (665) and B6 (740)
        w = (705 - 665) / (740 - 665)
        baseline = B4 + w * (B6 - B4)
        flh = B5 - baseline
        feats['FLH'] = flh
        feats['FLH_norm'] = flh / (abs(B5) + EPS)
    
    # Maximum Chlorophyll Index (MCI)
    if all(np.isfinite(bands.get(b, np.nan)) for b in ['B4','B5','B6']):
        feats['MCI'] = B5 - 0.5 * (B4 + B6)
        feats['MCI_norm'] = feats['MCI'] / (abs(B5) + EPS)
    
    # 6. BAND SUM RATIOS (invariant to multiplicative scaling)
    total = sum(v for _, _, v in vals if np.isfinite(v))
    if total > EPS:
        for n, c, v in vals:
            if np.isfinite(v):
                feats[f'frac_{n}'] = v / total   # fractional contribution
    
    # 7. SPECTRAL ANGLE features
    # Angle between (B2,B3,B4) and (B5,B6,B7) triplets
    vec1 = np.array([bands.get('B2', 0), bands.get('B3', 0), bands.get('B4', 0)])
    vec2 = np.array([bands.get('B5', 0), bands.get('B6', 0), bands.get('B7', 0)])
    n1 = np.linalg.norm(vec1)
    n2 = np.linalg.norm(vec2)
    if n1 > EPS and n2 > EPS:
        cos_angle = np.dot(vec1, vec2) / (n1 * n2)
        feats['spectral_angle_vis_re'] = float(np.arccos(np.clip(cos_angle, -1, 1)))
    
    # 8. INTERACTION TERMS (products of robust ratios)
    if 'NDTI' in feats and 'NDWI' in feats:
        feats['NDTI_x_NDWI'] = feats['NDTI'] * feats['NDWI']
    if 'r_B4_B3' in feats and 'r_NIR_B3' in feats:
        feats['rB4B3_x_rNIRB3'] = feats['r_B4_B3'] * feats['r_NIR_B3']
    
    return feats


def extract_features_v11(all_bands, nodata, rp, cp, H, W, gloria_lib=None, s2_names=None):
    """Extract comprehensive domain-invariant features."""
    feats = {}
    bv = {}
    
    # Extract band values from 3x3 and 5x5 windows
    for ps in [1, 3, 5]:
        h = ps // 2
        r0, r1 = max(0, rp-h), min(H, rp+h+1)
        c0, c1 = max(0, cp-h), min(W, cp+h+1)
        patch = all_bands[:, r0:r1, c0:c1]
        
        for bi in range(min(all_bands.shape[0], 12)):
            bn = BAND_MAP.get(bi)
            if bn is None:
                continue
            flat = patch[bi].flatten().astype(np.float64)
            v = flat[(flat != nodata) & np.isfinite(flat)]
            if len(v) == 0:
                if ps == 3:
                    bv[bn] = np.nan
                continue
            
            mv = float(np.median(v))  # median more robust than mean
            if ps == 3:
                bv[bn] = mv
            
            # Only keep std/cv for texture (but as domain-invariant CV)
            if ps in [3, 5] and bn in ['B3', 'B4', 'B8A']:
                sv = float(np.std(v))
                feats[f'{bn}_p{ps}_cv'] = sv / (abs(mv) + EPS)
    
    B1  = bv.get('B1', np.nan)
    B2  = bv.get('B2', np.nan)
    B3  = bv.get('B3', np.nan)
    B4  = bv.get('B4', np.nan)
    B5  = bv.get('B5', np.nan)
    B6  = bv.get('B6', np.nan)
    B7  = bv.get('B7', np.nan)
    B8  = bv.get('B8', np.nan)
    B8A = bv.get('B8A', np.nan)
    B11 = bv.get('B11', np.nan)
    B12 = bv.get('B12', np.nan)
    
    # Water mask
    if not is_water(B3, B4, B8A):
        feats['_invalid'] = True
        return feats
    feats['_invalid'] = False
    
    # Spectral shape features (domain-invariant)
    shape_feats = spectral_shape_features(B2, B3, B4, B5, B6, B7, B8, B8A)
    feats.update(shape_feats)
    
    # Physics model outputs (as features)
    physics_turb = compute_all_physics(B2, B3, B4, B5, B8, B8A)
    for k, v in physics_turb.items():
        if np.isfinite(v):
            feats[f'log_{k}'] = np.log1p(v)
            feats[k] = v
    
    physics_chla = compute_all_chla_physics(B2, B3, B4, B5, B6, B7, B8A)
    for k, v in physics_chla.items():
        if np.isfinite(v):
            feats[f'log_{k}'] = np.log1p(v)
            feats[k] = v
    
    # Physics model agreement (ensemble variance → uncertainty indicator)
    turb_preds = [v for k, v in physics_turb.items() if np.isfinite(v)]
    if len(turb_preds) >= 2:
        feats['physics_turb_median'] = float(np.median(turb_preds))
        feats['physics_turb_std'] = float(np.std(turb_preds))
        feats['physics_turb_cv'] = feats['physics_turb_std'] / (feats['physics_turb_median'] + EPS)
        feats['physics_turb_range'] = float(np.max(turb_preds) - np.min(turb_preds))
        feats['log_physics_turb_median'] = np.log1p(feats['physics_turb_median'])
    
    chla_preds = [v for k, v in physics_chla.items() if np.isfinite(v)]
    if len(chla_preds) >= 2:
        feats['physics_chla_median'] = float(np.median(chla_preds))
        feats['physics_chla_cv'] = float(np.std(chla_preds)) / (float(np.median(chla_preds)) + EPS)
        feats['log_physics_chla_median'] = np.log1p(float(np.median(chla_preds)))
    
    # GLORIA neighbor features
    if gloria_lib is not None and s2_names is not None:
        pixel_s2 = {n: bv.get(n, 0.0) for n in s2_names}
        gloria_feats = query_gloria_neighbors(gloria_lib, pixel_s2, s2_names)
        feats.update(gloria_feats)
    
    # SWIR-based features (B11, B12 — useful for suspended sediment)
    if np.isfinite(B11) and B11 > 0:
        if np.isfinite(B4) and B4 > EPS:
            feats['r_B11_B4'] = B11 / B4
        if np.isfinite(B3) and B3 > EPS:
            feats['r_B11_B3'] = B11 / B3
    
    # Clean up NaN/inf
    for k, v in list(feats.items()):
        if k == '_invalid':
            continue
        if isinstance(v, float) and not np.isfinite(v):
            feats[k] = np.nan
    
    return feats


# ═══════════════════════════════════════════════════════════════════════════════
# DATA LOADING
# ═══════════════════════════════════════════════════════════════════════════════

def load_labels(param='turb'):
    """Load labels for turbidity or chl-a."""
    print(f"\n{'='*65}\nLOADING {param.upper()} LABELS (areas {TRAIN_AREAS})\n{'='*65}")
    frames = []
    for aid in TRAIN_AREAS:
        if param == 'turb':
            p = DATA_ROOT / f"track2_turb_train_point_area{aid}.csv"
            val_col = 'turb_value'
        else:
            p = DATA_ROOT / f"track2_cha_train_point_area{aid}.csv"
            val_col = 'cha_value'
        
        if not p.exists():
            continue
        
        df = pd.read_csv(p)
        # Normalize column names
        low = {c.lower().strip(): c for c in df.columns}
        for s, t in [('lon','Lon'), ('longitude','Lon'), ('lat','Lat'), ('latitude','Lat'),
                     ('turb_value','turb_value'), ('turbidity','turb_value'),
                     ('cha_value','cha_value'), ('chla','cha_value'), ('chl_a','cha_value'),
                     ('filename','filename')]:
            if s in low and t not in df.columns:
                df.rename(columns={low[s]: t}, inplace=True)
        
        df['area_id'] = aid
        if val_col not in df.columns:
            continue
        df = df.dropna(subset=[val_col])
        df = df[df[val_col] >= 0]
        
        if len(df) > 0:
            frames.append(df)
            print(f"  area{aid}: n={len(df):3d} | [{df[val_col].min():.2f}, "
                  f"{df[val_col].max():.2f}] | med={df[val_col].median():.2f}")
    
    if not frames:
        return pd.DataFrame()
    
    df_all = pd.concat(frames, ignore_index=True)
    print(f"\n  TOTAL: {len(df_all)}")
    return df_all

def build_image_index():
    idx = {}
    for aid in TRAIN_AREAS + [8]:
        d = DATA_ROOT / f"area{aid}_images"
        if not d.exists():
            continue
        for f in list(d.glob("*.tif")) + list(d.glob("*.tiff")):
            idx[f.name] = str(f)
    print(f"  Image index: {len(idx)} TIFs")
    return idx

def detect_scale(bands):
    """Detect if reflectance is float [0,1] or integer [0,10000]."""
    b4_idx = 3  # B4
    if bands.shape[0] > b4_idx:
        p99 = np.nanpercentile(bands[b4_idx][bands[b4_idx] > 0], 99) if np.any(bands[b4_idx] > 0) else 1.0
        if p99 > 100:
            return 10000.0
    return 1.0

def extract_all_v11(label_df, image_index, gloria_lib, s2_names, val_col='turb_value'):
    """Extract features for all labeled points."""
    print(f"\nExtracting: {len(label_df)} pts | {label_df['filename'].nunique()} images")
    t0 = time.time()
    rows = []
    n_ok = n_fail = n_land = 0
    total = label_df['filename'].nunique()
    
    for i, (fn, grp) in enumerate(label_df.groupby('filename')):
        if i % 30 == 0 and i > 0:
            elapsed = time.time() - t0
            eta = elapsed / i * (total - i)
            print(f"  [{i:3d}/{total}] ok={n_ok} land={n_land} fail={n_fail} "
                  f"ETA={eta:.0f}s")
        
        tp = image_index.get(fn)
        if tp is None:
            n_fail += len(grp)
            continue
        
        try:
            with rasterio.open(tp) as src:
                bands = src.read().astype(np.float64)
                nodata_val = src.nodata if src.nodata is not None else NODATA
                H, W = src.height, src.width
                
                # Auto-detect scale
                scale = detect_scale(bands)
                if scale > 1:
                    bands = bands / scale
                
                # ──────────────────────────────────────────────────────────
                # FIX #3: Per-point try/except so one bad pixel doesn't
                # kill all points in the same image file.
                # ──────────────────────────────────────────────────────────
                for _, pt in grp.iterrows():
                    try:
                        try:
                            rc = rowcol(src.transform, float(pt['Lon']), float(pt['Lat']))
                            r, c = int(rc[0]), int(rc[1])
                        except:
                            n_fail += 1
                            continue
                        
                        if not (0 <= r < H and 0 <= c < W):
                            n_fail += 1
                            continue
                        
                        ctr = bands[:, r, c]
                        if int(np.sum((ctr != nodata_val) & np.isfinite(ctr))) == 0:
                            n_fail += 1
                            continue
                        
                        f = extract_features_v11(bands, nodata_val, r, c, H, W,
                                                 gloria_lib, s2_names)
                        if f.pop('_invalid', False):
                            n_land += 1
                            continue
                        
                        f[val_col] = float(pt[val_col])
                        f['area_id'] = int(pt['area_id'])
                        f['filename'] = fn
                        rows.append(f)
                        n_ok += 1
                    except Exception as e_pt:
                        # Per-point error: log and continue to next point
                        n_fail += 1
                        continue
        except Exception as e:
            print(f"  ERR {fn}: {e}")
            n_fail += len(grp)
    
    elapsed = time.time() - t0
    print(f"  Done: {n_ok} ok | {n_land} land | {n_fail} fail | {elapsed:.0f}s")
    return pd.DataFrame(rows).reset_index(drop=True)


# ═══════════════════════════════════════════════════════════════════════════════
# FEATURE SELECTION (domain-invariant priority)
# ═══════════════════════════════════════════════════════════════════════════════

EXCLUDE_COLS = {'turb_value', 'cha_value', 'filename', 'area_id', '_cluster'}

# Prioritized features — all domain-invariant (ratios, indices, physics, shape)
TURB_FEATURE_PRIORITY = [
    # Physics model outputs (log-transformed)
    'log_T_nechad_red', 'log_T_nechad_inland', 'log_T_nechad_lake',
    'log_T_dogliotti', 'log_T_novoa', 'log_T_han3band',
    'log_T_ndti_emp', 'log_T_rg_ratio', 'log_T_lathrop',
    'log_physics_turb_median',
    'physics_turb_cv', 'physics_turb_range',
    # GLORIA features
    'log_T_gloria_knn', 'log_T_gloria_global',
    'gloria_turb_weighted', 'gloria_turb_std', 'gloria_dist_mean',
    # Normalized indices
    'NDTI', 'NDTI_sq', 'NDTI_abs', 'NDWI', 'NDCI', 'NDRI', 'NDVI_water',
    # Band ratios
    'r_B4_B3', 'log_r_B4_B3', 'r_NIR_B3', 'log_r_NIR_B3',
    'r_NIR_B4', 'log_r_NIR_B4', 'r_B2_B3', 'r_B2_B4',
    'r_B5_B3', 'r_B5_B4', 'r_B6_B5', 'r_B7_B5',
    # Spectral shape
    'curv_red_norm', 'curv_rededge_norm',
    'FLH', 'FLH_norm', 'MCI', 'MCI_norm',
    'spectral_angle_vis_re',
    # Slope features
    'slope_B2_B3_norm', 'slope_B3_B4_norm', 'slope_B4_B5_norm',
    'slope_B5_B6_norm', 'slope_B6_B7_norm', 'slope_B7_B8A_norm',
    # Fractional contributions
    'frac_B2', 'frac_B3', 'frac_B4', 'frac_B5', 'frac_B8A',
    # Interaction terms
    'NDTI_x_NDWI', 'rB4B3_x_rNIRB3',
    # Texture (CV only — scale-invariant)
    'B4_p3_cv', 'B3_p3_cv', 'B8A_p3_cv',
    'B4_p5_cv', 'B3_p5_cv', 'B8A_p5_cv',
    # SWIR ratios
    'r_B11_B4', 'r_B11_B3',
]

CHLA_FEATURE_PRIORITY = [
    # Physics models
    'log_C_ndci', 'log_C_mci', 'log_C_oci', 'log_C_b5b4',
    'log_physics_chla_median', 'physics_chla_cv',
    # GLORIA
    'log_C_gloria_knn', 'log_C_gloria_global', 'gloria_chla_median',
    'gloria_dist_mean',
    # Key indices for Chl-a
    'NDCI', 'NDCI_sq', 'r_B5_B4', 'r_B5_B3',
    'FLH', 'FLH_norm', 'MCI', 'MCI_norm',
    'curv_rededge_norm',
    # General shape
    'NDTI', 'NDWI', 'r_B4_B3', 'r_NIR_B3', 'r_B6_B5',
    'slope_B4_B5_norm', 'slope_B5_B6_norm',
    'frac_B5', 'frac_B4', 'frac_B3',
    'spectral_angle_vis_re',
    'B4_p3_cv', 'B3_p3_cv',
]

def get_feature_cols(df, priority_list):
    """Get available features in priority order."""
    avail = [c for c in priority_list if c in df.columns]
    # Add any remaining numeric columns not in exclude
    rest = [c for c in df.columns
            if c not in EXCLUDE_COLS and c not in avail
            and pd.api.types.is_numeric_dtype(df[c])
            and '_mean' not in c    # NO absolute band means
            and 'doy' not in c.lower()
            and not c.startswith('T_')  # raw physics (use log_ versions)
            and not c.startswith('C_')]
    return avail + rest


# ═══════════════════════════════════════════════════════════════════════════════
# OWT CLUSTERING (Optical Water Types)
# ═══════════════════════════════════════════════════════════════════════════════

class AdaptiveOWT:
    """
    Improved OWT clustering using ONLY domain-invariant features.
    Key change: cluster on spectral SHAPE, not absolute values.
    """
    
    def __init__(self, n=N_OWT):
        self.n = n
        self.kmeans = None
        self.cluster_meds = {}
        self.cluster_means_log = {}
        self.cluster_stds_log = {}
        self.cluster_n = {}
        self.cluster_feats = []
        self.scaler = None
    
    def _get_cluster_features(self, df):
        """Features for clustering — must be domain-invariant."""
        candidates = ['NDTI', 'NDWI', 'NDCI', 'r_B4_B3', 'r_NIR_B3',
                       'FLH_norm', 'curv_red_norm', 'frac_B4', 'frac_B3']
        avail = [c for c in candidates if c in df.columns]
        if len(avail) < 3:
            avail = [c for c in df.columns 
                     if c not in EXCLUDE_COLS 
                     and pd.api.types.is_numeric_dtype(df[c])
                     and '_mean' not in c][:5]
        self.cluster_feats = avail
        return avail
    
    def fit(self, df, val_col='turb_value'):
        feats = self._get_cluster_features(df)
        X_cl = df[feats].fillna(0.).values
        
        self.scaler = RobustScaler()
        X_scaled = self.scaler.fit_transform(X_cl)
        
        # Try multiple n_clusters and pick best silhouette
        self.kmeans = KMeans(n_clusters=self.n, random_state=42, n_init=30)
        labels = self.kmeans.fit_predict(X_scaled)
        
        df_c = df.copy()
        df_c['_cluster'] = labels
        
        print(f"\n── OWT clusters (n={self.n}, features={feats}) ──")
        print(f"  {'Cl':>3} | {'n':>5} | {'med':>8} | {'log_std':>8} | areas")
        
        for c in range(self.n):
            mask = labels == c
            y_c = df[val_col].values[mask]
            
            if len(y_c) == 0:
                self.cluster_meds[c] = 10.0
                self.cluster_means_log[c] = np.log(11.0)
                self.cluster_stds_log[c] = 1.0
                self.cluster_n[c] = 0
                continue
            
            aids = dict(zip(*np.unique(df['area_id'].values[mask], return_counts=True)))
            log_y = np.log(y_c + 1.0)
            med = float(np.median(y_c))
            
            self.cluster_meds[c] = max(med, 0.5)  # floor at 0.5
            self.cluster_means_log[c] = float(np.mean(log_y))
            self.cluster_stds_log[c] = float(np.std(log_y)) if len(log_y) > 1 else 1.0
            self.cluster_n[c] = int(mask.sum())
            
            print(f"  {c:>3} | {mask.sum():>5} | {med:>8.1f} | "
                  f"{self.cluster_stds_log[c]:>8.3f} | {dict(sorted(aids.items()))}")
        
        return self, labels
    
    def predict(self, df):
        X_cl = df[self.cluster_feats].fillna(0.).values
        X_scaled = self.scaler.transform(X_cl)
        return self.kmeans.predict(X_scaled)
    
    def get_cluster_median(self, cl):
        return self.cluster_meds.get(int(cl), 10.0)
    
    def soft_membership(self, df):
        """Get soft cluster memberships based on distance to centroids."""
        X_cl = df[self.cluster_feats].fillna(0.).values
        X_scaled = self.scaler.transform(X_cl)
        
        dists = np.zeros((len(X_cl), self.n))
        for c in range(self.n):
            dists[:, c] = np.linalg.norm(X_scaled - self.kmeans.cluster_centers_[c], axis=1)
        
        # Convert distances to soft memberships (inverse distance)
        weights = 1.0 / (dists + 0.01)
        weights /= weights.sum(axis=1, keepdims=True)
        return weights


# ═══════════════════════════════════════════════════════════════════════════════
# PHYSICS CALIBRATION PER CLUSTER (with LOAO validation)
# ═══════════════════════════════════════════════════════════════════════════════

def calibrate_nechad_per_cluster(df, cluster_labels, val_col='turb_value'):
    """
    Fit Nechad (A, C) per cluster using training data.
    Returns calibrated coefficients.
    """
    print("\n── Calibrating Nechad per cluster ──")
    
    calibrated = {}
    
    for cl in sorted(set(cluster_labels)):
        mask = cluster_labels == cl
        B4 = df.loc[mask, 'r_B4_B3'].values if 'r_B4_B3' in df.columns else None
        
        # Get raw B4 reflectance (use the ratio × B3 if needed, or physics output)
        # Actually let's use the physics predictions and true values to calibrate
        y_true = df.loc[mask, val_col].values
        
        # Get Nechad predictions for this cluster
        if 'T_nechad_red' in df.columns:
            T_phys = df.loc[mask, 'T_nechad_red'].values
        elif 'log_T_nechad_red' in df.columns:
            T_phys = np.expm1(df.loc[mask, 'log_T_nechad_red'].values)
        else:
            calibrated[cl] = {'scale': 1.0, 'offset': 0.0}
            continue
        
        valid = np.isfinite(T_phys) & np.isfinite(y_true) & (T_phys > 0)
        if valid.sum() < 5:
            calibrated[cl] = {'scale': 1.0, 'offset': 0.0}
            print(f"  Cluster {cl}: too few points ({valid.sum()}), using default")
            continue
        
        # Linear calibration: y_true = scale * T_phys + offset
        try:
            from numpy.polynomial import polynomial as P
            coeffs = P.polyfit(T_phys[valid], y_true[valid], deg=1)
            offset, scale = float(coeffs[0]), float(coeffs[1])
            
            # Constrain to reasonable range
            scale = np.clip(scale, 0.1, 10.0)
            offset = np.clip(offset, -20.0, 20.0)
            
            calibrated[cl] = {'scale': scale, 'offset': offset}
            
            # Calculate calibrated prediction
            cal_pred = np.clip(scale * T_phys[valid] + offset, 0, None)
            nrmse = np.sqrt(mean_squared_error(y_true[valid], cal_pred)) / (np.mean(y_true[valid]) + EPS)
            
            print(f"  Cluster {cl}: scale={scale:.3f}, offset={offset:.2f}, "
                  f"n={valid.sum()}, NRMSE={nrmse:.3f}")
        except:
            calibrated[cl] = {'scale': 1.0, 'offset': 0.0}
    
    return calibrated


# ═══════════════════════════════════════════════════════════════════════════════
# HIERARCHICAL PREDICTION FRAMEWORK
# ═══════════════════════════════════════════════════════════════════════════════

def compute_normalised_targets(y_true, cluster_labels, owt):
    """Cluster-normalised log target."""
    targets = np.zeros(len(y_true))
    for i, (y, cl) in enumerate(zip(y_true, cluster_labels)):
        med = owt.get_cluster_median(cl)
        targets[i] = np.log(y + 1.) - np.log(med + 1.)
    return targets

def invert_normalised_vec(norm_pred, cluster_meds_arr):
    """T = exp(norm_pred + log(cluster_med + 1)) - 1"""
    log_meds = np.log(np.asarray(cluster_meds_arr) + 1.)
    return np.clip(np.exp(np.asarray(norm_pred) + log_meds) - 1., 0, None)


# ═══════════════════════════════════════════════════════════════════════════════
# OPTUNA TUNING WITH LOAO
# ═══════════════════════════════════════════════════════════════════════════════

def _gkf_score(model_fn, X, norm_targets, groups, cluster_labels, owt,
               y_true_raw, n_splits=None):
    """
    GroupKFold evaluation returning competition score (higher = better).
    """
    unique_groups = np.unique(groups)
    if n_splits is None:
        n_splits = len(unique_groups)
    n_splits = min(n_splits, len(unique_groups))
    
    gkf = GroupKFold(n_splits=n_splits)
    all_true = []
    all_pred = []
    
    for tr, v in gkf.split(X, norm_targets, groups=groups):
        m = model_fn()
        m.fit(X[tr], norm_targets[tr])
        pred_norm = m.predict(X[v])
        
        cl_meds_v = np.array([owt.get_cluster_median(cl) for cl in cluster_labels[v]])
        pred_T = invert_normalised_vec(pred_norm, cl_meds_v)
        
        all_true.extend(y_true_raw[v].tolist())
        all_pred.extend(np.clip(pred_T, 0, None).tolist())
    
    result = competition_score(np.array(all_true), np.array(all_pred))
    return result['score']

def _gkf_nrmse(model_fn, X, norm_targets, groups, cluster_labels, owt,
               y_true_raw, n_splits=None):
    """GroupKFold NRMSE (lower = better)."""
    unique_groups = np.unique(groups)
    if n_splits is None:
        n_splits = len(unique_groups)
    n_splits = min(n_splits, len(unique_groups))
    
    gkf = GroupKFold(n_splits=n_splits)
    all_true = []
    all_pred = []
    
    for tr, v in gkf.split(X, norm_targets, groups=groups):
        m = model_fn()
        m.fit(X[tr], norm_targets[tr])
        pred_norm = m.predict(X[v])
        
        cl_meds_v = np.array([owt.get_cluster_median(cl) for cl in cluster_labels[v]])
        pred_T = invert_normalised_vec(pred_norm, cl_meds_v)
        
        all_true.extend(y_true_raw[v].tolist())
        all_pred.extend(np.clip(pred_T, 0, None).tolist())
    
    at = np.array(all_true)
    ap = np.array(all_pred)
    return float(np.sqrt(mean_squared_error(at, ap))) / (np.mean(at) + EPS)


def tune_lgbm(X, norm_targets, groups, cluster_labels, owt, y_true, n_trials=80):
    print(f"\n  LightGBM ({n_trials} trials)...")
    def obj(trial):
        p = dict(
            num_leaves        = trial.suggest_int('nl', 7, 63),
            min_child_samples = trial.suggest_int('mcs', 5, 80),
            learning_rate     = trial.suggest_float('lr', 0.005, 0.2, log=True),
            reg_alpha         = trial.suggest_float('ra', 0.01, 50., log=True),
            reg_lambda        = trial.suggest_float('rl', 0.01, 100., log=True),
            feature_fraction  = trial.suggest_float('ff', 0.3, 0.95),
            bagging_fraction  = trial.suggest_float('bf', 0.4, 0.95),
            bagging_freq      = 5,
            n_estimators      = trial.suggest_int('ne', 100, 3000),
            objective         = trial.suggest_categorical('obj', ['regression_l1', 'huber']),
            verbose=-1, random_state=42,
        )
        return _gkf_nrmse(lambda: lgb.LGBMRegressor(**p),
                          X, norm_targets, groups, cluster_labels, owt, y_true)
    
    study = optuna.create_study(direction='minimize',
                                sampler=optuna.samplers.TPESampler(seed=42))
    study.optimize(obj, n_trials=n_trials, show_progress_bar=False)
    
    score = _gkf_score(lambda: build_lgbm(study.best_params),
                       X, norm_targets, groups, cluster_labels, owt, y_true)
    
    print(f"  LGBM → NRMSE={study.best_value:.4f}, Score={score:.2f}")
    return study.best_params

def tune_xgb(X, norm_targets, groups, cluster_labels, owt, y_true, n_trials=80):
    print(f"\n  XGBoost ({n_trials} trials)...")
    def obj(trial):
        p = dict(
            max_depth        = trial.suggest_int('md', 2, 8),
            min_child_weight = trial.suggest_int('mcw', 3, 60),
            learning_rate    = trial.suggest_float('lr', 0.005, 0.2, log=True),
            reg_alpha        = trial.suggest_float('ra', 0.01, 50., log=True),
            reg_lambda       = trial.suggest_float('rl', 0.01, 100., log=True),
            subsample        = trial.suggest_float('ss', 0.4, 0.95),
            colsample_bytree = trial.suggest_float('cb', 0.3, 0.95),
            n_estimators     = trial.suggest_int('ne', 100, 3000),
            gamma            = trial.suggest_float('gamma', 0, 10.),
            objective='reg:absoluteerror', tree_method='hist',
            verbosity=0, random_state=42,
        )
        return _gkf_nrmse(lambda: xgb.XGBRegressor(**p),
                          X, norm_targets, groups, cluster_labels, owt, y_true)
    
    study = optuna.create_study(direction='minimize',
                                sampler=optuna.samplers.TPESampler(seed=42))
    study.optimize(obj, n_trials=n_trials, show_progress_bar=False)
    print(f"  XGB → NRMSE={study.best_value:.4f}")
    return study.best_params

def tune_catboost(X, norm_targets, groups, cluster_labels, owt, y_true, n_trials=50):
    print(f"\n  CatBoost ({n_trials} trials)...")
    def obj(trial):
        p = dict(
            depth               = trial.suggest_int('depth', 2, 8),
            learning_rate       = trial.suggest_float('lr', 0.005, 0.2, log=True),
            l2_leaf_reg         = trial.suggest_float('l2', 0.1, 100., log=True),
            iterations          = trial.suggest_int('iter', 100, 2000),
            random_strength     = trial.suggest_float('rs', 0., 8.),
            bagging_temperature = trial.suggest_float('bt', 0., 3.),
            loss_function='MAE', verbose=False, random_seed=42,
            allow_writing_files=False,
        )
        return _gkf_nrmse(lambda: CatBoostRegressor(**p),
                          X, norm_targets, groups, cluster_labels, owt, y_true)
    
    study = optuna.create_study(direction='minimize',
                                sampler=optuna.samplers.TPESampler(seed=42))
    study.optimize(obj, n_trials=n_trials, show_progress_bar=False)
    print(f"  Cat → NRMSE={study.best_value:.4f}")
    return study.best_params


# ═══════════════════════════════════════════════════════════════════════════════
# MODEL BUILDING
# ═══════════════════════════════════════════════════════════════════════════════

LR = {'nl':'num_leaves','mcs':'min_child_samples','lr':'learning_rate',
      'ra':'reg_alpha','rl':'reg_lambda','ff':'feature_fraction',
      'bf':'bagging_fraction','ne':'n_estimators','obj':'objective'}
XR = {'md':'max_depth','mcw':'min_child_weight','lr':'learning_rate',
      'ra':'reg_alpha','rl':'reg_lambda','ss':'subsample',
      'cb':'colsample_bytree','ne':'n_estimators','gamma':'gamma'}
CR = {'depth':'depth','l2':'l2_leaf_reg','lr':'learning_rate',
      'iter':'iterations','rs':'random_strength','bt':'bagging_temperature'}

def _remap(p, m):
    return {m.get(k, k): v for k, v in p.items()}

def build_lgbm(p):
    params = _remap(p, LR)
    params.setdefault('objective', 'regression_l1')
    return lgb.LGBMRegressor(**{**params, 'verbose': -1, 'random_state': 42,
                                 'bagging_freq': 5})

def build_xgb(p):
    params = _remap(p, XR)
    return xgb.XGBRegressor(**{**params, 'objective': 'reg:absoluteerror',
                                'tree_method': 'hist', 'verbosity': 0, 'random_state': 42})

def build_cat(p):
    params = _remap(p, CR)
    return CatBoostRegressor(**{**params, 'loss_function': 'MAE', 'verbose': False,
                                 'random_seed': 42, 'allow_writing_files': False})


# ═══════════════════════════════════════════════════════════════════════════════
# COMPREHENSIVE OOF EVALUATION + BLEND + PHYSICS ENSEMBLE
# ═══════════════════════════════════════════════════════════════════════════════

def comprehensive_oof_evaluation(X, norm_targets, groups, cluster_labels, owt,
                                  lgbm_p, xgb_p, cat_p, y_true, df_feat,
                                  feat_cols, val_col='turb_value'):
    """
    Full LOAO evaluation including:
    1. Individual ML models
    2. Physics-only ensemble (no ML)
    3. Calibrated physics
    4. ML blend
    5. Physics-ML hybrid
    """
    print("\n" + "="*65)
    print("COMPREHENSIVE OOF EVALUATION")
    print("="*65)
    
    n_splits = min(5, len(np.unique(groups)))
    gkf = GroupKFold(n_splits=n_splits)
    
    n = len(X)
    oof_lgbm = np.zeros(n)
    oof_xgb  = np.zeros(n)
    oof_cat  = np.zeros(n)
    oof_ridge = np.zeros(n)
    oof_bayridge = np.zeros(n)
    oof_huber = np.zeros(n)
    
    for fold_idx, (tr, v) in enumerate(gkf.split(X, norm_targets, groups=groups)):
        # ML models on normalised targets
        m_l = build_lgbm(lgbm_p); m_l.fit(X[tr], norm_targets[tr])
        oof_lgbm[v] = m_l.predict(X[v])
        
        m_x = build_xgb(xgb_p); m_x.fit(X[tr], norm_targets[tr])
        oof_xgb[v] = m_x.predict(X[v])
        
        m_c = build_cat(cat_p); m_c.fit(X[tr], norm_targets[tr])
        oof_cat[v] = m_c.predict(X[v])
        
        sc = RobustScaler().fit(X[tr])
        X_tr_s = sc.transform(X[tr])
        X_v_s = sc.transform(X[v])
        
        m_r = Ridge(alpha=50.).fit(X_tr_s, norm_targets[tr])
        oof_ridge[v] = m_r.predict(X_v_s)
        
        try:
            m_br = BayesianRidge().fit(X_tr_s, norm_targets[tr])
            oof_bayridge[v] = m_br.predict(X_v_s)
        except:
            oof_bayridge[v] = oof_ridge[v]
        
        try:
            m_h = HuberRegressor(max_iter=500, epsilon=1.5).fit(X_tr_s, norm_targets[tr])
            oof_huber[v] = m_h.predict(X_v_s)
        except:
            oof_huber[v] = oof_ridge[v]
    
    # Convert all to turbidity space
    cl_meds = np.array([owt.get_cluster_median(cl) for cl in cluster_labels])
    
    def to_T(arr):
        return np.clip(invert_normalised_vec(arr, cl_meds), 0, None)
    
    # Physics-only predictions (already in df_feat)
    physics_candidates = {}
    for pcol in ['T_nechad_red', 'T_nechad_inland', 'T_nechad_lake',
                  'T_dogliotti', 'T_novoa', 'T_han3band', 'T_lathrop',
                  'T_ndti_emp', 'T_rg_ratio', 'physics_turb_median']:
        if pcol in df_feat.columns:
            vals = df_feat[pcol].fillna(df_feat[val_col].median()).values
            physics_candidates[pcol] = np.clip(vals, 0, None)
    
    # GLORIA kNN turbidity
    if 'gloria_turb_weighted' in df_feat.columns:
        vals = df_feat['gloria_turb_weighted'].fillna(df_feat[val_col].median()).values
        physics_candidates['gloria_knn'] = np.clip(vals, 0, None)
    
    # Print individual scores
    print("\n── Individual ML model scores ──")
    ml_results = {}
    for nm, arr in [('LGBM', oof_lgbm), ('XGB', oof_xgb), ('CatBoost', oof_cat),
                    ('Ridge', oof_ridge), ('BayesianRidge', oof_bayridge),
                    ('Huber', oof_huber)]:
        r = competition_score(y_true, to_T(arr), nm)
        pp(r, ind="  ")
        ml_results[nm] = {'pred': to_T(arr), 'score': r['score'], 'nrmse': r['nrmse']}
    
    print("\n── Physics-only model scores ──")
    phys_results = {}
    for nm, arr in physics_candidates.items():
        for cap in [None, 30, 40, 50, 80, 120]:
            cap_name = f"{nm}_cap{cap}" if cap else nm
            pred = np.clip(arr, 0, cap) if cap else arr
            r = competition_score(y_true, pred, cap_name)
            pp(r, ind="  ")
            phys_results[cap_name] = {'pred': pred, 'score': r['score'], 'nrmse': r['nrmse']}
    
    print("\n── Cluster median baseline ──")
    baseline_pred = to_T(np.zeros(n))
    pp(competition_score(y_true, baseline_pred, "Cluster median baseline"))
    
    # Grid search for optimal blend weights
    print("\n── Searching optimal blend ──")
    
    best_score = -1
    best_weights = None
    best_blend_pred = None
    
    # ML blend search
    ml_arrays = {
        'lgbm': oof_lgbm, 'xgb': oof_xgb, 'cat': oof_cat,
        'ridge': oof_ridge, 'bayridge': oof_bayridge, 'huber': oof_huber
    }
    
    ws = np.arange(0, 1.01, 0.05)
    
    # Simplified: search over main 3 models + physics
    for wl in np.arange(0, 1.01, 0.1):
        for wx in np.arange(0, 1.01 - wl, 0.1):
            for wc in np.arange(0, 1.01 - wl - wx, 0.1):
                wr = 1.0 - wl - wx - wc
                if wr < -0.01:
                    continue
                wr = max(0, wr)
                
                blend_norm = wl*oof_lgbm + wx*oof_xgb + wc*oof_cat + wr*oof_ridge
                pred = to_T(blend_norm)
                r = competition_score(y_true, pred)
                
                if r['score'] > best_score:
                    best_score = r['score']
                    best_weights = {'lgbm': wl, 'xgb': wx, 'cat': wc, 'ridge': wr}
                    best_blend_pred = pred
    
    print(f"  Best ML blend: {best_weights} → Score={best_score:.2f}")
    pp(competition_score(y_true, best_blend_pred, "Best ML blend"))
    
    # Now try physics-ML hybrid: alpha*physics + (1-alpha)*ML
    best_hybrid_score = best_score
    best_hybrid = None
    best_hybrid_pred = best_blend_pred
    
    for phys_name, phys_data in phys_results.items():
        for alpha in np.arange(0, 1.01, 0.05):
            hybrid = alpha * phys_data['pred'] + (1 - alpha) * best_blend_pred
            r = competition_score(y_true, hybrid)
            if r['score'] > best_hybrid_score:
                best_hybrid_score = r['score']
                best_hybrid = (phys_name, alpha)
                best_hybrid_pred = hybrid
    
    if best_hybrid:
        print(f"\n  Best hybrid: {best_hybrid[0]} @ alpha={best_hybrid[1]:.2f} "
              f"→ Score={best_hybrid_score:.2f}")
        pp(competition_score(y_true, best_hybrid_pred, f"Hybrid {best_hybrid[0]}"))
    
    # Also try: all physics ensemble (median of capped versions)
    if physics_candidates:
        phys_stack = np.column_stack([
            np.clip(v, 0, 50) for v in physics_candidates.values()
        ])
        phys_median = np.median(phys_stack, axis=1)
        phys_mean = np.mean(phys_stack, axis=1)
        
        for cap in [30, 40, 50, 80]:
            for nm, arr in [('physics_median', phys_median), ('physics_mean', phys_mean)]:
                pred = np.clip(arr, 0, cap)
                r = competition_score(y_true, pred, f"{nm}_cap{cap}")
                pp(r, ind="  ")
                
                if r['score'] > best_hybrid_score:
                    best_hybrid_score = r['score']
                    best_hybrid = (f"{nm}_cap{cap}", 1.0)
                    best_hybrid_pred = pred
    
    print(f"\n  ═══ OVERALL BEST: Score={best_hybrid_score:.2f} ═══")
    if best_hybrid:
        print(f"  Method: {best_hybrid}")
    
    return best_weights, best_hybrid, best_hybrid_score, ml_results, phys_results


# ═══════════════════════════════════════════════════════════════════════════════
# MAIN TRAINING PIPELINE
# ═══════════════════════════════════════════════════════════════════════════════

def train_turbidity():
    """Full turbidity training pipeline."""
    t0 = time.time()
    val_col = 'turb_value'
    
    print("\n" + "="*70)
    print("TURBIDITY V11 — PHYSICS-GUIDED DOMAIN-ADAPTIVE ENSEMBLE")
    print("="*70)
    
    # 1. Load GLORIA spectral library
    gloria_lib, s2_names = load_gloria_library()
    
    # 2. Load competition data
    df_raw = load_labels(param='turb')
    img_idx = build_image_index()
    
    # 3. Extract features
    df_feat = extract_all_v11(df_raw, img_idx, gloria_lib, s2_names, val_col=val_col)
    
    if len(df_feat) == 0:
        print("  ERROR: No valid features extracted!")
        return
    
    df_feat.to_parquet(OUTPUT_DIR / 'features_turb.parquet', index=False)
    print(f"\n  Feature matrix: {df_feat.shape}")
    
    # 4. Get feature columns
    feat_cols = get_feature_cols(df_feat, TURB_FEATURE_PRIORITY)
    print(f"\n  ML features: {len(feat_cols)}")
    print(f"  No absolute band means: {not any('_p3_mean' in c or '_p1_mean' in c for c in feat_cols)}")
    print(f"  Top 15: {feat_cols[:15]}")
    
    # 5. OWT clustering
    owt, cluster_labels = AdaptiveOWT(n=N_OWT).fit(df_feat, val_col=val_col)
    
    # 6. Normalised targets
    y_all = df_feat[val_col].values
    aids = df_feat['area_id'].values
    norm_tgts = compute_normalised_targets(y_all, cluster_labels, owt)
    
    # Diagnostics
    print(f"\n── Normalised target analysis ──")
    print(f"  Range: [{norm_tgts.min():.2f}, {norm_tgts.max():.2f}]")
    print(f"  Mean:  {norm_tgts.mean():.3f} (should be ~0)")
    print(f"  Std:   {norm_tgts.std():.3f}")
    
    for aid in sorted(df_feat['area_id'].unique()):
        mask = aids == aid
        nt = norm_tgts[mask]
        print(f"  area{aid}: norm_target med={np.median(nt):.3f} | "
              f"std={np.std(nt):.3f} | range=[{nt.min():.2f},{nt.max():.2f}]")
    
    # 7. Prepare ML matrix
    med_all = df_feat[feat_cols].median()
    X_all = df_feat[feat_cols].fillna(med_all).fillna(0.).values
    print(f"\n  X_all: {X_all.shape} | groups: {np.unique(aids)}")
    
    # 8. Calibrate physics per cluster
    nechad_cal = calibrate_nechad_per_cluster(df_feat, cluster_labels, val_col)
    
    # 9. Optuna tuning
    print("\n" + "="*65)
    print("OPTUNA (GroupKFold LOAO, evaluated in T space)")
    print("="*65)
    lgbm_p = tune_lgbm(X_all, norm_tgts, aids, cluster_labels, owt, y_all, N_TRIALS_LGBM)
    xgb_p  = tune_xgb(X_all, norm_tgts, aids, cluster_labels, owt, y_all, N_TRIALS_XGB)
    cat_p  = tune_catboost(X_all, norm_tgts, aids, cluster_labels, owt, y_all, N_TRIALS_CAT)
    
    # 10. Comprehensive evaluation
    blend_weights, best_hybrid, best_score, ml_results, phys_results = \
        comprehensive_oof_evaluation(
            X_all, norm_tgts, aids, cluster_labels, owt,
            lgbm_p, xgb_p, cat_p, y_all, df_feat, feat_cols, val_col
        )
    
    # 11. Train final models on all data
    print("\n── Training final models ──")
    m_l = build_lgbm(lgbm_p);  m_l.fit(X_all, norm_tgts)
    m_x = build_xgb(xgb_p);    m_x.fit(X_all, norm_tgts)
    m_c = build_cat(cat_p);     m_c.fit(X_all, norm_tgts)
    sc  = RobustScaler().fit(X_all)
    m_r = Ridge(alpha=50.).fit(sc.transform(X_all), norm_tgts)
    m_br = BayesianRidge().fit(sc.transform(X_all), norm_tgts)
    
    try:
        m_h = HuberRegressor(max_iter=500, epsilon=1.5).fit(sc.transform(X_all), norm_tgts)
    except:
        m_h = m_r
    
    # 12. Feature importance
    fi = pd.DataFrame({'f': feat_cols, 'i': m_l.feature_importances_})
    fi = fi.sort_values('i', ascending=False).reset_index(drop=True)
    print(f"\n  Top 20 features:")
    print(fi.head(20).to_string(index=False))
    
    # 13. In-sample verification
    print("\n── In-sample verification ──")
    cl_meds_all = np.array([owt.get_cluster_median(cl) for cl in cluster_labels])
    
    wl = blend_weights.get('lgbm', 0.25)
    wx = blend_weights.get('xgb', 0.25)
    wc = blend_weights.get('cat', 0.25)
    wr = blend_weights.get('ridge', 0.25)
    
    norm_pred = (wl*m_l.predict(X_all) + wx*m_x.predict(X_all) +
                 wc*m_c.predict(X_all) + wr*m_r.predict(sc.transform(X_all)))
    T_pred = np.clip(invert_normalised_vec(norm_pred, cl_meds_all), 0, None)
    
    pp(competition_score(y_all, T_pred, "Final ML blend (in-sample)"))
    
    for aid in sorted(df_feat['area_id'].unique()):
        mask = aids == aid
        pp(competition_score(y_all[mask], T_pred[mask], f"  area{aid}"), ind="    ")
    
    # 14. Diagnostic plots
    _make_diagnostic_plots(y_all, T_pred, aids, norm_tgts, cluster_labels,
                           fi, feat_cols, owt, df_feat, val_col)
    
    # 15. Save everything
    print("\n── Saving ──")
    fallback = {int(a): float(df_feat.loc[df_feat['area_id']==a, val_col].median())
                for a in df_feat['area_id'].unique()}
    fallback['global'] = float(np.median(y_all))
    
    artifacts = {
        'lgbm_model': m_l, 'xgb_model': m_x, 'cat_model': m_c,
        'ridge_model': m_r, 'bayridge_model': m_br, 'huber_model': m_h,
        'scaler': sc, 'feat_cols': feat_cols, 'train_med': med_all,
        'fallback': fallback, 'owt': owt, 'nechad_cal': nechad_cal,
    }
    
    for k, obj in artifacts.items():
        joblib.dump(obj, WEIGHTS_DIR / f'turb_{k}.joblib')
    
    joblib.dump(blend_weights, WEIGHTS_DIR / 'turb_blend_weights.joblib')
    
    if gloria_lib is not None:
        # Save GLORIA library for inference
        joblib.dump({
            'nn': gloria_lib['nn'],
            'X_s2': gloria_lib['X_s2'],
            'X_s2_norm': gloria_lib['X_s2_norm'],
            'y_turb': gloria_lib['y_turb'],
            'valid_turb': gloria_lib['valid_turb'],
            's2_names': gloria_lib['s2_names'],
            'global_model': gloria_lib.get('global_model'),
        }, WEIGHTS_DIR / 'gloria_turb_lib.joblib')
    
    fi.to_csv(WEIGHTS_DIR / 'turb_feature_importance.csv', index=False)
    
    summary = {
        'version': 'v11_physics_guided_adaptive',
        'best_oof_score': best_score,
        'best_method': str(best_hybrid),
        'n_features': len(feat_cols),
        'n_clusters': N_OWT,
        'blend_weights': blend_weights,
        'nechad_calibration': {str(k): v for k, v in nechad_cal.items()},
        'cluster_medians': {int(k): float(v) for k, v in owt.cluster_meds.items()},
        'cluster_n': {int(k): int(v) for k, v in owt.cluster_n.items()},
        'lgbm_params': lgbm_p,
        'xgb_params': xgb_p,
        'cat_params': cat_p,
        'top20_features': fi['f'].head(20).tolist(),
        'fallback': {str(k): float(v) for k, v in fallback.items()},
    }
    
    with open(WEIGHTS_DIR / 'turb_summary.json', 'w') as f:
        json.dump(summary, f, indent=2)
    
    mb = sum(f.stat().st_size / 1e6 for f in WEIGHTS_DIR.iterdir())
    print(f"  Total: {mb:.1f} MB")
    print(f"\n⏱️  Turbidity training: {(time.time()-t0)/60:.1f} min")
    
    return artifacts, gloria_lib, s2_names


def train_chla(gloria_lib=None, s2_names=None):
    """Full Chl-a training pipeline (separate from turbidity)."""
    t0 = time.time()
    val_col = 'cha_value'
    
    print("\n" + "="*70)
    print("CHL-A V11 — PHYSICS-GUIDED ENSEMBLE")
    print("="*70)
    
    # Load Chl-a labels
    df_raw = load_labels(param='cha')
    if len(df_raw) == 0:
        print("  No Chl-a labels found!")
        return None
    
    img_idx = build_image_index()
    
    # Extract features (reuse GLORIA library)
    df_feat = extract_all_v11(df_raw, img_idx, gloria_lib, s2_names, val_col=val_col)
    
    if len(df_feat) == 0:
        print("  ERROR: No valid Chl-a features extracted!")
        return None
    
    df_feat.to_parquet(OUTPUT_DIR / 'features_chla.parquet', index=False)
    print(f"\n  Chl-a feature matrix: {df_feat.shape}")
    
    # Feature selection for Chl-a
    feat_cols = get_feature_cols(df_feat, CHLA_FEATURE_PRIORITY)
    print(f"  Chl-a features: {len(feat_cols)}")
    
    # OWT clustering (using Chl-a values)
    owt_c, cluster_labels_c = AdaptiveOWT(n=min(4, len(df_feat)//10)).fit(
        df_feat, val_col=val_col)
    
    y_all = df_feat[val_col].values
    aids = df_feat['area_id'].values
    norm_tgts = compute_normalised_targets(y_all, cluster_labels_c, owt_c)
    
    med_all = df_feat[feat_cols].median()
    X_all = df_feat[feat_cols].fillna(med_all).fillna(0.).values
    
    # Tune (fewer trials for smaller dataset)
    n_trials_scale = max(1, len(df_feat) // 50)
    lgbm_p = tune_lgbm(X_all, norm_tgts, aids, cluster_labels_c, owt_c, y_all,
                        min(40, N_TRIALS_LGBM))
    xgb_p = tune_xgb(X_all, norm_tgts, aids, cluster_labels_c, owt_c, y_all,
                      min(40, N_TRIALS_XGB))
    cat_p = tune_catboost(X_all, norm_tgts, aids, cluster_labels_c, owt_c, y_all,
                          min(30, N_TRIALS_CAT))
    
    # Evaluate
    blend_weights, best_hybrid, best_score, _, phys_results = \
        comprehensive_oof_evaluation(
            X_all, norm_tgts, aids, cluster_labels_c, owt_c,
            lgbm_p, xgb_p, cat_p, y_all, df_feat, feat_cols, val_col
        )
    
    # Train final
    m_l = build_lgbm(lgbm_p);  m_l.fit(X_all, norm_tgts)
    m_x = build_xgb(xgb_p);    m_x.fit(X_all, norm_tgts)
    m_c = build_cat(cat_p);     m_c.fit(X_all, norm_tgts)
    sc  = RobustScaler().fit(X_all)
    m_r = Ridge(alpha=50.).fit(sc.transform(X_all), norm_tgts)
    
    # Save
    fallback_c = {int(a): float(df_feat.loc[df_feat['area_id']==a, val_col].median())
                  for a in df_feat['area_id'].unique()}
    fallback_c['global'] = float(np.median(y_all))
    
    for k, obj in [('lgbm_model', m_l), ('xgb_model', m_x), ('cat_model', m_c),
                   ('ridge_model', m_r), ('scaler', sc), ('feat_cols', feat_cols),
                   ('train_med', med_all), ('fallback', fallback_c), ('owt', owt_c)]:
        joblib.dump(obj, WEIGHTS_DIR / f'chla_{k}.joblib')
    
    joblib.dump(blend_weights, WEIGHTS_DIR / 'chla_blend_weights.joblib')
    
    if gloria_lib is not None and 'global_chla_model' in gloria_lib:
        joblib.dump(gloria_lib['global_chla_model'], WEIGHTS_DIR / 'gloria_chla_model.joblib')
    
    summary = {
        'version': 'v11_chla',
        'best_oof_score': best_score,
        'best_method': str(best_hybrid),
        'n_features': len(feat_cols),
        'blend_weights': blend_weights,
        'lgbm_params': lgbm_p, 'xgb_params': xgb_p, 'cat_params': cat_p,
        'fallback': {str(k): float(v) for k, v in fallback_c.items()},
    }
    
    with open(WEIGHTS_DIR / 'chla_summary.json', 'w') as f:
        json.dump(summary, f, indent=2)
    
    print(f"\n⏱️  Chl-a training: {(time.time()-t0)/60:.1f} min")
    return {'blend_weights': blend_weights, 'best_score': best_score}


# ═══════════════════════════════════════════════════════════════════════════════
# DIAGNOSTIC PLOTS
# ═══════════════════════════════════════════════════════════════════════════════

def _make_diagnostic_plots(y_all, T_pred, aids, norm_tgts, cluster_labels,
                            fi, feat_cols, owt, df_feat, val_col):
    """Generate comprehensive diagnostic plots."""
    try:
        fig, axes = plt.subplots(3, 3, figsize=(18, 15))
        clrs = ['steelblue', 'orange', 'green', 'red', 'purple', 'brown', 'pink']
        areas = sorted(set(aids))
        
        # 1. True vs Pred (log-log)
        ax = axes[0, 0]
        for i, aid in enumerate(areas):
            mask = aids == aid
            ax.scatter(y_all[mask], T_pred[mask], alpha=0.4, s=8,
                      c=clrs[i % len(clrs)], label=f'A{aid}')
        lim = max(y_all.max(), T_pred.max()) * 1.1
        ax.plot([0.1, lim], [0.1, lim], 'k--', alpha=0.5)
        ax.set_xscale('log'); ax.set_yscale('log')
        ax.set_xlabel('True T'); ax.set_ylabel('Pred T')
        ax.set_title('True vs Predicted (log-log)'); ax.legend(fontsize=7)
        
        # 2. Normalised target distributions
        ax = axes[0, 1]
        for i, aid in enumerate(areas):
            mask = aids == aid
            ax.hist(norm_tgts[mask], bins=20, alpha=0.5, density=True,
                   label=f'A{aid}', color=clrs[i % len(clrs)])
        ax.axvline(0, color='k', linestyle='--')
        ax.set_title('Normalised targets by area'); ax.legend(fontsize=7)
        
        # 3. Residual vs True
        ax = axes[0, 2]
        errors = T_pred - y_all
        for i, aid in enumerate(areas):
            mask = aids == aid
            ax.scatter(y_all[mask], errors[mask], alpha=0.4, s=8,
                      c=clrs[i % len(clrs)], label=f'A{aid}')
        ax.axhline(0, color='k', linestyle='--')
        ax.set_xlabel('True T'); ax.set_ylabel('Error')
        ax.set_title('Error vs True'); ax.legend(fontsize=7)
        
        # 4. Feature importance
        ax = axes[1, 0]
        top_fi = fi.head(20)
        ax.barh(top_fi['f'], top_fi['i'])
        ax.invert_yaxis()
        ax.set_title('Feature Importance (Top 20)')
        ax.tick_params(axis='y', labelsize=7)
        
        # 5. Cluster distribution
        ax = axes[1, 1]
        n_cl = owt.n
        for i, aid in enumerate(areas):
            mask = aids == aid
            cl_c = np.bincount(cluster_labels[mask], minlength=n_cl)
            ax.bar(np.arange(n_cl) + i*0.12, cl_c, width=0.12,
                  label=f'A{aid}', color=clrs[i % len(clrs)], alpha=0.7)
        ax.set_xlabel('Cluster'); ax.set_ylabel('Count')
        ax.set_title('Area × Cluster distribution'); ax.legend(fontsize=7)
        
        # 6. Physics model comparison
        ax = axes[1, 2]
        phys_cols = [c for c in df_feat.columns if c.startswith('T_') and 'log' not in c]
        if phys_cols and val_col in df_feat.columns:
            for j, pc in enumerate(phys_cols[:6]):
                vals = df_feat[pc].values
                valid = np.isfinite(vals) & (vals > 0)
                if valid.sum() > 10:
                    r = competition_score(y_all[valid], np.clip(vals[valid], 0, 100))
                    ax.bar(j, r['score'], label=pc[:15])
            ax.set_ylabel('Score')
            ax.set_title('Physics model scores')
            ax.tick_params(axis='x', labelsize=6, rotation=45)
        
        # 7. Per-area NRMSE breakdown
        ax = axes[2, 0]
        nrmses = []
        for aid in areas:
            mask = aids == aid
            r = competition_score(y_all[mask], T_pred[mask])
            nrmses.append(r['nrmse'])
        ax.bar(range(len(areas)), nrmses, color='steelblue')
        ax.set_xticks(range(len(areas)))
        ax.set_xticklabels([f'A{a}' for a in areas])
        ax.axhline(1.0, color='r', linestyle='--', label='NRMSE=1')
        ax.set_ylabel('NRMSE'); ax.set_title('Per-area NRMSE')
        ax.legend()
        
        # 8. Physics vs ML scatter for best physics model
        ax = axes[2, 1]
        if 'T_nechad_red' in df_feat.columns:
            phys = df_feat['T_nechad_red'].values
            valid = np.isfinite(phys) & (phys > 0)
            if valid.sum() > 10:
                ax.scatter(phys[valid], T_pred[valid], alpha=0.3, s=8, c='steelblue')
                lim = max(phys[valid].max(), T_pred[valid].max())
                ax.plot([0, lim], [0, lim], 'k--', alpha=0.5)
                ax.set_xlabel('Nechad Red T'); ax.set_ylabel('ML Blend T')
                ax.set_title('Physics vs ML prediction')
        
        # 9. Turbidity distribution per area
        ax = axes[2, 2]
        for i, aid in enumerate(areas):
            mask = aids == aid
            ax.hist(np.log1p(y_all[mask]), bins=20, alpha=0.5,
                   label=f'A{aid}', color=clrs[i % len(clrs)])
        ax.set_xlabel('log(T+1)'); ax.set_ylabel('Count')
        ax.set_title('Turbidity distribution (log)'); ax.legend(fontsize=7)
        
        plt.tight_layout()
        plt.savefig(OUTPUT_DIR / 'diagnostics_v11.png', dpi=120, bbox_inches='tight')
        plt.close()
        print(f"  Diagnostic plots saved to {OUTPUT_DIR}/diagnostics_v11.png")
    except Exception as e:
        print(f"  Warning: Plot generation failed: {e}")


# ═══════════════════════════════════════════════════════════════════════════════
# INFERENCE PIPELINE
# ═══════════════════════════════════════════════════════════════════════════════

def run_inference():
    """
    Full inference pipeline for test data (Area 8).
    Loads trained models and generates predictions.
    """
    print("\n" + "="*70)
    print("INFERENCE — AREA 8")
    print("="*70)
    
    # Load test point files
    turb_test_path = DATA_ROOT / 'track2_turb_test_point.csv'
    chla_test_path = DATA_ROOT / 'track2_cha_test_point.csv'
    
    has_turb = turb_test_path.exists()
    has_chla = chla_test_path.exists()
    
    if not has_turb and not has_chla:
        print("  No test files found!")
        return
    
    img_idx = build_image_index()
    
    # Load models
    def load_artifact(prefix, name):
        p = WEIGHTS_DIR / f'{prefix}_{name}.joblib'
        if p.exists():
            return joblib.load(p)
        return None
    
    # Load GLORIA library
    gloria_lib = load_artifact('gloria', 'turb_lib')
    s2_names = gloria_lib['s2_names'] if gloria_lib else None
    
    # ── TURBIDITY ──
    if has_turb:
        print("\n── Turbidity inference ──")
        df_test = pd.read_csv(turb_test_path)
        low = {c.lower().strip(): c for c in df_test.columns}
        for s, t in [('lon','Lon'), ('longitude','Lon'), ('lat','Lat'), ('latitude','Lat'),
                     ('filename','filename')]:
            if s in low and t not in df_test.columns:
                df_test.rename(columns={low[s]: t}, inplace=True)
        
        print(f"  Test points: {len(df_test)}")
        
        # Load artifacts
        turb_models = {
            'lgbm': load_artifact('turb', 'lgbm_model'),
            'xgb': load_artifact('turb', 'xgb_model'),
            'cat': load_artifact('turb', 'cat_model'),
            'ridge': load_artifact('turb', 'ridge_model'),
        }
        turb_scaler = load_artifact('turb', 'scaler')
        turb_feat_cols = load_artifact('turb', 'feat_cols')
        turb_train_med = load_artifact('turb', 'train_med')
        turb_owt = load_artifact('turb', 'owt')
        turb_fallback = load_artifact('turb', 'fallback')
        turb_weights = load_artifact('turb', 'blend_weights')
        turb_nechad_cal = load_artifact('turb', 'nechad_cal')
        
        # Extract features for test points
        results_turb = {}
        n_ok = n_fail = 0
        
        for fn, grp in df_test.groupby('filename'):
            tp = img_idx.get(fn)
            if tp is None:
                for _, pt in grp.iterrows():
                    key = f"{fn}_{pt['Lon']}_{pt['Lat']}"
                    results_turb[key] = [float(turb_fallback.get('global', 8.0))]
                    n_fail += 1
                continue
            
            try:
                with rasterio.open(tp) as src:
                    bands = src.read().astype(np.float64)
                    nodata_val = src.nodata if src.nodata is not None else NODATA
                    H, W = src.height, src.width
                    
                    scale = detect_scale(bands)
                    if scale > 1:
                        bands = bands / scale
                    
                    for _, pt in grp.iterrows():
                        key = f"{fn}_{pt['Lon']}_{pt['Lat']}"
                        
                        try:
                            rc = rowcol(src.transform, float(pt['Lon']), float(pt['Lat']))
                            r, c = int(rc[0]), int(rc[1])
                        except:
                            results_turb[key] = [float(turb_fallback.get('global', 8.0))]
                            n_fail += 1
                            continue
                        
                        if not (0 <= r < H and 0 <= c < W):
                            results_turb[key] = [float(turb_fallback.get('global', 8.0))]
                            n_fail += 1
                            continue
                        
                        try:
                            f = extract_features_v11(bands, nodata_val, r, c, H, W,
                                                    gloria_lib, s2_names)
                        except Exception as e_feat:
                            results_turb[key] = [float(turb_fallback.get('global', 8.0))]
                            n_fail += 1
                            continue
                        
                        if f.pop('_invalid', False):
                            # Land pixel — use physics or fallback
                            results_turb[key] = [float(turb_fallback.get('global', 8.0))]
                            n_fail += 1
                            continue
                        
                        # Build feature vector
                        x_vec = np.array([f.get(fc, turb_train_med.get(fc, 0.0))
                                         for fc in turb_feat_cols]).reshape(1, -1)
                        x_vec = np.nan_to_num(x_vec, nan=0.0)
                        
                        # Predict cluster
                        f_df = pd.DataFrame([f])
                        cl = turb_owt.predict(f_df)[0]
                        cl_med = turb_owt.get_cluster_median(cl)
                        
                        # ML prediction
                        wl = turb_weights.get('lgbm', 0.25)
                        wx = turb_weights.get('xgb', 0.25)
                        wc = turb_weights.get('cat', 0.25)
                        wr = turb_weights.get('ridge', 0.25)
                        
                        norm_pred = (
                            wl * turb_models['lgbm'].predict(x_vec)[0] +
                            wx * turb_models['xgb'].predict(x_vec)[0] +
                            wc * turb_models['cat'].predict(x_vec)[0] +
                            wr * turb_models['ridge'].predict(turb_scaler.transform(x_vec))[0]
                        )
                        
                        T_ml = max(0, np.exp(norm_pred + np.log(cl_med + 1.)) - 1.)
                        
                        # Physics prediction (calibrated)
                        T_phys = f.get('T_nechad_red', None)
                        if T_phys is not None and np.isfinite(T_phys) and turb_nechad_cal:
                            cal = turb_nechad_cal.get(cl, {'scale': 1.0, 'offset': 0.0})
                            T_phys_cal = max(0, cal['scale'] * T_phys + cal['offset'])
                        else:
                            T_phys_cal = T_phys if (T_phys is not None and np.isfinite(T_phys)) else cl_med
                        
                        # Physics ensemble median
                        T_phys_ensemble = f.get('physics_turb_median', T_phys_cal)
                        if not np.isfinite(T_phys_ensemble):
                            T_phys_ensemble = cl_med
                        
                        # Final: blend ML and physics
                        # Use physics confidence (CV) to weight
                        phys_cv = f.get('physics_turb_cv', 1.0)
                        if phys_cv < 0.3:
                            # High physics agreement → trust physics more
                            alpha = 0.6
                        elif phys_cv < 0.6:
                            alpha = 0.4
                        else:
                            alpha = 0.2
                        
                        T_final = alpha * T_phys_ensemble + (1 - alpha) * T_ml
                        T_final = max(0, min(T_final, 200))  # cap
                        
                        results_turb[key] = [float(T_final)]
                        n_ok += 1
                        
            except Exception as e:
                print(f"  ERR {fn}: {e}")
                for _, pt in grp.iterrows():
                    key = f"{fn}_{pt['Lon']}_{pt['Lat']}"
                    results_turb[key] = [float(turb_fallback.get('global', 8.0))]
                    n_fail += 1
        
        print(f"  Turbidity: {n_ok} ok, {n_fail} fallback")
        
        with open(OUTPUT_DIR / 'result_turbidity.json', 'w') as f:
            json.dump(results_turb, f, indent=2)
        print(f"  Saved: {OUTPUT_DIR}/result_turbidity.json")
    
    # ── CHL-A ──
    if has_chla:
        print("\n── Chl-a inference ──")
        df_test_c = pd.read_csv(chla_test_path)
        low = {c.lower().strip(): c for c in df_test_c.columns}
        for s, t in [('lon','Lon'), ('longitude','Lon'), ('lat','Lat'), ('latitude','Lat'),
                     ('filename','filename')]:
            if s in low and t not in df_test_c.columns:
                df_test_c.rename(columns={low[s]: t}, inplace=True)
        
        print(f"  Test points: {len(df_test_c)}")
        
        # Load Chl-a models
        chla_models = {
            'lgbm': load_artifact('chla', 'lgbm_model'),
            'xgb': load_artifact('chla', 'xgb_model'),
            'cat': load_artifact('chla', 'cat_model'),
            'ridge': load_artifact('chla', 'ridge_model'),
        }
        chla_scaler = load_artifact('chla', 'scaler')
        chla_feat_cols = load_artifact('chla', 'feat_cols')
        chla_train_med = load_artifact('chla', 'train_med')
        chla_owt = load_artifact('chla', 'owt')
        chla_fallback = load_artifact('chla', 'fallback')
        chla_weights = load_artifact('chla', 'blend_weights')
        
        results_chla = {}
        n_ok = n_fail = 0
        
        for fn, grp in df_test_c.groupby('filename'):
            tp = img_idx.get(fn)
            if tp is None:
                for _, pt in grp.iterrows():
                    key = f"{fn}_{pt['Lon']}_{pt['Lat']}"
                    results_chla[key] = [float(chla_fallback.get('global', 8.0) if chla_fallback else 8.0)]
                    n_fail += 1
                continue
            
            try:
                with rasterio.open(tp) as src:
                    bands = src.read().astype(np.float64)
                    nodata_val = src.nodata if src.nodata is not None else NODATA
                    H, W = src.height, src.width
                    
                    scale = detect_scale(bands)
                    if scale > 1:
                        bands = bands / scale
                    
                    for _, pt in grp.iterrows():
                        key = f"{fn}_{pt['Lon']}_{pt['Lat']}"
                        
                        try:
                            rc = rowcol(src.transform, float(pt['Lon']), float(pt['Lat']))
                            r, c = int(rc[0]), int(rc[1])
                        except:
                            results_chla[key] = [float(chla_fallback.get('global', 8.0) if chla_fallback else 8.0)]
                            n_fail += 1
                            continue
                        
                        if not (0 <= r < H and 0 <= c < W):
                            results_chla[key] = [float(chla_fallback.get('global', 8.0) if chla_fallback else 8.0)]
                            n_fail += 1
                            continue
                        
                        try:
                            f = extract_features_v11(bands, nodata_val, r, c, H, W,
                                                    gloria_lib, s2_names)
                        except Exception as e_feat:
                            results_chla[key] = [float(chla_fallback.get('global', 8.0) if chla_fallback else 8.0)]
                            n_fail += 1
                            continue
                        
                        if f.pop('_invalid', False):
                            results_chla[key] = [float(chla_fallback.get('global', 8.0) if chla_fallback else 8.0)]
                            n_fail += 1
                            continue
                        
                        # Physics ensemble for Chl-a
                        chla_phys = []
                        for pc in ['C_ndci', 'C_mci', 'C_oci', 'C_b5b4']:
                            v = f.get(pc, np.nan)
                            if np.isfinite(v) and v > 0:
                                chla_phys.append(v)
                        
                        if chla_phys:
                            C_phys = float(np.median(chla_phys))
                        else:
                            C_phys = float(chla_fallback.get('global', 8.0) if chla_fallback else 8.0)
                        
                        # ML prediction (if models available)
                        if chla_models['lgbm'] is not None and chla_feat_cols is not None:
                            x_vec = np.array([f.get(fc, chla_train_med.get(fc, 0.0) if chla_train_med is not None else 0.0)
                                             for fc in chla_feat_cols]).reshape(1, -1)
                            x_vec = np.nan_to_num(x_vec, nan=0.0)
                            
                            f_df = pd.DataFrame([f])
                            cl = chla_owt.predict(f_df)[0] if chla_owt is not None else 0
                            cl_med = chla_owt.get_cluster_median(cl) if chla_owt is not None else 8.0
                            
                            wl = chla_weights.get('lgbm', 0.25) if chla_weights else 0.25
                            wx = chla_weights.get('xgb', 0.25) if chla_weights else 0.25
                            wc_w = chla_weights.get('cat', 0.25) if chla_weights else 0.25
                            wr = chla_weights.get('ridge', 0.25) if chla_weights else 0.25
                            
                            norm_pred = (
                                wl * chla_models['lgbm'].predict(x_vec)[0] +
                                wx * chla_models['xgb'].predict(x_vec)[0] +
                                wc_w * chla_models['cat'].predict(x_vec)[0] +
                                wr * chla_models['ridge'].predict(chla_scaler.transform(x_vec))[0]
                            )
                            
                            C_ml = max(0, np.exp(norm_pred + np.log(cl_med + 1.)) - 1.)
                            
                            # Blend
                            C_final = 0.4 * C_phys + 0.6 * C_ml
                        else:
                            C_final = C_phys
                        
                        C_final = max(0, min(C_final, 500))
                        results_chla[key] = [float(C_final)]
                        n_ok += 1
                        
            except Exception as e:
                print(f"  ERR {fn}: {e}")
                for _, pt in grp.iterrows():
                    key = f"{fn}_{pt['Lon']}_{pt['Lat']}"
                    results_chla[key] = [float(chla_fallback.get('global', 8.0) if chla_fallback else 8.0)]
                    n_fail += 1
        
        print(f"  Chl-a: {n_ok} ok, {n_fail} fallback")
        
        with open(OUTPUT_DIR / 'result_chla.json', 'w') as f:
            json.dump(results_chla, f, indent=2)
        print(f"  Saved: {OUTPUT_DIR}/result_chla.json")


# ═══════════════════════════════════════════════════════════════════════════════
# MAIN
# ═══════════════════════════════════════════════════════════════════════════════

def main():
    t_start = time.time()
    
    print("╔══════════════════════════════════════════════════════════════════╗")
    print("║  ITU TRACK 2 — V11 PHYSICS-GUIDED DOMAIN-ADAPTIVE ENSEMBLE     ║")
    print("╚══════════════════════════════════════════════════════════════════╝")
    print()
    print("  Improvements over V8+GLORIA:")
    print("    1. Multi-physics ensemble (10 models)")
    print("    2. Domain-invariant spectral shape features")
    print("    3. GLORIA spectral neighbor matching (kNN)")
    print("    4. Per-cluster physics calibration")
    print("    5. Comprehensive evaluation framework")
    print("    6. Chl-a pipeline (NDCI/MCI/OCI ensemble)")
    print("    7. Physics-ML hybrid with confidence weighting")
    print()
    
    # Train turbidity
    turb_result = train_turbidity()
    
    if turb_result is not None:
        artifacts, gloria_lib, s2_names = turb_result
        
        # Train Chl-a
        chla_result = train_chla(gloria_lib, s2_names)
        
        # Run inference
        run_inference()
    
    total_time = (time.time() - t_start) / 60
    print(f"\n{'='*70}")
    print(f"  TOTAL TIME: {total_time:.1f} min")
    print(f"  Outputs: {OUTPUT_DIR}")
    print(f"  Weights: {WEIGHTS_DIR}")
    print(f"{'='*70}")


if __name__ == '__main__':
    main()