"""
All 5 evaluations for Campus Weather VAE.
Run: python evaluate.py
"""
import os, sys, json
sys.path.insert(0, os.path.dirname(__file__))
import numpy as np
import torch
from torch.utils.data import DataLoader
from sklearn.metrics import r2_score, mean_absolute_error, mean_squared_error
from sklearn.cluster import KMeans
from sklearn.linear_model import Ridge
from sklearn.decomposition import PCA
from model import WeatherVAE, get_config
from train import load_nus40, FlatDataset, VAR_NAMES, VAR_UNITS, load_trained, train_model

SAVE = '/app/campus_weather/results'
DATA = '/app/data_tmp/imputed'
DEVICE = 'cpu'


def eval_reconstruction(model, data, t_split):
    """Basic reconstruction quality on test set."""
    test = data[t_split:]
    x = torch.from_numpy(test.reshape(-1, 6))
    with torch.no_grad():
        out = model(x)
    pred, true = out['x_hat'].numpy(), x.numpy()
    
    results = {}
    for v, (name, unit) in enumerate(zip(VAR_NAMES, VAR_UNITS)):
        results[name] = {
            'MAE': float(mean_absolute_error(true[:, v], pred[:, v])),
            'RMSE': float(np.sqrt(mean_squared_error(true[:, v], pred[:, v]))),
            'R2': float(r2_score(true[:, v], pred[:, v])),
        }
    return results


def eval1_spatial_interpolation(data, coords):
    """
    Hold out 5 geographically spread stations. Train on 35. 
    Reconstruct held-out stations using embeddings from neighbouring stations.
    """
    print("\n" + "="*60)
    print("EVAL 1: Spatial Interpolation (hold-out 5 stations)")
    print("="*60)
    
    N = data.shape[1]
    # Pick 5 spread across campus (south, central, north, east, west)
    holdout_idx = [4, 12, 20, 30, 37]  # WS05, WS13, WS21, WS31, WS38
    train_idx = [i for i in range(N) if i not in holdout_idx]
    
    print(f"Hold-out stations: {[i+1 for i in holdout_idx]}")
    print(f"Training stations: {len(train_idx)}")
    
    # Train model on 35 stations
    T = data.shape[0]
    t_tr = int(T * 0.7)
    data_35 = data[:, train_idx, :]
    
    cfg = get_config('base')
    model_35 = WeatherVAE(**cfg)
    mean = torch.tensor(data_35[:t_tr].mean(axis=(0,1)), dtype=torch.float32)
    std = torch.tensor(data_35[:t_tr].std(axis=(0,1)), dtype=torch.float32)
    model_35.set_normalisation(mean, std)
    
    # Quick train (50 epochs — enough to converge on this simple model)
    import torch.optim as optim
    opt = optim.AdamW(model_35.parameters(), lr=5e-4, weight_decay=0.01)
    tr_ds = FlatDataset(data_35[:t_tr])
    tr_ld = DataLoader(tr_ds, batch_size=256, shuffle=True, drop_last=True)
    
    model_35.train()
    for ep in range(50):
        for batch in tr_ld:
            out = model_35(batch)
            out['loss'].backward()
            torch.nn.utils.clip_grad_norm_(model_35.parameters(), 1.0)
            opt.step(); opt.zero_grad()
    model_35.eval()
    
    # For each held-out station, reconstruct using k nearest training stations
    from sklearn.neighbors import NearestNeighbors
    train_coords = coords[train_idx]
    holdout_coords = coords[holdout_idx]
    
    nn_model = NearestNeighbors(n_neighbors=5).fit(train_coords)
    _, nn_idx = nn_model.kneighbors(holdout_coords)
    
    t_test = int(T * 0.85)
    test_data = data[t_test:]
    
    results = {}
    for hi, ho_station in enumerate(holdout_idx):
        # Get embeddings of 5 nearest training stations
        neighbour_stations = [train_idx[j] for j in nn_idx[hi]]
        neighbour_data = test_data[:, neighbour_stations, :]  # (T_test, 5, 6)
        
        # Average neighbour embeddings as proxy for held-out station
        with torch.no_grad():
            n_flat = torch.from_numpy(neighbour_data.reshape(-1, 6))
            n_emb = model_35.get_embedding(n_flat).numpy()
            n_emb = n_emb.reshape(test_data.shape[0], 5, -1).mean(axis=1)  # (T_test, d)
            # Decode averaged embedding
            pred = model_35.decode(torch.from_numpy(n_emb)).numpy()
        
        true = test_data[:, ho_station, :]
        
        station_results = {}
        for v, name in enumerate(VAR_NAMES):
            station_results[name] = {
                'MAE': float(mean_absolute_error(true[:, v], pred[:, v])),
                'R2': float(r2_score(true[:, v], pred[:, v])),
            }
        results[f'WS{ho_station+1:02d}'] = station_results
        print(f"  WS{ho_station+1:02d}: AirTemp MAE={station_results['AirTemp']['MAE']:.3f}°C, "
              f"R²={station_results['AirTemp']['R2']:.4f}")
    
    # Average across held-out stations
    avg = {}
    for name in VAR_NAMES:
        maes = [results[s][name]['MAE'] for s in results]
        r2s = [results[s][name]['R2'] for s in results]
        avg[name] = {'MAE': float(np.mean(maes)), 'R2': float(np.mean(r2s))}
    results['average'] = avg
    print(f"\n  Average: AirTemp MAE={avg['AirTemp']['MAE']:.3f}°C R²={avg['AirTemp']['R2']:.4f} | "
          f"RelHum MAE={avg['RelHum']['MAE']:.3f}% R²={avg['RelHum']['R2']:.4f}")
    
    return results


