#!/usr/bin/env python3
"""Generate a LightGBM forecasting script with leak-free features, a rolling-origin backtest and intervals.

The script it writes, for one series (time + value columns; --id-col for many series in long format):
  - puts the data on a complete grid (--freq D|W|MS|H), filling gaps with --fill zero|interpolate
  - builds features that only use values at least --horizon periods old ("direct" strategy with
    minimum lag = horizon: one model forecasts every step 1..H without recursion and without leakage):
    lags H..H+season-1, seasonal lags, rolling mean/std over windows ending H periods back, calendar
    features (day of week, month, day of year as sin/cos, week of year) and optional --holidays dates
  - optionally detrends with a linear trend fitted on each training window (--detrend: tree models cannot
    extrapolate trends) and adds it back to the predictions
  - backtests at --folds rolling origins (step = horizon), refitting each time, and writes
    backtest_forecasts.csv (origin,date,model,forecast) so ts_backtest.py --forecasts can rank it against
    the baselines on identical origins; prints MAE / WAPE and the seasonal-naive comparison
  - builds split-conformal prediction intervals from the backtest residuals per horizon step, pooled over
    +-3 neighbouring steps and scaled by each series' recent level (--coverage, default 0.8) and writes forecast.csv (date, forecast, lo, hi) from a model refit on all data
Needs pandas, numpy, lightgbm.

Usage:
    python gen_forecast_script.py --data sales.csv --time-col date --value-col units --freq D --horizon 28 -o forecast.py
    python gen_forecast_script.py --data sales.csv --freq D --horizon 28 --detrend --holidays holidays.csv
    python gen_forecast_script.py --data stores.csv --id-col store --value-col sales --freq W --horizon 8 --fill zero

With --id-col the script fits one model per series and pools the (level-relative) backtest residuals across
series for the intervals; for hundreds of related series, a single global model with the id as a feature
usually forecasts better (see references/ml-forecasting.md).

Exit codes: 0 written, 2 usage error.
"""

import argparse
import os
import sys

SEASON = {"D": 7, "W": 52, "MS": 12, "H": 24}

