"""Train the EnergyPredictor (XGBoost + LR fallback) from
``Data/energy_weather_merged.csv`` and write the bundle to
``Data/energy_predictor_model.pkl``.

The bundle schema matches what ``src/energy_predictor.py:EnergyPredictor.__init__``
loads with pickle:
  - xgb_model        : fitted xgboost.XGBRegressor
  - xgb_features     : list of feature column names
  - lr_fallback      : fitted sklearn.linear_model.LinearRegression
  - lr_features      : list of feature column names
  - test_mape_pct, test_mae_kwh, train_rows, test_rows, trained_at, source

Temporal hold-out: last ``--holdout-days`` of the merged CSV (default 14).

Usage:
    python -m scripts.train_energy_predictor
    python -m scripts.train_energy_predictor --holdout-days 21
"""

from __future__ import annotations

import argparse
import pickle
import sys
from datetime import datetime, timezone
from pathlib import Path

import numpy as np
import pandas as pd

_PROJECT_ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(_PROJECT_ROOT))

_DATA_CSV = _PROJECT_ROOT / "Data" / "energy_weather_merged.csv"
_OUT_PKL = _PROJECT_ROOT / "Data" / "energy_predictor_model.pkl"

# Features must match what EnergyPredictor.predict_hourly constructs at inference.
# Tracker features were trialed but worsened the May hold-out bias from +50%
# to +75% because the training distribution (Dec-Apr, avg tracker -15°) is
# operationally different from the hold-out (May, avg tracker -51° due to
# active shading interventions). Reverted until 2026 growing-season data
# accumulates — see 2026-bigml-tracker-feature in context/3_todo.md.
_XGB_FEATURES = ["GSR", "airTemperature", "sin_elevation", "hour", "clearness", "windSpeed"]
_LR_FEATURES = ["GSR", "airTemperature"]
_TARGET = "production_kwh"


def _load() -> pd.DataFrame:
    if not _DATA_CSV.exists():
        raise FileNotFoundError(
            f"{_DATA_CSV} not found. Run scripts.refresh_energy_data first."
        )
    df = pd.read_csv(_DATA_CSV)
    df["timestamp_utc"] = pd.to_datetime(df["timestamp_utc"], utc=True)
    df = df.set_index("timestamp_utc").sort_index()

    # Derived helper kept for future re-enabling of tracker features.
    if "tracker_angle_mean" in df.columns:
        df["abs_tracker_angle"] = df["tracker_angle_mean"].abs()

    needed = set(_XGB_FEATURES) | {_TARGET}
    missing = needed - set(df.columns)
    if missing:
        raise ValueError(f"Merged CSV is missing columns: {sorted(missing)}")
    df = df.dropna(subset=list(needed))
    return df


def _mape(y_true: np.ndarray, y_pred: np.ndarray, floor_kwh: float = 1.0) -> float:
    """Mean absolute percentage error, excluding hours below `floor_kwh`
    (nighttime / very low production where small absolute errors blow up %)."""
    mask = y_true >= floor_kwh
    if not mask.any():
        return float("nan")
    return float(np.mean(np.abs(y_pred[mask] - y_true[mask]) / y_true[mask]) * 100)