def eval2_temporal_forecasting(embeddings, data):
    """
    Forecast T+1/6/24 using linear regression on embeddings.
    Compare vs persistence and climatology baselines.
    """
    print("\n" + "="*60)
    print("EVAL 2: Temporal Forecasting (embedding vs baselines)")
    print("="*60)
    
    T, N, V = data.shape
    t_tr, t_te = int(T * 0.7), int(T * 0.85)
    hours = np.arange(T) % 24
    
    # Compute hourly climatology from training data
    climatology = np.zeros((24, N, V))
    for h in range(24):
        mask = hours[:t_tr] == h
        climatology[h] = data[:t_tr][mask].mean(axis=0)
    
    results = {}
    for horizon in [1, 6, 24]:
        # Build train/test pairs: embedding at t → weather at t+horizon
        tr_X, tr_Y = [], []
        for t in range(0, t_tr - horizon):
            tr_X.append(embeddings[t].reshape(N, -1))  # (N, d)
            tr_Y.append(data[t + horizon])              # (N, V)
        tr_X = np.array(tr_X).reshape(-1, embeddings.shape[-1])  # (samples*N, d)
        tr_Y = np.array(tr_Y).reshape(-1, V)
        
        te_X, te_Y = [], []
        te_persist, te_clim = [], []
        for t in range(t_te, T - horizon):
            te_X.append(embeddings[t].reshape(N, -1))
            te_Y.append(data[t + horizon])
            te_persist.append(data[t])  # persistence: weather at t
            te_clim.append(climatology[(t + horizon) % 24])
        te_X = np.array(te_X).reshape(-1, embeddings.shape[-1])
        te_Y = np.array(te_Y).reshape(-1, V)
        te_persist = np.array(te_persist).reshape(-1, V)
        te_clim = np.array(te_clim).reshape(-1, V)
        
        horizon_results = {}
        for v, name in enumerate(VAR_NAMES):
            reg = Ridge(alpha=1.0)
            reg.fit(tr_X, tr_Y[:, v])
            pred = reg.predict(te_X)
            
            mae_emb = mean_absolute_error(te_Y[:, v], pred)
            mae_persist = mean_absolute_error(te_Y[:, v], te_persist[:, v])
            mae_clim = mean_absolute_error(te_Y[:, v], te_clim[:, v])
            
            horizon_results[name] = {
                'MAE_embedding': float(mae_emb),
                'MAE_persistence': float(mae_persist),
                'MAE_climatology': float(mae_clim),
            }
        
        results[f'T+{horizon}'] = horizon_results
        print(f"\n  T+{horizon}h:")
        print(f"  {'Variable':>12s} {'Embedding':>10s} {'Persist':>10s} {'Climat':>10s} {'Skill':>8s}")
        for name in ['AirTemp', 'RelHum', 'GlobalRad', 'WindSpeed']:
            r = horizon_results[name]
            skill = 1 - r['MAE_embedding'] / r['MAE_persistence']
            print(f"  {name:>12s} {r['MAE_embedding']:>10.3f} {r['MAE_persistence']:>10.3f} "
                  f"{r['MAE_climatology']:>10.3f} {skill:>7.1%}")
    
    return results


