# Copyright 2026 Kemal Yaylali
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Where the blend weights and the prediction intervals are allowed to learn.

Both are statements about how this head's issued forecasts actually turned out.
A refit is not an outcome, it is the same week of weather being read again, and
a matured forecast that stays readable for an hour is one outcome and not
twelve. Getting either wrong does not raise: it quietly multiplies the evidence
until Hedge saturates and the ACI integrator pins against its clips.
"""

from __future__ import annotations

import numpy as np

from ashvale.config import CONFIG, load_config
from ashvale.models.nowcast import MEMBERS, NowcastEnsemble
from ashvale.station import Station
from ashvale.storage import Store


def _synthetic(n: int = 700, seed: int = 5):
    rng = np.random.default_rng(seed)
    g = CONFIG.model.grid_s
    ts = np.arange(n) * g + 1.7554e9
    cols = {
        "temperature": 22 + 3 * np.sin(np.arange(n) / 288.0) + 0.1 * rng.normal(size=n),
        "humidity": 50 + 9 * np.cos(np.arange(n) / 288.0),
        "pressure": 1013 + 2 * np.sin(np.arange(n) / 600.0),
        "lux": np.clip(400 * np.sin(np.arange(n) / 288.0), 0, None),
    }
    X = rng.normal(size=(n, 33))
    X[:, 0] = 1.0
    return ts, cols, X, np.ones(n, dtype=bool)


# ------------------------------------------------ a refit is not an outcome

def test_refitting_does_not_touch_the_weights_or_the_calibrator():
    """The defect this file exists for.

    Measured on 8.2 days of a real station, fit() had put 977,078 Hedge updates
    through the 15 minute head from 758 distinct supervised pairs, and 296,715
    through the 1 day head from 12. Hedge is multiplicative, so an edge far too
    small to be real compounds to certainty.
    """
    _, cols, X, valid = _synthetic()
    ens = NowcastEnsemble(CONFIG.model.targets, CONFIG.model.horizons_s, CONFIG.model)
    before = {k: h.weights.copy() for k, h in ens.heads.items()}

    for _ in range(5):
        ens.fit(X, valid, cols)

    for k, head in ens.heads.items():
        assert np.allclose(head.weights, before[k]), \
            f"{k} blend weights moved during a refit"
        assert len(head.conformal.scores) == 0, \
            f"{k} fed {len(head.conformal.scores)} scores to the calibrator from a refit"
        assert head.n_scored == 0
        # The regression itself must still be learning.
        assert head.model.n_updates > 0


def test_an_outcome_does_move_them():
    """The other half: verify()'s channel has to work, or nothing ever learns."""
    _, cols, X, valid = _synthetic()
    ens = NowcastEnsemble(CONFIG.model.targets, CONFIG.model.horizons_s, CONFIG.model)
    ens.fit(X, valid, cols)
    head = ens.heads[("temperature", 900)]
    before = head.weights.copy()

    # Persistence right, everything else wrong, so the weights must move to it.
    members = np.array([22.0, 30.0, 30.0, 30.0])
    for k in range(20):
        head.observe_outcome(members, 22.0, covered=True, valid_ts=1e9 + k * 900)

    assert head.n_scored == 20
    assert len(head.conformal.scores) == 20
    assert head.weights[MEMBERS.index("persistence")] > before[MEMBERS.index("persistence")]
    assert head.weights[MEMBERS.index("climatology")] < before[MEMBERS.index("climatology")]


# ------------------------------------------- overlapping outcomes are one fact

