"""
Run full pipeline: Stage 1 (sensors -> Farquhar -> A labels), validation, Stage 2 (IMS + labels -> train/eval), plots.
Requires: Data/Seymour/sensors_wide.csv, Data/ims/ims_merged_15min.csv (or run download_ims_data first).

Usage:
  python -m scripts.run_pipeline
"""

from __future__ import annotations

import sys
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parent.parent
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

try:
    from dotenv import load_dotenv
    load_dotenv(PROJECT_ROOT / ".env")
except ImportError:
    pass

import pandas as pd


def run_stage1():
    """Load sensors, filter daytime, compute A with Farquhar, save labels. Returns (df_sensor, labels)."""
    from config import settings
    from src.sensor_data_loader import SensorDataLoader
    from src.farquhar_model import FarquharModel

    settings.PROCESSED_DIR.mkdir(parents=True, exist_ok=True)
    loader = SensorDataLoader()
    df = loader.load()
    df = loader.filter_daytime(df, par_threshold=50.0)
    # Keep growing season only (May–September); vine dormant Oct–April, no photosynthesis
    ts_col = "time" if "time" in df.columns else df.columns[0]
    t = pd.to_datetime(df[ts_col], utc=True)
    df = df.loc[t.dt.month.isin(settings.GROWING_SEASON_MONTHS)].copy()
    if df.empty:
        return None, None
    model = FarquharModel()
    A = model.compute_all(df)
    # Index labels by timestamp; round to 15min and aggregate to align with IMS for Stage 2 merge
    ts_col = "time" if "time" in df.columns else df.columns[0]
    t = pd.to_datetime(df[ts_col], utc=True)
    A = A.copy()
    A.index = t.dt.floor("15min")
    A = A.groupby(level=0).mean()
    A.name = "A"
    path = settings.PROCESSED_DIR / "stage1_labels.csv"
    A.to_csv(path)
    print(f"Stage 1: saved {len(A)} labels to {path}")
    return df, A


def validate_stage1(df_sensor, labels, outputs_dir: Path):
    """Sanity check A (0–25 typical), diurnal plot, correlation with PAR and T."""
    import numpy as np
    import pandas as pd

    if labels is None or labels.empty:
        return
    outputs_dir.mkdir(parents=True, exist_ok=True)
    A = labels.dropna()
    stats = {"count": len(A), "mean": A.mean(), "std": A.std(), "min": A.min(), "max": A.max()}
    print("Stage 1 validation – A (umol m-2 s-1):", stats)
    if A.max() > 50 or A.min() < -5:
        print("  Note: typical grapevine A is 0–25; check inputs if out of range.")
    try:
        import matplotlib.pyplot as plt
        fig, axes = plt.subplots(2, 1, figsize=(10, 8))
        # Diurnal: A vs hour
        ts = pd.to_datetime(A.index, utc=True)
        hour = ts.hour + ts.minute / 60
        axes[0].scatter(hour, A.values, alpha=0.3, s=5)
        axes[0].set_xlabel("Hour (UTC)")
        axes[0].set_ylabel("A (umol m-2 s-1)")
        axes[0].set_title("Diurnal pattern (Stage 1)")
        axes[0].set_xlim(0, 24)
        if df_sensor is not None and not df_sensor.empty and "Air1_PAR_ref" in df_sensor.columns:
            ts_col = "time" if "time" in df_sensor.columns else df_sensor.columns[0]
            par_series = df_sensor.set_index(ts_col)["Air1_PAR_ref"]
            common = A.reindex(par_series.index).dropna()
            par = par_series.reindex(common.index).dropna()
            common = common.reindex(par.index).dropna()
            if len(par) > 10:
                axes[1].scatter(par.values, common.values, alpha=0.3, s=5)
                axes[1].set_xlabel("PAR (umol m-2 s-1)")
                axes[1].set_ylabel("A (umol m-2 s-1)")
                axes[1].set_title("A vs PAR")
        plt.tight_layout()
        plt.savefig(outputs_dir / "stage1_validation.png", dpi=150)
        plt.close()
        print(f"  Saved {outputs_dir / 'stage1_validation.png'}")
    except Exception as e:
        print("  Plot skip:", e)