def eval3_clustering(embeddings, coords):
    """
    Unsupervised microclimate zone discovery from station embeddings.
    """
    print("\n" + "="*60)
    print("EVAL 3: Microclimate Clustering")
    print("="*60)
    
    # Average embedding per station across all time
    station_emb = embeddings.mean(axis=0)  # (N, d)
    
    # Try K=3,4,5 clusters
    results = {}
    for k in [3, 4, 5]:
        km = KMeans(n_clusters=k, random_state=42, n_init=10)
        labels = km.fit_predict(station_emb)
        inertia = km.inertia_
        
        # Silhouette score
        from sklearn.metrics import silhouette_score
        sil = silhouette_score(station_emb, labels) if k < len(station_emb) else 0
        
        results[f'K={k}'] = {
            'labels': labels.tolist(),
            'inertia': float(inertia),
            'silhouette': float(sil),
        }
        print(f"  K={k}: silhouette={sil:.3f}")
        for c in range(k):
            stations = [i+1 for i in range(len(labels)) if labels[i] == c]
            lat_mean = coords[labels == c, 0].mean()
            lng_mean = coords[labels == c, 1].mean()
            print(f"    Cluster {c}: {len(stations)} stations, "
                  f"centroid=({lat_mean:.4f}, {lng_mean:.4f}), stations={stations[:8]}...")
    
    # PCA for visualisation
    pca = PCA(n_components=2)
    station_2d = pca.fit_transform(station_emb)
    results['pca_2d'] = station_2d.tolist()
    results['explained_var'] = pca.explained_variance_ratio_.tolist()
    
    return results


def eval4_anomaly_detection(model, data, embeddings):
    """
    Anomaly detection via reconstruction error.
    High error = unusual weather the model hasn't learned.
    """
    print("\n" + "="*60)
    print("EVAL 4: Anomaly Detection")
    print("="*60)
    
    T, N, V = data.shape
    
    # Compute per-timestep reconstruction error
    recon_errors = np.zeros((T, N))
    with torch.no_grad():
        for t in range(0, T, 500):
            chunk = data[t:t+500]  # (chunk, N, V)
            ct, cn = chunk.shape[0], chunk.shape[1]
            x = torch.from_numpy(chunk.reshape(-1, V))
            out = model(x)
            err = (out['x_hat'].numpy() - x.numpy()) ** 2
            err = err.mean(axis=1)  # mean across variables
            recon_errors[t:t+ct] = err.reshape(ct, cn)
    
    # Per-station mean error
    station_mean_err = recon_errors.mean(axis=0)
    
    # Find top anomalous hours (campus-wide)
    campus_err = recon_errors.mean(axis=1)  # (T,)
    top_anomalies_idx = np.argsort(campus_err)[-20:][::-1]
    
    # Threshold at 95th percentile
    threshold = np.percentile(campus_err, 95)
    anomaly_mask = campus_err > threshold
    n_anomalies = anomaly_mask.sum()
    
    # Check if anomalies correlate with rainfall
    rain = data[:, :, 5]  # GlobalRad as proxy (low = cloudy/rainy)
    campus_rain = rain.mean(axis=1)
    anomaly_rain = campus_rain[anomaly_mask].mean()
    normal_rain = campus_rain[~anomaly_mask].mean()
    
    hours_of_day = np.arange(T) % 24
    anomaly_hours = hours_of_day[anomaly_mask]
    hour_dist = np.bincount(anomaly_hours.astype(int), minlength=24)
    
    results = {
        'station_mean_error': station_mean_err.tolist(),
        'threshold_95': float(threshold),
        'n_anomalies': int(n_anomalies),
        'anomaly_rate': float(n_anomalies / T),
        'mean_globalrad_anomaly': float(anomaly_rain),
        'mean_globalrad_normal': float(normal_rain),
        'anomaly_hour_distribution': hour_dist.tolist(),
        'top_20_indices': top_anomalies_idx.tolist(),
        'campus_error_timeseries': campus_err.tolist(),
    }
    
    print(f"  Threshold (95th pct): {threshold:.4f}")
    print(f"  Anomalous hours: {n_anomalies}/{T} ({n_anomalies/T*100:.1f}%)")
    print(f"  Mean GlobalRad during anomalies: {anomaly_rain:.1f} vs normal: {normal_rain:.1f} W/m²")
    print(f"  Peak anomaly hours: {np.argsort(hour_dist)[-3:][::-1].tolist()}")
    
    return results