def test_hedge_ignores_outcomes_that_overlap_the_last_one_it_took():
    """Forecasts are issued every retrain tick, so at the 1 day horizon about
    two hundred a day resolve against very nearly the same outcome. The
    conformal window can absorb that, being an order statistic. Exponentiated
    gradient cannot."""
    _, cols, X, valid = _synthetic()
    ens = NowcastEnsemble(CONFIG.model.targets, CONFIG.model.horizons_s, CONFIG.model)
    head = ens.heads[("temperature", 86400)]
    members = np.array([22.0, 30.0, 30.0, 30.0])

    t0 = 1.7554e9
    for k in range(50):                    # 50 forecasts, 10 minutes apart
        head.observe_outcome(members, 22.0, covered=True, valid_ts=t0 + k * 600)

    assert head.n_scored == 1, \
        f"Hedge took {head.n_scored} of 50 overlapping outcomes at a 1 d horizon"
    # The calibrator still sees all of them: a quantile over duplicated scores
    # is over-confident about its sample size, not wrong about its value.
    assert len(head.conformal.scores) == 50

    head.observe_outcome(members, 22.0, covered=True, valid_ts=t0 + 86400)
    assert head.n_scored == 2, "a genuinely new outcome must be taken"


def test_the_decimation_clock_survives_a_restart():
    """Otherwise every deploy hands the weights a free duplicate."""
    ens = NowcastEnsemble(CONFIG.model.targets, CONFIG.model.horizons_s, CONFIG.model)
    head = ens.heads[("pressure", 3600)]
    head.observe_outcome(np.full(len(MEMBERS), 1013.0), 1013.5,
                         covered=True, valid_ts=1.7554e9)

    back = NowcastEnsemble(CONFIG.model.targets, CONFIG.model.horizons_s, CONFIG.model)
    back.load_dict(ens.to_dict())
    assert back.heads[("pressure", 3600)].last_hedge_ts == head.last_hedge_ts

    # A head that has never scored must not refuse its first outcome.
    fresh = NowcastEnsemble(CONFIG.model.targets, CONFIG.model.horizons_s, CONFIG.model)
    fresh.load_dict(NowcastEnsemble(CONFIG.model.targets,
                                    CONFIG.model.horizons_s, CONFIG.model).to_dict())
    h = fresh.heads[("pressure", 3600)]
    h.observe_outcome(np.full(len(MEMBERS), 1013.0), 1013.5,
                      covered=True, valid_ts=1.7554e9)
    assert h.n_scored == 1


# -------------------------------------------------------------- persistence

def test_forecast_members_round_trip(tmp_path):
    """The learned member cannot be recomputed at maturity, the RLS has moved
    on, so it has to be written down at issue time."""
    store = Store(str(tmp_path / "f.db"))
    members = {"persistence": 21.0, "climatology": 22.0,
               "learned": 23.0, "setpoint": 24.0}
    store.insert_forecast(1000.0, 900, "temperature", 22.5, 21.0, 24.0,
                          "ensemble", members=members)
    row = store.due_forecasts(now=2000.0)[0]
    assert Station._members_of(row).tolist() == [21.0, 22.0, 23.0, 24.0]
    assert row["scored"] == 0

    store.mark_forecasts_scored([(1000.0, 900, "temperature")])
    assert store.due_forecasts(now=2000.0)[0]["scored"] == 1


def test_a_forecast_from_before_the_columns_existed_is_not_fatal(tmp_path):
    """It still deserves a conformal score. It just cannot move the weights."""
    store = Store(str(tmp_path / "old.db"))
    store.insert_forecast(1000.0, 900, "temperature", 22.5, 21.0, 24.0, "ensemble")
    row = store.due_forecasts(now=2000.0)[0]
    assert Station._members_of(row) is None


def test_migration_adds_the_forecast_columns_to_a_live_table(tmp_path):
    """Same hazard as the telemetry columns: CREATE TABLE IF NOT EXISTS is a
    no-op against a table that already exists."""
    import sqlite3

    path = str(tmp_path / "legacy.db")
    with sqlite3.connect(path) as c:
        c.execute("CREATE TABLE forecasts ("
                  "issued_ts REAL NOT NULL, valid_ts REAL NOT NULL, "
                  "horizon_s INTEGER NOT NULL, target TEXT NOT NULL, "
                  "mu REAL, lo REAL, hi REAL, model TEXT, "
                  "PRIMARY KEY (issued_ts, horizon_s, target))")
        c.execute("INSERT INTO forecasts VALUES (1.0, 901.0, 900, 'temperature',"
                  " 22.0, 21.0, 23.0, 'ensemble')")

    store = Store(path)
    with sqlite3.connect(path) as c:
        have = {r[1] for r in c.execute("PRAGMA table_info(forecasts)")}
    for col in ("m_persistence", "m_climatology", "m_learned", "m_setpoint", "scored"):
        assert col in have, f"migration missed {col}"

    rows = store.due_forecasts(now=2000.0)
    assert len(rows) == 1, "the existing row must survive the migration"
    assert rows[0]["scored"] == 0
    assert Station._members_of(rows[0]) is None


