#!/usr/bin/env python3
"""
Training Pipeline for Spatial-Temporal Transformer + GCN Model.

Simplified approach: use global weather mean per day as node features,
scaled by per-node spatial features (elevation, population density).
"""

import os
import sys
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import warnings
warnings.filterwarnings('ignore')

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
import mlflow
from pathlib import Path
from datetime import datetime

from models.spatiotemporal_gcn.model import SpatialTemporalGCN

DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {DEVICE}")

PROCESSED_DIR = Path('processed')
MODEL_DIR = Path('models/spatiotemporal_gcn')
MODEL_DIR.mkdir(exist_ok=True)

LEARNING_RATE = 1e-4
WEIGHT_DECAY = 0.01
PATIENCE = 15
MAX_EPOCHS = 200
BATCH_SIZE = 1024

# Data split (medical data only available in December)
TRAIN_START = '2022-12-01'
TRAIN_END = '2022-12-31'
VAL_START = '2023-12-01'
VAL_END = '2023-12-31'

# Baseline MAE from compute_baseline_mae.py
BASELINE_MAE = {'1-day': 0.2314, '3-day': 0.5424, '7-day': 0.6391}

SEED = 42
np.random.seed(SEED)
torch.manual_seed(SEED)


def load_all_data():
    """Load all processed data."""
    print("Loading data...")

    # Graph
    adj = np.load(PROCESSED_DIR / 'graph' / 'adjacency_matrix.npz')
    from scipy.sparse import csr_matrix
    sp_adj = csr_matrix((adj['data'], adj['indices'], adj['indptr']), shape=tuple(adj['shape']))
    sp_adj_coo = sp_adj.tocoo()
    edge_index = torch.tensor(
        np.stack([sp_adj_coo.row, sp_adj_coo.col]),
        dtype=torch.long
    )  # Keep on CPU for subgraph operations

    # Node metadata
    nodes = pd.read_parquet(PROCESSED_DIR / 'graph' / 'node_features.parquet')
    n_nodes = len(nodes)
    print(f"  Graph: {n_nodes} nodes, {edge_index.shape[1]} edges")

    # Weather lag features (station-level daily)
    lf = pd.read_parquet(PROCESSED_DIR / 'weather' / 'lag_features.parquet')
    lf['date'] = pd.to_datetime(lf['date'])
    lf = lf.sort_values('date')
    print(f"  Weather: {len(lf)} records, {lf['station_id'].nunique()} stations")

    # Medical targets (district-level daily)
    out = pd.read_csv(PROCESSED_DIR / 'medical' / 'outpatient_daily.csv', parse_dates=['date'])
    inp = pd.read_csv(PROCESSED_DIR / 'medical' / 'inpatient_daily.csv', parse_dates=['date'])
    out['weight'] = 1
    inp['weight'] = 3
    combined = pd.concat([out, inp])
    combined['weighted_cases'] = combined['case_count'] * combined['weight']
    medical = combined.groupby(['date', 'district']).agg(
        weighted_cases=('weighted_cases', 'sum')
    ).reset_index()
    medical['risk'] = medical.groupby('district')['weighted_cases'].transform(
        lambda x: x / x.mean()
    )
    print(f"  Medical: {len(medical)} district-day records")

    return edge_index, nodes, lf, medical


def build_global_weather_timeseries(lf):
    """
    Build global mean weather per day: [T, 48]
    """
    feat_cols = [c for c in lf.columns if c not in ('date', 'station_id')]
    daily_mean = lf.groupby('date')[feat_cols].mean()
    daily_mean = daily_mean.sort_index()
    dates = daily_mean.index.tolist()
    x_global = daily_mean.values.astype(np.float32)  # [T, 48]
    return x_global, dates


def build_node_targets(nodes, medical, dates):
    """
    Build per-node risk target per day: [N, T]
    Use district-level medical risk, tiled to all nodes in district.
    District assignment from node lat/lon nearest centroid (simplified: use 'unknown').
    For nodes with no district match, use global mean risk.
    """
    n_nodes = len(nodes)
    n_days = len(dates)

    # Global mean risk per day
    global_risk = medical.groupby('date')['risk'].mean()
    global_risk_dict = global_risk.to_dict()

    # For each node, assign a district based on nearest centroid
    # (simplified: just use global risk for all nodes)
    targets = np.full((n_nodes, n_days), np.nan, dtype=np.float32)

    for i, d in enumerate(dates):
        if d in global_risk_dict:
            targets[:, i] = global_risk_dict[d]

    # Normalize per node
    node_means = np.nanmean(targets, axis=1, keepdims=True)
    node_means[node_means == 0] = 1
    targets = targets / (node_means + 1e-8)

    return targets, dates