def eval5_future_prediction(embeddings, data):
    """
    Multi-step rolling forecast. At each test hour, predict next 24 hours.
    Compare vs persistence and climatology.
    """
    print("\n" + "="*60)
    print("EVAL 5: Rolling 24h Future Prediction")
    print("="*60)
    
    T, N, V = data.shape
    t_tr, t_te = int(T * 0.7), int(T * 0.85)
    hours = np.arange(T) % 24
    
    # Climatology
    climatology = np.zeros((24, N, V))
    for h in range(24):
        mask = hours[:t_tr] == h
        climatology[h] = data[:t_tr][mask].mean(axis=0)
    
    # Train one Ridge per (horizon, variable) on training data
    models = {}  # (horizon, var) → Ridge
    for horizon in range(1, 25):
        tr_X = embeddings[:t_tr - horizon].reshape(-1, embeddings.shape[-1])
        tr_Y = data[horizon:t_tr].reshape(-1, V)
        for v in range(V):
            reg = Ridge(alpha=1.0)
            reg.fit(tr_X, tr_Y[:, v])
            models[(horizon, v)] = reg
    
    # Rolling forecast on test set
    n_test = T - t_te - 24
    all_pred = np.zeros((n_test, 24, N, V))
    all_true = np.zeros((n_test, 24, N, V))
    all_persist = np.zeros((n_test, 24, N, V))
    all_clim = np.zeros((n_test, 24, N, V))
    
    for i, t in enumerate(range(t_te, t_te + n_test)):
        emb_t = embeddings[t].reshape(-1, embeddings.shape[-1])  # (N, d)
        for h in range(24):
            for v in range(V):
                all_pred[i, h, :, v] = models[(h + 1, v)].predict(emb_t)
            all_true[i, h] = data[t + h + 1]
            all_persist[i, h] = data[t]  # persistence: current hour
            all_clim[i, h] = climatology[(t + h + 1) % 24]
    
    # Compute MAE per horizon
    results = {'per_hour': {}}
    for h in range(24):
        hr_results = {}
        for v, name in enumerate(VAR_NAMES):
            mae_emb = mean_absolute_error(all_true[:, h, :, v].flatten(), all_pred[:, h, :, v].flatten())
            mae_per = mean_absolute_error(all_true[:, h, :, v].flatten(), all_persist[:, h, :, v].flatten())
            mae_clm = mean_absolute_error(all_true[:, h, :, v].flatten(), all_clim[:, h, :, v].flatten())
            hr_results[name] = {
                'MAE_embedding': float(mae_emb),
                'MAE_persistence': float(mae_per),
                'MAE_climatology': float(mae_clm),
            }
        results['per_hour'][f'h+{h+1}'] = hr_results
    
    # Summary at key horizons
    for h_show in [0, 5, 11, 23]:
        hr = results['per_hour'][f'h+{h_show+1}']
        print(f"\n  h+{h_show+1}:")
        for name in ['AirTemp', 'RelHum', 'GlobalRad']:
            r = hr[name]
            skill = 1 - r['MAE_embedding'] / r['MAE_persistence']
            print(f"    {name:>12s}: Emb={r['MAE_embedding']:.3f} Pers={r['MAE_persistence']:.3f} "
                  f"Clim={r['MAE_climatology']:.3f} Skill={skill:.1%}")
    
    return results


def run_all():
    """Run all 5 evaluations."""
    os.makedirs(SAVE, exist_ok=True)
    
    # Load model and data
    model, data, coords, embeddings, ckpt = load_trained(f'{SAVE}/checkpoints')
    T = data.shape[0]
    t_te = int(T * 0.85)
    
    all_results = {}
    
    # Reconstruction baseline
    print("="*60)
    print("BASELINE: Reconstruction Quality")
    print("="*60)
    recon = eval_reconstruction(model, data, t_te)
    all_results['reconstruction'] = recon
    for name in VAR_NAMES:
        r = recon[name]
        print(f"  {name:>12s}: MAE={r['MAE']:.4f} RMSE={r['RMSE']:.4f} R²={r['R2']:.4f}")
    
    # Eval 1
    all_results['spatial_interpolation'] = eval1_spatial_interpolation(data, coords)
    
    # Eval 2
    all_results['temporal_forecasting'] = eval2_temporal_forecasting(embeddings, data)
    
    # Eval 3
    all_results['clustering'] = eval3_clustering(embeddings, coords)
    
    # Eval 4
    all_results['anomaly_detection'] = eval4_anomaly_detection(model, data, embeddings)
    
    # Eval 5
    all_results['future_prediction'] = eval5_future_prediction(embeddings, data)
    
    # Save
    # Remove large arrays for JSON serialisation
    save_results = {}
    for k, v in all_results.items():
        if k == 'anomaly_detection':
            v2 = {kk: vv for kk, vv in v.items() if kk != 'campus_error_timeseries'}
            save_results[k] = v2
        else:
            save_results[k] = v
    
    with open(f'{SAVE}/all_results.json', 'w') as f:
        json.dump(save_results, f, indent=2, default=str)
    
    # Save anomaly timeseries separately (large)
    np.save(f'{SAVE}/anomaly_errors.npy', np.array(all_results['anomaly_detection']['campus_error_timeseries']))
    
    print(f"\n{'='*60}")
    print(f"All results saved to {SAVE}/")
    return all_results


if __name__ == '__main__':
    run_all()