# --------------------------------------------------------------- end to end

def test_verify_teaches_from_a_matured_forecast_exactly_once(tmp_path):
    """A matured forecast stays readable for an hour so the scorecard can
    aggregate a rolling window, which means verify() sees it about twelve
    times. Aggregating it twelve times is harmless. Teaching from it twelve
    times is how the calibrator ends up believing it holds four hundred
    independent scores when it holds one."""
    import time

    cfg = load_config()
    cfg.storage.db_path = str(tmp_path / "verify.db")
    # Without this the station loads the developer's own saved state, whose
    # heads already carry six figures of n_scored.
    cfg.storage.state_dir = str(tmp_path)
    st = Station(cfg)

    now = time.time()
    # The record has to run past the forecast's validity time, or verify()
    # finds no truth to score it against.
    for k in range(61):
        ts = now - 1800 + k * 30
        st.store.insert_telemetry({"ts": ts, "temp_raw": 21.0, "temp_smooth": 21.0,
                                   "hum": 50.0, "hum_smooth": 50.0,
                                   "press": 1013.0, "press_slp": 1013.0})

    issued = now - 1500
    st.store.insert_forecast(issued, 900, "temperature", 21.4, 20.9, 21.9,
                             "ensemble",
                             members={"persistence": 21.0, "climatology": 21.6,
                                      "learned": 21.5, "setpoint": 21.0})

    head = st.nowcast.heads[("temperature", 900)]
    first = st.verify()
    assert first["learned"] == 1, first
    assert head.n_scored == 1
    assert len(head.conformal.scores) == 1

    for _ in range(5):
        again = st.verify()
        assert again["learned"] == 0, "the same outcome was taught again"
    assert head.n_scored == 1
    assert len(head.conformal.scores) == 1
    # and it is still being aggregated into the scorecard
    assert again["scored"] >= 1


def test_state_from_before_the_fix_is_not_trusted():
    """The saturated weights do not decay on their own.

    A pre-fix state file holds weights produced by the same week of weather
    read about a thousand times. Hedge needs roughly twenty independent
    outcomes to climb back off its 1e-4 floor, and the 1 d head sees one a day,
    so leaving them in place would mean three weeks of a forecast pinned to
    whichever member the replay happened to favour.
    """
    ens = NowcastEnsemble(CONFIG.model.targets, CONFIG.model.horizons_s, CONFIG.model)
    head = ens.heads[("temperature", 86400)]
    for k in range(400):
        head.observe_outcome(np.array([22.0, 30.0, 30.0, 30.0]), 22.0,
                             covered=False, valid_ts=1.7554e9 + k * 86400)
    saved = ens.to_dict()
    assert head.weights.max() > 0.9 and head.conformal.alpha < 0.02

    for h in saved["heads"]:
        h.pop("last_hedge_ts")                     # a state file from before

    back = NowcastEnsemble(CONFIG.model.targets, CONFIG.model.horizons_s, CONFIG.model)
    back.load_dict(saved)
    b = back.heads[("temperature", 86400)]
    assert np.allclose(b.weights, 1.0 / len(MEMBERS)), "saturated weights were trusted"
    assert b.n_scored == 0
    assert b.conformal.alpha == b.conformal.alpha_target, "the ACI integrator kept its wind-up"
    # but the residual magnitudes are kept, or the long horizons spend nine
    # days on the Gaussian fallback
    assert len(b.conformal.scores) == 400

    # A state file written after the fix must survive untouched.
    keep = NowcastEnsemble(CONFIG.model.targets, CONFIG.model.horizons_s, CONFIG.model)
    keep.load_dict(ens.to_dict())
    assert np.allclose(keep.heads[("temperature", 86400)].weights, head.weights)