def build_spatial_scalars(nodes):
    """
    Pre-compute per-node spatial scaling factors (small, O(N)).
    Returns: elev_scale [N], pop_scale [N]
    """
    elev = nodes['elevation_m'].values
    pop = nodes['pop_density'].values
    elev_norm = (elev - elev.mean()) / (elev.std() + 1e-8)
    pop_norm = (pop - pop.mean()) / (pop.std() + 1e-8)
    
    # Scaling factors
    elev_scale = 1.0 + 0.1 * elev_norm
    elev_scale = np.clip(elev_scale, 0.5, 2.0).astype(np.float32)
    
    # Pop scale (optional, can be 1.0 if not used)
    pop_scale = np.ones_like(elev_scale)  # or add similar modulation if needed
    
    return elev_scale, pop_scale


def get_batch_features(elev_scale, x_global, node_indices):
    """
    Compute features for a batch of nodes on-the-fly.
    elev_scale: [N] pre-computed spatial scalars
    x_global: [T, F] global weather per day
    node_indices: list/array of node indices to fetch
    
    Returns: [len(node_indices), T, F]
    """
    batch_size = len(node_indices)
    T, F = x_global.shape
    
    # Get spatial scales for batch
    batch_elev = elev_scale[node_indices]
    
    # Tile global weather for batch: [T, F] -> [batch, T, F]
    x = np.tile(x_global[np.newaxis, :, :], (batch_size, 1, 1))
    
    # Apply spatial scaling
    x = x * batch_elev[:, np.newaxis, np.newaxis]
    
    return x.astype(np.float32)


def make_time_windows_lazy(x_global, elev_scale, targets, dates, window=14):
    """
    Create time window metadata without materializing full [N, T, F] tensor.
    Returns list of (time_start, node_indices) tuples for lazy feature fetching.
    
    x_global: [T, F] global weather
    elev_scale: [N] spatial scaling per node
    targets: [N, T] target values
    window: input window size
    """
    N = len(elev_scale)
    T = x_global.shape[0]
    
    # Store window metadata: which time steps and which nodes
    windows_meta = []
    for t in range(T - window + 1):
        # All nodes for this time window
        windows_meta.append({
            'time_start': t,
            'time_end': t + window,
            'target_time': t + window - 1,
        })
    
    return windows_meta


def train_epoch_lazy(model, windows_meta, x_global, elev_scale, y, edge_index, 
                     optimizer, criterion, batch_size=1024):
    """
    Train one epoch using lazy feature computation.
    For each window, sample a batch of nodes and compute features on-the-fly.
    """
    from torch_geometric.utils import subgraph
    
    model.train()
    total_loss = 0
    n_batches = 0
    n_windows = len(windows_meta)
    n_nodes = len(elev_scale)
    
    # Process each time window
    for window_meta in windows_meta:
        t_start = window_meta['time_start']
        t_end = window_meta['time_end']
        t_target = window_meta['target_time']
        
        node_indices = np.random.choice(n_nodes, size=min(batch_size, n_nodes), replace=False)
        node_indices_torch = torch.tensor(node_indices, dtype=torch.long)
        
        x_batch = get_batch_features(elev_scale, x_global[t_start:t_end], node_indices)
        y_batch = y[node_indices, t_target]
        
        # Filter out NaN targets
        valid_mask = ~np.isnan(y_batch)
        if valid_mask.sum() == 0:
            continue
        
        # Extract subgraph and manually remap indices to ensure correctness
        sub_edge_index, edge_mask = subgraph(node_indices_torch, edge_index, relabel_nodes=False)
        
        # Create remapping: global_id -> local_idx (0 to batch_size-1)
        # Use index_put for efficient remapping
        local_idx = torch.arange(len(node_indices), dtype=torch.long)
        remap_tensor = torch.full((n_nodes,), -1, dtype=torch.long)
        remap_tensor[node_indices_torch] = local_idx
        
        # Remap edge indices
        sub_edge_index = remap_tensor[sub_edge_index]
        
        # Validate: all indices should be in [0, batch_size)
        assert sub_edge_index.min() >= 0 and sub_edge_index.max() < len(node_indices), \
            f"Edge index out of bounds: min={sub_edge_index.min()}, max={sub_edge_index.max()}"
        
        x_batch = torch.FloatTensor(x_batch).to(DEVICE)
        y_batch = torch.FloatTensor(y_batch).to(DEVICE)
        sub_edge_index = sub_edge_index.to(DEVICE)
        
        # Forward pass - only compute loss on valid samples
        optimizer.zero_grad()
        out = model(x_batch, sub_edge_index)
        loss = criterion(out[valid_mask, 1], y_batch[valid_mask])
        loss.backward()
        optimizer.step()
        
        total_loss += loss.item()
        n_batches += 1
    
    return total_loss / max(n_batches, 1)