def run_stage2(outputs_dir: Path):
    """Load IMS + stage1 labels, preprocess, train, evaluate, save metrics and plots."""
    from config import settings
    from src.ims_client import IMSClient
    from src.preprocessor import Preprocessor
    from src.predictor import PhotosynthesisPredictor
    import pandas as pd
    import numpy as np

    path_labels = settings.PROCESSED_DIR / "stage1_labels.csv"
    if not path_labels.exists():
        print("Stage 2: no stage1_labels.csv; run Stage 1 first.")
        return
    labels = pd.read_csv(path_labels, index_col=0, parse_dates=True).iloc[:, 0]
    labels.index = pd.to_datetime(labels.index, utc=True)
    labels = labels.dropna()

    client = IMSClient()
    ims_df = client.load_cached()
    if ims_df.empty:
        print("Stage 2: no IMS cache; run download_ims_data first.")
        return

    preproc = Preprocessor()
    merged = preproc.merge_ims_with_labels(ims_df, labels, timestamp_index_labels=True)
    if merged.empty:
        print("Stage 2: no overlap between IMS and labels. Ensure IMS cache covers sensor date range (e.g. run download_ims_data --from YYYY-MM-DD --to YYYY-MM-DD for that range).")
        return
    merged = preproc.create_time_features(merged)
    X_train, y_train, X_test, y_test = preproc.temporal_split(merged)
    if X_train.empty or X_test.empty:
        print("Stage 2: insufficient data after split.")
        return
    X_train_s = preproc.fit_transform_train(X_train)
    X_test_s = preproc.transform_test(X_test)
    # Ensure no sensor columns in features
    assert "A" not in X_train_s.columns and "timestamp_utc" not in X_train_s.columns

    predictor = PhotosynthesisPredictor()
    predictor.train(X_train_s, y_train)
    eval_df = predictor.evaluate(X_test_s, y_test)
    print("Stage 2 metrics:")
    print(eval_df.to_string(index=False))
    eval_df.to_csv(outputs_dir / "stage2_metrics.csv", index=False)

    predictor.plot_results(y_test, save_path=outputs_dir / "stage2_predictions.png")
    imp = predictor.get_feature_importance()
    if not imp.empty:
        imp.to_csv(outputs_dir / "stage2_feature_importance.csv", index=False)
        try:
            import matplotlib.pyplot as plt
            imp.set_index("feature")["importance"].plot(kind="barh", legend=False, figsize=(8, 4))
            plt.tight_layout()
            plt.savefig(outputs_dir / "stage2_feature_importance.png", dpi=150)
            plt.close()
        except Exception:
            pass
    print(f"  Saved metrics and plots to {outputs_dir}")


def run_stage2_timeseries(outputs_dir: Path):
    """Load merged IMS+labels, build lag features, train/eval per horizon."""
    from config import settings
    from src.ims_client import IMSClient
    from src.preprocessor import Preprocessor
    from src.ts_predictor import TimeSeriesPredictor
    import pandas as pd

    path_labels = settings.PROCESSED_DIR / "stage1_labels.csv"
    if not path_labels.exists():
        print("TS: no stage1_labels.csv; run Stage 1 first.")
        return
    labels = pd.read_csv(path_labels, index_col=0, parse_dates=True).iloc[:, 0]
    labels.index = pd.to_datetime(labels.index, utc=True)
    labels = labels.dropna()

    client = IMSClient()
    ims_df = client.load_cached()
    if ims_df.empty:
        print("TS: no IMS cache; run download_ims_data first.")
        return

    preproc = Preprocessor()
    merged = preproc.merge_ims_with_labels(ims_df, labels, timestamp_index_labels=True)
    if merged.empty:
        print("TS: no overlap between IMS and labels.")
        return
    merged = preproc.create_time_features(merged)

    ts_pred = TimeSeriesPredictor()
    merged = ts_pred.assign_season(merged)
    merged = ts_pred.identify_sessions(merged)
    merged = ts_pred.create_lag_features(merged)

    print("Training time-series models per horizon (per season)...")
    ts_metrics = ts_pred.train_all_horizons(merged)

    # Save full per-season breakdown
    ts_metrics.to_csv(outputs_dir / "ts_per_season_metrics.csv", index=False)

    # Load cross-sectional baseline for comparison (overall only)
    baseline_path = outputs_dir / "stage2_metrics.csv"
    if baseline_path.exists():
        baseline = pd.read_csv(baseline_path)
        comparison = ts_pred.get_comparison_with_baseline(baseline)
    else:
        comparison = ts_metrics[ts_metrics["season"] == "all"].copy()

    out_path = outputs_dir / "ts_horizon_comparison.csv"
    comparison.to_csv(out_path, index=False)
    print(f"  Saved {out_path}")
    print(f"  Saved {outputs_dir / 'ts_per_season_metrics.csv'}")


def main():
    from config import settings

    outputs_dir = settings.OUTPUTS_DIR
    outputs_dir.mkdir(parents=True, exist_ok=True)

    print("Running Stage 1 (sensors -> Farquhar -> labels)...")
    try:
        df_sensor, labels = run_stage1()
    except FileNotFoundError as e:
        print(f"Stage 1 skipped (missing data): {e}")
        df_sensor, labels = None, None
    except Exception as e:
        print(f"Stage 1 error: {e}")
        raise

    if labels is not None:
        print("Stage 1 validation...")
        validate_stage1(df_sensor, labels, outputs_dir)

    print("Running Stage 2 (IMS + labels -> train/eval)...")
    run_stage2(outputs_dir)

    print("Running Stage 2 time-series forecasting...")
    run_stage2_timeseries(outputs_dir)
    print("Done.")


if __name__ == "__main__":
    main()