TEMPLATE = '''"""LightGBM forecast for {value!r} ({freq}, horizon {horizon}). Generated by gen_forecast_script.py.

Run:  python {script}
Then: python ts_backtest.py {data} --horizon {horizon} --folds {folds} --forecasts backtest_forecasts.csv
"""
import lightgbm as lgb
import numpy as np
import pandas as pd

DATA, TIME, VALUE, ID = {data!r}, {time!r}, {value!r}, {id_col!r}
FREQ, H, SEASON, FOLDS, FILL, DETREND, COVERAGE, SEED = {freq!r}, {horizon}, {season}, {folds}, {fill!r}, {detrend}, {coverage}, 0
HOLIDAYS = {holidays!r}

raw = pd.read_csv(DATA, parse_dates=[TIME])
if ID is None:
    raw["__id"] = "series"
    ID_COL = "__id"
else:
    ID_COL = ID
holidays = set(pd.to_datetime(pd.read_csv(HOLIDAYS).iloc[:, 0])) if HOLIDAYS else set()


def regular(g):
    s = g.set_index(TIME)[VALUE].sort_index()
    s = s[~s.index.duplicated(keep="last")].asfreq(FREQ)
    return s.fillna(0) if FILL == "zero" else s.interpolate(limit_direction="both")


series = {{k: regular(g) for k, g in raw.groupby(ID_COL)}}
LAGS = sorted(set(range(H, H + SEASON)) | {{SEASON * k for k in range(1, 9) if SEASON * k >= H}}
              | ({{364}} if FREQ == "D" else {{52}} if FREQ == "W" else set()))
WINDOWS = sorted({{SEASON, 2 * SEASON, 4 * SEASON}})


def features(s):
    df = pd.DataFrame({{"y": s}})
    for lag in LAGS:
        df[f"lag_{{lag}}"] = s.shift(lag)
    base = s.shift(H)                                    # every window ends H periods back: no leakage
    for w in WINDOWS:
        df[f"mean_{{w}}"] = base.rolling(w, min_periods=max(2, w // 2)).mean()
        df[f"std_{{w}}"] = base.rolling(w, min_periods=max(2, w // 2)).std()
    idx = s.index
    if FREQ in ("D", "H"):
        df["dow"] = idx.dayofweek
    df["month"] = idx.month
    doy = idx.dayofyear
    df["doy_sin"], df["doy_cos"] = np.sin(2 * np.pi * doy / 365.25), np.cos(2 * np.pi * doy / 365.25)
    df["week"] = idx.isocalendar().week.astype(int).values
    if holidays:
        df["holiday"] = idx.normalize().isin(holidays).astype(int)
    return df


def fit_predict(train_s, future_index):
    """Fit on train_s (a Series) and predict the H periods in future_index."""
    full = pd.concat([train_s, pd.Series(np.nan, index=future_index)])
    t = np.arange(len(full), dtype=float)
    trend = np.zeros(len(full))
    if DETREND:
        coef = np.polyfit(t[: len(train_s)], train_s.to_numpy(dtype=float), 1)
        trend = np.polyval(coef, t)
    target = full - trend
    feats = features(target)
    feats["t"] = t if not DETREND else 0.0
    X = feats.drop(columns="y")
    tr = feats.iloc[: len(train_s)].dropna(subset=["y"] + [f"lag_{{H}}"])
    model = lgb.LGBMRegressor(n_estimators=400, learning_rate=0.03, num_leaves=15, min_child_samples=10,
                              subsample=0.8, subsample_freq=1, colsample_bytree=0.8, random_state=SEED, verbose=-1)
    model.fit(tr.drop(columns="y"), tr["y"])
    pred = model.predict(X.iloc[len(train_s):]) + trend[len(train_s):]
    if (train_s >= 0).all():
        pred = np.maximum(pred, 0)
    return pd.Series(pred, index=future_index)


def level_of(s):
    recent = s.iloc[-4 * SEASON:].abs().mean()
    return recent if recent > 0 else 1.0


rows, resid = [], {{k: [] for k in range(H)}}
summary = []
for key, s in series.items():
    n = len(s)
    origins = [n - H - (FOLDS - 1 - i) * H for i in range(FOLDS)]
    origins = [o for o in origins if o > max(LAGS) + 2 * SEASON]
    for o in origins:
        train_s, test = s.iloc[:o], s.iloc[o:o + H]
        pred = fit_predict(train_s, test.index)
        snaive = np.array([train_s.iloc[len(train_s) - SEASON + (k % SEASON)] for k in range(len(test))])
        level = level_of(train_s)
        for k, (d, p, a) in enumerate(zip(test.index, pred, test)):
            rows.append((train_s.index[-1].date().isoformat(), d.date().isoformat(), "lightgbm", float(p)))
            resid[k].append((a - p) / level)                 # relative residuals: comparable across series
        summary.append((np.abs(pred.values - test.values).sum(), np.abs(snaive - test.values).sum(), np.abs(test.values).sum()))
    if not origins:
        print(f"{{key}}: too short for {{FOLDS}} folds with lags up to {{max(LAGS)}}")

pd.DataFrame(rows, columns=["origin", "date", "model", "forecast"]).to_csv("backtest_forecasts.csv", index=False)
if summary:
    e_m, e_s, tot = (sum(x) for x in zip(*summary))
    print(f"backtest: LightGBM WAPE {{e_m / tot:.1%}} vs seasonal naive {{e_s / tot:.1%}} "
          f"(skill {{1 - e_m / e_s:+.1%}}); wrote backtest_forecasts.csv")

alpha = (1 - COVERAGE) / 2
pooled = [sum((resid[j] for j in range(max(0, k - 3), min(H, k + 4))), []) for k in range(H)]   # +-3 steps
q_lo = [np.quantile(p, alpha) if p else 0.0 for p in pooled]
q_hi = [np.quantile(p, 1 - alpha) if p else 0.0 for p in pooled]
if min(len(p) for p in pooled) < 20:
    print(f"warning: only {{min(len(p) for p in pooled)}} residuals per horizon step for the intervals; "
          "add folds or series for reliable coverage")
out = []
for key, s in series.items():
    future = pd.date_range(s.index[-1], periods=H + 1, freq=FREQ)[1:]
    pred = fit_predict(s, future)
    level = level_of(s)
    for k, (d, p) in enumerate(pred.items()):
        lo, hi = p + q_lo[k] * level, p + q_hi[k] * level
        if (s >= 0).all():
            lo = max(lo, 0.0)
        out.append((key, d.date().isoformat(), p, lo, hi))
fc = pd.DataFrame(out, columns=["series", "date", "forecast", "lo", "hi"])
if ID is None:
    fc = fc.drop(columns="series")
fc.to_csv("forecast.csv", index=False)
print(f"wrote forecast.csv ({{H}} periods, {{COVERAGE:.0%}} conformal intervals from about {{len(pooled[H // 2])}} "
      "backtest residuals per step, pooled over +-3 steps)")
'''


def main():
    ap = argparse.ArgumentParser(description="Generate a LightGBM forecasting script.", epilog=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--data", required=True)
    ap.add_argument("--time-col", default="date")
    ap.add_argument("--value-col", default="y")
    ap.add_argument("--id-col", help="series id column for long-format panels")
    ap.add_argument("--freq", choices=sorted(SEASON), default="D")
    ap.add_argument("--horizon", type=int, required=True)
    ap.add_argument("--season", type=int)
    ap.add_argument("--folds", type=int, default=5)
    ap.add_argument("--fill", choices=["interpolate", "zero"], default="interpolate")
    ap.add_argument("--detrend", action="store_true", help="linear detrending per training window")
    ap.add_argument("--coverage", type=float, default=0.8)
    ap.add_argument("--holidays", help="CSV whose first column lists holiday dates")
    ap.add_argument("-o", "--output")
    a = ap.parse_args()
    if a.horizon < 1 or a.folds < 1 or not 0.5 <= a.coverage < 1:
        print("error: need --horizon >= 1, --folds >= 1 and 0.5 <= --coverage < 1", file=sys.stderr)
        return 2
    text = TEMPLATE.format(data=a.data, time=a.time_col, value=a.value_col, id_col=a.id_col, freq=a.freq,
                           horizon=a.horizon, season=a.season or SEASON[a.freq], folds=a.folds, fill=a.fill,
                           detrend=a.detrend, coverage=a.coverage, holidays=a.holidays,
                           script=os.path.basename(a.output or "forecast.py"))
    if a.output:
        with open(a.output, "w", encoding="utf-8") as f:
            f.write(text)
        print(f"wrote {a.output}: {a.freq}, horizon {a.horizon}, {a.folds} folds"
              + (", detrended" if a.detrend else ""), file=sys.stderr)
    else:
        sys.stdout.write(text)
    return 0


if __name__ == "__main__":
    sys.exit(main())