def evaluate_lazy(model, x_global, elev_scale, y, edge_index, window=14):
    """
    Evaluate MAE per horizon using lazy feature computation.
    """
    from torch_geometric.utils import subgraph
    
    model.eval()
    T = x_global.shape[0]
    n_nodes = len(elev_scale)
    horizons = {'1-day': 1, '3-day': 3, '7-day': 7}
    results = {}
    
    with torch.no_grad():
        for name, h in horizons.items():
            if h > T - window:
                results[name] = float('nan')
                continue
            
            preds_all = []
            acts_all = []
            
            for t in range(window, T - h + 1, 5):
                node_indices = np.random.choice(n_nodes, size=min(100, n_nodes), replace=False)
                node_indices_torch = torch.tensor(node_indices, dtype=torch.long)
                
                x_win = get_batch_features(elev_scale, x_global[t-window:t], node_indices)
                x_win = torch.FloatTensor(x_win).to(DEVICE)
                
                y_actual = y[node_indices, t+h-1]
                y_actual = torch.FloatTensor(y_actual).to(DEVICE)
                
                # Extract subgraph and manually remap indices
                sub_edge_index, _ = subgraph(node_indices_torch, edge_index, relabel_nodes=False)
                
                # Remap global IDs to local indices
                local_idx = torch.arange(len(node_indices), dtype=torch.long)
                remap_tensor = torch.full((n_nodes,), -1, dtype=torch.long)
                remap_tensor[node_indices_torch] = local_idx
                sub_edge_index = remap_tensor[sub_edge_index]
                
                sub_edge_index = sub_edge_index.to(DEVICE)
                
                # Filter out NaN targets
                valid_mask = ~torch.isnan(y_actual)
                if valid_mask.sum() == 0:
                    continue
                
                pred = model(x_win, sub_edge_index)[valid_mask, 1]
                preds_all.append(pred.mean())
                acts_all.append(y_actual[valid_mask].mean())
            
            if preds_all:
                preds = torch.stack(preds_all).mean()
                acts = torch.stack(acts_all).mean()
                results[name] = torch.mean(torch.abs(preds - acts)).item()
            else:
                results[name] = float('nan')
    
    return results


