import pickle
from pathlib import Path
import numpy as np
import pandas as pd
import warnings
warnings.simplefilter("ignore")

from statsmodels.tsa.stattools import adfuller, kpss
from arch.unitroot import DFGLS

DATA = Path(__file__).resolve().parent.parent / "data" / "netlw_10min_djf.pkl"


def load():
    d = pickle.load(open(DATA, "rb"))
    return pd.Series(d["netlw_10min"], index=pd.DatetimeIndex(d["time"]))


def detrend_keepmean(x, t_hours):
    """Mean-preserving linear detrend (removes only the slope)."""
    sl, ic = np.polyfit(t_hours, x, 1)
    return x - sl * (t_hours - t_hours.mean()), sl


def run(series, label):
    s = series.dropna()
    x = s.values
    t_h = (s.index.astype("int64") / 1e9 / 3600.0).values
    xd, slope = detrend_keepmean(x, t_h)
    # detrend correlation r (fitted line vs data) - support for trend-stationary
    fit = slope * (t_h - t_h.mean())
    r = np.corrcoef(fit, x - x.mean())[0, 1] if np.std(fit) > 0 else 0.0

    adf_stat, adf_p, *_ = adfuller(xd, regression="ct", autolag="AIC")
    kp_stat, kp_p, *_ = kpss(xd, regression="ct", nlags="auto")
    gls = DFGLS(xd, trend="ct")

    print(f"\n== {label}  (n={len(x):,}) ==")
    print(f"  ADF     stat={adf_stat:8.2f}  p={adf_p:.2e}   "
          f"({'reject unit root' if adf_p < 0.05 else 'cannot reject'})")
    print(f"  ADF-GLS stat={gls.stat:8.2f}  p={gls.pvalue:.2e}   "
          f"({'reject unit root' if gls.pvalue < 0.05 else 'cannot reject'})")
    print(f"  KPSS    stat={kp_stat:8.2f}  p~{kp_p:.3f}    "
          f"({'reject stationarity' if kp_p < 0.05 else 'cannot reject stationarity'})")
    print(f"  detrend slope={slope*24*1e3:.3f} mW m^-2 day^-1,  detrend r={r:.3f}")


def main():
    s10 = load()
    run(s10, "10-min")
    run(s10.resample("1h").mean(), "1-hour")
    print("\nReading: ADF/ADF-GLS reject the unit root (conventional support); KPSS is "
          "the load-bearing test and marginally rejects strict stationarity at both "
          "cadences -> trend-stationary with a slow climate trend (detrend r small).")


if __name__ == "__main__":
    main()