"""
Predict chart points from sample BAS data.

Run:
    python forecast.py

Output:
    forecast_output.json  (test.php reads this for dashed forecast lines)
"""

import json
import re
from pathlib import Path

import pandas as pd
from sklearn.linear_model import Ridge
from sklearn.metrics import mean_absolute_error
from skforecast.recursive import ForecasterRecursive

# --- settings you can change ---
FUTURE_STEPS = 100                # how many readings to predict ahead
LAGS = 25                        # how many past readings the model looks at each step

# (sample file, point names) — one model per point
JOBS = [
    ("js/veris-sample.js", ["P", "Ia", "Ib", "Ic"]),
    ("js/rtu-sample.js", ["ServRmTmp"]),  # RTU room temp (spark chart later)
]

folder = Path(__file__).parent


def load_sample_js(path: str) -> dict:
    text = (folder / path).read_text(encoding="utf-8")
    return json.loads(re.search(r"=\s*(\{.*\})\s*;?\s*$", text, re.DOTALL).group(1))


def forecast_point(data: dict, point: str) -> dict:
    """Train one model for one point. Returns mae + future [{x, y}, ...]."""

    # 1) BUILD TIME SERIES from [{x, y}, ...] rows
    rows = data[point]
    times = pd.to_datetime([r["x"] for r in rows], unit="ms")  # x = timestamp in milliseconds
    values = [float(r["y"]) for r in rows]                     # y = value as a number
    series = pd.Series(values, index=times, name=point).sort_index()  # oldest first
    step = series.index.to_series().diff().median()            # typical gap (~2 min Veris, ~32s RTU)
    series = series.resample(step).mean().interpolate()        # even spacing so lags line up

    print(f"\n{point} — {len(series)} readings")

    # 2) TRAIN + TEST on the last 20% (see how wrong the model is)
    split = int(len(series) * 0.8)
    train, test = series.iloc[:split], series.iloc[split:]  # first 80% train, last 20% test

    model = ForecasterRecursive(estimator=Ridge(), lags=LAGS)  # uses the lag values to make new predictions make sense
    model.fit(y=train)
    guesses = model.predict(steps=len(test))  # try to predict the hidden 20%
    guesses.index = test.index
    mae = mean_absolute_error(test, guesses)  # average error from the test data
    print(f"  MAE (last 20%): {mae:.2f}")

    # 3) PREDICT THE FUTURE (retrain on all data, then forecast ahead)
    model.fit(y=series)
    future = model.predict(steps=FUTURE_STEPS)
    future.index = pd.date_range(series.index[-1] + step, periods=FUTURE_STEPS, freq=step)

    # 4) FORMAT for charts — same {x, y} shape as veris-sample.js
    return {
        "mae": round(mae, 2),
        "future": [
            {"x": int(t.timestamp() * 1000), "y": round(float(v), 2)}
            for t, v in future.items()
        ],
    }


output = {}
for js_path, points in JOBS:
    data = load_sample_js(js_path)
    for point in points:
        output[point] = forecast_point(data, point)

out_file = folder / "forecast_output.json"
out_file.write_text(json.dumps(output, indent=2), encoding="utf-8")
print(f"\nSaved forecasts for: {', '.join(output.keys())}")