def main():
    print(f"\n=== Training Pipeline ===  {datetime.now()}")

    edge_index, nodes, lf, medical = load_all_data()
    n_nodes = len(nodes)

    # Build time series - global weather only (small: [T, 48])
    x_global, weather_dates = build_global_weather_timeseries(lf)
    targets, _ = build_node_targets(nodes, medical, weather_dates)
    
    # Pre-compute spatial scalars (small: O(N))
    elev_scale, pop_scale = build_spatial_scalars(nodes)

    print(f"\nGlobal weather: {x_global.shape} [T, F]")
    print(f"Targets: {targets.shape} [N, T]")
    print(f"Spatial scalars: {len(elev_scale)} nodes")

    # Align to training period
    dates_arr = pd.to_datetime(weather_dates)
    train_mask = (dates_arr >= TRAIN_START) & (dates_arr <= TRAIN_END)
    val_mask = (dates_arr >= VAL_START) & (dates_arr <= VAL_END)

    x_global_train = x_global[train_mask]
    y_train = targets[:, train_mask]
    x_global_val = x_global[val_mask]
    y_val = targets[:, val_mask]
    
    train_days = len(x_global_train)
    val_days = len(x_global_val)
    print(f"Train: {train_days} steps, Val: {val_days} steps")

    # Make training windows (metadata only, no large tensors)
    WINDOW = 14
    print("Building training windows (metadata)...")
    windows_meta = make_time_windows_lazy(x_global_train, elev_scale, y_train, 
                                          pd.to_datetime(weather_dates)[train_mask].tolist(), 
                                          window=WINDOW)
    print(f"  {len(windows_meta)} windows")

    # Model
    model = SpatialTemporalGCN(
        node_features=48,
        temporal_heads=4,
        temporal_layers=3,
        gcn_hidden=128,
        gcn_output=64,
        dropout=0.2
    ).to(DEVICE)
    print(f"\nModel params: {sum(p.numel() for p in model.parameters()):,}")

    optimizer = optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY)
    criterion = nn.L1Loss()
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', patience=5, factor=0.5)

    # Truncate edge_index for batch processing
    edge_idx_trunc = edge_index[:, :min(edge_index.shape[1], n_nodes * 4)].contiguous()

    mlflow.set_experiment("wuhan_respiratory_training")
    with mlflow.start_run(run_name=f"train_{datetime.now().strftime('%Y%m%d_%H%M%S')}"):
        mlflow.log_params({
            "learning_rate": LEARNING_RATE,
            "weight_decay": WEIGHT_DECAY,
            "patience": PATIENCE,
            "max_epochs": MAX_EPOCHS,
            "window": WINDOW,
            "n_nodes": n_nodes,
            "train_start": TRAIN_START, "train_end": TRAIN_END,
            "val_start": VAL_START, "val_end": VAL_END,
            "baseline_mae_1d": BASELINE_MAE['1-day'],
            "baseline_mae_3d": BASELINE_MAE['3-day'],
            "baseline_mae_7d": BASELINE_MAE['7-day'],
        })

        best_val_mae = float('inf')
        patience_counter = 0
        best_state = None

        for epoch in range(1, MAX_EPOCHS + 1):
            train_loss = train_epoch_lazy(model, windows_meta, x_global_train, elev_scale, 
                                          y_train, edge_idx_trunc, optimizer, criterion, BATCH_SIZE)
            val_mae_h = evaluate_lazy(model, x_global_val, elev_scale, y_val, edge_idx_trunc, WINDOW)
            val_mae = np.nanmean(list(val_mae_h.values()))

            scheduler.step(val_mae)

            if epoch % 5 == 0 or val_mae < best_val_mae:
                print(f"Epoch {epoch:3d} | Loss: {train_loss:.4f} | Val MAE: {val_mae:.4f} "
                      f"| 1d:{val_mae_h.get('1-day', 0):.4f} "
                      f"3d:{val_mae_h.get('3-day', 0):.4f} "
                      f"7d:{val_mae_h.get('7-day', 0):.4f}")

            mlflow.log_metrics({
                "train_loss": train_loss,
                f"val_mae_1d": val_mae_h.get('1-day', float('nan')),
                f"val_mae_3d": val_mae_h.get('3-day', float('nan')),
                f"val_mae_7d": val_mae_h.get('7-day', float('nan')),
            }, step=epoch)

            if val_mae < best_val_mae:
                best_val_mae = val_mae
                best_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}
                patience_counter = 0
            else:
                patience_counter += 1
                if patience_counter >= PATIENCE:
                    print(f"\nEarly stopping at epoch {epoch}")
                    break

        # Save
        model.load_state_dict(best_state)
        torch.save(best_state, MODEL_DIR / 'best_model.pt')
        mlflow.log_artifact(MODEL_DIR / 'best_model.pt')

        # Beat-baseline check
        beat_count = sum(
            val_mae_h.get(h, float('inf')) < 0.9 * BASELINE_MAE[h]
            for h in ('1-day', '3-day', '7-day')
        )
        print(f"\nBest Val MAE: {best_val_mae:.4f}")
        print(f"Baseline 1d/3d/7d: {BASELINE_MAE['1-day']:.4f}/{BASELINE_MAE['3-day']:.4f}/{BASELINE_MAE['7-day']:.4f}")
        print(f"Beats baseline at 0.9x: {beat_count}/3 horizons")

    print(f"\nDone! {datetime.now()}")


if __name__ == '__main__':
    main()