def main() -> None:
    p = argparse.ArgumentParser(description="Train the EnergyPredictor.")
    p.add_argument("--holdout-days", type=int, default=14,
                   help="Last N days held out for validation (default 14).")
    p.add_argument("--n-estimators", type=int, default=400)
    p.add_argument("--max-depth", type=int, default=6)
    p.add_argument("--learning-rate", type=float, default=0.05)
    args = p.parse_args()

    from sklearn.ensemble import GradientBoostingRegressor
    from sklearn.linear_model import LinearRegression
    from sklearn.metrics import mean_absolute_error

    # XGBoost requires libomp on macOS which isn't installed in this venv;
    # GradientBoostingRegressor has the same fit/predict API and is bundled
    # into the same `xgb_model` slot for backward compatibility with
    # EnergyPredictor.__init__.

    df = _load()
    print(f"Loaded {len(df):,} hourly rows  {df.index.min()}  →  {df.index.max()}")

    split = df.index.max() - pd.Timedelta(days=args.holdout_days)
    train_df = df.loc[df.index < split]
    test_df = df.loc[df.index >= split]
    print(f"Train: {len(train_df):,} rows   Hold-out: {len(test_df):,} rows  "
          f"(split at {split})")

    if len(test_df) < 24:
        raise RuntimeError("Hold-out window too small. Increase --holdout-days "
                           "or refresh more recent data.")

    # ---------- Gradient Boosting (stored under xgb_model slot) ----------
    xgb = GradientBoostingRegressor(
        n_estimators=args.n_estimators,
        max_depth=args.max_depth,
        learning_rate=args.learning_rate,
        random_state=42,
    )
    xgb.fit(train_df[_XGB_FEATURES], train_df[_TARGET])
    pred_test = np.clip(xgb.predict(test_df[_XGB_FEATURES]), 0, None)

    mae_xgb = mean_absolute_error(test_df[_TARGET], pred_test)
    mape_xgb = _mape(test_df[_TARGET].to_numpy(), pred_test)

    # ---------- Linear fallback ----------
    lr = LinearRegression()
    lr.fit(train_df[_LR_FEATURES], train_df[_TARGET])
    pred_test_lr = np.clip(lr.predict(test_df[_LR_FEATURES]), 0, None)
    mae_lr = mean_absolute_error(test_df[_TARGET], pred_test_lr)
    mape_lr = _mape(test_df[_TARGET].to_numpy(), pred_test_lr)

    print("\n=== Hold-out validation ===")
    print(f"  XGBoost   MAE = {mae_xgb:.3f} kWh   MAPE (>1 kWh) = {mape_xgb:.2f}%")
    print(f"  Linear    MAE = {mae_lr:.3f} kWh   MAPE (>1 kWh) = {mape_lr:.2f}%")

    # Sanity: daily totals (filter days where actual is <10 kWh — likely partial)
    test_daily = test_df.assign(pred=pred_test).resample("1D").agg(
        actual_kwh=(_TARGET, "sum"),
        predicted_kwh=("pred", "sum"),
    )
    test_daily = test_daily[test_daily["actual_kwh"] >= 10]
    test_daily["err_pct"] = (
        (test_daily["predicted_kwh"] - test_daily["actual_kwh"]) / test_daily["actual_kwh"] * 100
    )
    print("\n=== Daily hold-out (XGBoost, actual ≥ 10 kWh) ===")
    print(test_daily.round(1).to_string())
    if len(test_daily) > 0:
        mean_abs_err = float(test_daily["err_pct"].abs().mean())
        mean_signed_err = float(test_daily["err_pct"].mean())
        print(f"\nMean abs daily error: {mean_abs_err:.1f}%  "
              f"(signed: {mean_signed_err:+.1f}% — positive = over-prediction)")
    else:
        mean_abs_err = mean_signed_err = float("nan")

    bundle = {
        "xgb_model":      xgb,
        "xgb_features":   _XGB_FEATURES,
        "lr_fallback":    lr,
        "lr_features":    _LR_FEATURES,
        "test_mape_pct":          round(mape_xgb, 3),
        "test_mae_kwh":           round(mae_xgb, 3),
        "daily_mae_pct":          (round(mean_abs_err, 1) if mean_abs_err == mean_abs_err else None),
        "daily_signed_err_pct":   (round(mean_signed_err, 1) if mean_signed_err == mean_signed_err else None),
        "train_rows":     int(len(train_df)),
        "test_rows":      int(len(test_df)),
        "train_range":    (str(train_df.index.min()), str(train_df.index.max())),
        "test_range":     (str(test_df.index.min()), str(test_df.index.max())),
        "trained_at":     datetime.now(tz=timezone.utc).isoformat(),
        "source":         str(_DATA_CSV.relative_to(_PROJECT_ROOT)),
    }
    with open(_OUT_PKL, "wb") as f:
        pickle.dump(bundle, f)
    print(f"\nSaved bundle → {_OUT_PKL}")


if __name__ == "__main__":
    main()