# -*- coding: utf-8 -*-
"""R1-R — CONTROLLED MODEL EXECUTION at FROZEN origin lists.

Reads the frozen R1-R origin CSVs (NEVER regenerates them), runs the frozen
models (Naive, Drift, ARIMA(1,0,0), SAX, Hybrid) at exactly those origins,
and persists per-origin prediction evidence.

NO metrics / verdicts / optimization are produced here.

The fast SAX path is the committed equivalence-validated path (matches the
frozen SaxAnalogForecaster after DEFECT-001 z-normalization fix) and is used
for tractability exactly as in the historical frozen runner; a small equality
spot check (M1 prefix) confirms parity.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import os
import subprocess
import sys

import numpy as np

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if ROOT not in sys.path:
    sys.path.insert(0, ROOT)

from src.forecasting.target import ForecastContext, data_atr
from src.forecasting.interface import Series
from src.arima import ArimaConfig, ArimaModel
from src.sax import SaxConfig, SaxAnalogForecaster
from src.sax.transform import sax_encode, z_normalize, _bp_for
from src.baselines import NaiveBaseline, NaiveConfig, DriftBaseline, DriftConfig
from src.hybrid import make_hybrid_record
from src.forecasting.record import ForecastRecord

RAW = os.path.join(ROOT, "results", "R1_real_data", "XAUUSDc_M1_raw.json")
ORIG_DIR = os.path.join(ROOT, "results", "R1_R")
CONFIG = os.path.join(ROOT, "configs", "default.json")
ATR_PERIOD = 20
PROTOCOL_COMMIT = "4da3788"
EXPECTED_CONFIG_HASH = "ef1e3fd55b9808cf"
EXPECTED_DATASET_HASH = "80e4b52b0df0c6e348ce2002348279eee9f6e8a61c3109efe82fe475ed3cbc19"

PROTOCOLS = {"M1": {"H": 15, "tf": 1, "validate": True},
             "M5": {"H": 3, "tf": 5, "validate": False},
             "M15": {"H": 1, "tf": 15, "validate": False}}
MODELS = ["naive", "drift", "arima", "sax", "hybrid"]


def model_state(m, rec):
    return {"naive": rec.forecast_direction,
            "drift": rec.forecast_direction,
            "arima": rec.arima_state,
            "sax": rec.sax_state,
            "hybrid": rec.hybrid_state}[m]


def sha256(path):
    with open(path, "rb") as fh:
        return hashlib.sha256(fh.read()).hexdigest()


def fnum(v):
    if v is None:
        return ""
    try:
        f = float(v)
        if f != f:      # NaN
            return ""
        return repr(round(f, 12))
    except Exception:
        return ""


def load_raw(path):
    with open(path, "r", encoding="utf-8") as fh:
        return json.load(fh)


def sorted_unique(bars):
    seen = {}
    for b in bars:
        seen[b["time"]] = b
    return [seen[k] for k in sorted(seen.keys())]


def aggregate(bars, k):
    out = []
    for i in range(0, len(bars), k):
        if i + k > len(bars):
            break
        block = bars[i:i + k]
        out.append({"time": block[-1]["time"],
                    "open": float(block[0]["open"]),
                    "high": max(float(x["high"]) for x in block),
                    "low": min(float(x["low"]) for x in block),
                    "close": float(block[-1]["close"])})
    return out


def mk_series(sub, symbol, timeframe):
    return Series(symbol=symbol, timeframe=timeframe,
                  timestamp=[b["time"] for b in sub],
                  open=[float(b["open"]) for b in sub],
                  high=[float(b["high"]) for b in sub],
                  low=[float(b["low"]) for b in sub],
                  close=[float(b["close"]) for b in sub])


def read_origins(tf):
    rows = []
    with open(os.path.join(ORIG_DIR, f"origins_{tf}.csv"), encoding="utf-8") as fh:
        header = fh.readline().strip().split(",")
        for line in fh:
            line = line.strip()
            if not line:
                continue
            p2 = dict(zip(header, line.split(",")))
            rows.append({
                "origin_index": int(p2["origin_index"]),
                "origin_timestamp": p2["origin_timestamp"],
                "training_start": int(p2["training_start"]),
                "training_end": int(p2["training_end"]),
                "target_start": int(p2["target_start"]),
                "target_end": int(p2["target_end"]),
                "horizon": int(p2["horizon"]),
            })
    return rows


# ---- fast SAX path (committed, equivalence-validated) ----
def precompute_words(close, W, word_len, alpha):
    bp = _bp_for(alpha)
    bounds = np.concatenate(([-2.0], bp, [2.0]))
    mid = (bounds[:-1] + bounds[1:]) / 2.0
    n_w = len(close) - W + 1
    sym_idx = np.zeros((n_w, word_len), dtype=np.int32)
    for j in range(n_w):
        q = z_normalize(close[j:j + W])
        out = sax_encode(q, word_len, alpha)
        sym_idx[j] = [ord(ch) - ord("a") for ch in out]
    return sym_idx, mid


def fast_sax_record(series, origin, ctx, sax_cfg, sym_idx, mid):
    close = series.closes()
    W = sax_cfg.window_length
    H = ctx.horizon
    rec = ForecastRecord(symbol=ctx.symbol, timeframe=ctx.timeframe, horizon=H,
                         variant="C_sax", model_id="C_sax", model_version="0.1.0",
                         forecast_origin=origin, outcome_boundary=origin + H)
    if origin + 1 < W:
        rec.sax_state, rec.sax_analog_count = "insufficient", 0
        return rec
    q = z_normalize(close[origin - W + 1:origin + 1])
    q_word = sax_encode(q, sax_cfg.word_length, sax_cfg.alphabet_size)
    rec.sax_word = q_word
    q_idx = np.asarray([ord(ch) - ord("a") for ch in q_word], np.int32)
    last_allowed = origin - H
    if last_allowed < W - 1:
        rec.sax_state, rec.sax_analog_count = "insufficient", 0
        return rec
    cand = sym_idx[:last_allowed - (W - 1) + 1]
    Mdiff = np.abs(mid[None, :] - mid[:, None]).astype(float)
    dist = np.zeros(cand.shape[0])
    for w in range(sax_cfg.word_length):
        dist += Mdiff[q_idx[w], cand[:, w]]
    idx = np.where(dist <= sax_cfg.distance_threshold)[0]
    if len(idx) < sax_cfg.min_analogs:
        rec.sax_state, rec.sax_analog_count = "insufficient", 0
        return rec
    k = min(sax_cfg.top_k_analogs, len(idx))
    order = np.argsort(dist[idx])[:k]
    sel_rows = idx[order]
    sel_j = [(W - 1) + int(r) for r in sel_rows]
    atr = data_atr(close, ctx.atr_period)
    out_v = np.asarray([(close[j + H] - close[j]) / atr[j] for j in sel_j], float)
    rec.sax_analog_count = int(len(out_v))
    rec.sax_distance_metric = "midpoint_mindist"
    rec.sax_distance_median = float(np.median(dist[sel_rows]))
    rec.sax_median = float(np.median(out_v))
    rec.sax_P25 = float(np.percentile(out_v, 25))
    rec.sax_P75 = float(np.percentile(out_v, 75))
    rec.sax_up_rate = float(np.mean(out_v > 0))
    med = rec.sax_median
    thr = sax_cfg.neutral_threshold
    if med > thr:
        rec.sax_state = "bullish"
    elif med < -thr:
        rec.sax_state = "bearish"
    else:
        rec.sax_state = "no_edge" if abs(rec.sax_up_rate - 0.5) <= 0.10 else "neutral"
    rec.normalized_expected_return = rec.sax_median
    rec.forecast_direction = {"bullish": "LONG", "bearish": "SHORT"}.get(rec.sax_state, "NEUTRAL")
    rec.confidence = float(abs(rec.sax_up_rate - 0.5))
    rec.uncertainty = float(rec.sax_P75 - rec.sax_P25) if len(out_v) >= 2 else None
    return rec


def validate_fast_sax(series, ctx, sax_cfg, sym_idx, mid, n=40, seed=7, cap=8000):
    full_close = np.asarray(series.close, float)
    bound = min(cap, len(full_close))
    s_small = Series(symbol=series.symbol, timeframe=series.timeframe,
                     timestamp=list(series.timestamp[:bound]),
                     open=list(float(x) for x in series.open[:bound]),
                     high=list(float(x) for x in series.high[:bound]),
                     low=list(float(x) for x in series.low[:bound]),
                     close=list(float(x) for x in series.close[:bound]))
    frozen = SaxAnalogForecaster(sax_cfg)
    rng = np.random.default_rng(seed)
    lo = sax_cfg.window_length + ctx.horizon
    pool = np.arange(lo, bound - ctx.horizon)
    origins = rng.choice(pool, size=min(n, len(pool)), replace=False)
    checked = 0
    for o in origins:
        o = int(o)
        a = frozen.forecast(s_small, o, ctx, "", "")
        b = fast_sax_record(s_small, o, ctx, sax_cfg, sym_idx, mid)
        da, db = a.to_dict(), b.to_dict()
        for key in ("sax_state", "sax_analog_count", "sax_median", "sax_P25",
                    "sax_P75", "sax_up_rate", "sax_word", "forecast_direction",
                    "normalized_expected_return", "sax_distance_median"):
            va, vb = da.get(key), db.get(key)
            if isinstance(va, float) and isinstance(vb, float):
                if not (abs(va - vb) < 1e-9):
                    return False
            elif va != vb:
                return False
        checked += 1
    return checked


def mutation_check(s, ctx, sax_cfg, arima_cfg, sym_idx, mid, origin_idx):
    """Mutate bars strictly after origin (well beyond target start); model must not change."""
    close = np.array(s.close, dtype=float)
    mutated = close.copy()
    if origin_idx + ctx.horizon + 3 < len(mutated):
        # mutate bars strictly after origin + horizon (never used in any fit)
        mutated[origin_idx + ctx.horizon: origin_idx + ctx.horizon + 10] = \
            mutated[origin_idx + ctx.horizon: origin_idx + ctx.horizon + 10] * 7.0 + 100.0

    s0 = Series(symbol=s.symbol, timeframe=s.timeframe, timestamp=list(s.timestamp),
                open=list(s.open), high=list(s.high), low=list(s.low), close=list(close))
    arima0 = ArimaModel(arima_cfg)
    rec0_ar = arima0.forecast(s0, origin_idx, ctx, "", "")
    rec0_sx = fast_sax_record(s0, origin_idx, ctx, sax_cfg, sym_idx, mid)

    sm = Series(symbol=s.symbol, timeframe=s.timeframe, timestamp=list(s.timestamp),
                open=list(s.open), high=list(s.high), low=list(s.low), close=list(mutated))
    arima1 = ArimaModel(arima_cfg)
    rec1_ar = arima1.forecast(sm, origin_idx, ctx, "", "")
    sym2, mid2 = precompute_words(mutated, sax_cfg.window_length,
                                  sax_cfg.word_length, sax_cfg.alphabet_size)
    rec1_sx = fast_sax_record(sm, origin_idx, ctx, sax_cfg, sym2, mid2)

    def eq(a, b):
        if a is None and b is None:
            return True
        if a is None or b is None:
            return False
        return abs(float(a) - float(b)) < 1e-9

    return {
        "origin": origin_idx,
        "arima_prediction_unchanged": eq(rec0_ar.normalized_expected_return, rec1_ar.normalized_expected_return),
        "sax_prediction_unchanged": eq(rec0_sx.sax_median, rec1_sx.sax_median),
        "future_bars_mutated": True,
    }


def git_head():
    try:
        out = subprocess.run(["git", "rev-parse", "--short", "HEAD"],
                             capture_output=True, text=True, cwd=ROOT).stdout.strip()
        if not out:
            out = subprocess.run(["git", "rev-parse", "HEAD"], capture_output=True,
                                 text=True, cwd=ROOT).stdout.strip()
        return out
    except Exception:
        return ""


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--code-commit", default="")
    args = ap.parse_args()
    code_commit = args.code_commit or git_head()

    cfg = json.load(open(CONFIG, encoding="utf-8"))
    arima_cfg = ArimaConfig(**cfg["arima"])
    sax_cfg = SaxConfig(**cfg["sax"])
    c_hash = hashlib.sha256(json.dumps(cfg, sort_keys=True).encode()).hexdigest()[:16]
    if c_hash != EXPECTED_CONFIG_HASH:
        print("FATAL: config hash changed", c_hash)
        return 2

    m1 = sorted_unique(load_raw(RAW))
    m5 = aggregate(m1, 5)
    m15 = aggregate(m1, 15)
    raw = {"M1": m1, "M5": m5, "M15": m15}
    series = {"M1": mk_series(m1, "XAUUSDc", "M1"),
              "M5": mk_series(m5, "XAUUSDc", "M5"),
              "M15": mk_series(m15, "XAUUSDc", "M15")}
    closes = {tf: np.array([float(b["close"]) for b in raw[tf]], float) for tf in ("M1", "M5", "M15")}
    close_hash = {tf: hashlib.sha256(closes[tf].astype("<f8").tobytes()).hexdigest() for tf in ("M1", "M5", "M15")}
    dataset_hash = close_hash["M1"]
    if dataset_hash != EXPECTED_DATASET_HASH:
        print("FATAL: dataset hash changed")
        sys.exit(2)
    raw_json_hash = sha256(RAW)

    origin_hashes = {tf: sha256(os.path.join(ORIG_DIR, f"origins_{tf}.csv")) for tf in ("M1", "M5", "M15")}
    proto = json.load(open(os.path.join(ORIG_DIR, "R1_R_PROTOCOL_MANIFEST.json"), encoding="utf-8"))
    for tf in ("M1", "M5", "M15"):
        if origin_hashes[tf] != proto["origin_list_hashes"][tf]:
            print(f"FATAL: {tf} origin hash mismatch")
            sys.exit(2)

    naive = NaiveBaseline(NaiveConfig())
    drift = DriftBaseline(DriftConfig())

    col_list = ["experiment_id", "timeframe", "origin_index", "origin_timestamp",
                "training_start", "training_end", "target_start", "target_end",
                "horizon", "model", "prediction", "actual", "actual_abs", "ATR",
                "expected_return", "forecast_direction", "model_state",
                "arima_state", "sax_state", "hybrid_state",
                "sax_word", "sax_analog_count", "sax_median", "sax_P25", "sax_P75",
                "sax_up_rate", "arima_p", "arima_d", "arima_q",
                "arima_lower", "arima_upper", "config_hash", "dataset_hash",
                "code_commit"]

    row_counts = {}
    pred_hashes = {}
    validate_result = {}
    no_lookahead = {}

    for tf, p in PROTOCOLS.items():
        ctx = ForecastContext(symbol="XAUUSDc", timeframe=tf, horizon=p["H"], atr_period=ATR_PERIOD)
        s = series[tf]
        close = closes[tf]
        n = len(close)
        atr = data_atr(close, ATR_PERIOD)

        sym_idx, mid = precompute_words(close, sax_cfg.window_length,
                                        sax_cfg.word_length, sax_cfg.alphabet_size)
        vc = None
        if p["validate"]:
            vc = validate_fast_sax(s, ctx, sax_cfg, sym_idx, mid)
        validate_result[tf] = vc

        arima_model = ArimaModel(arima_cfg)
        origins = read_origins(tf)
        oset = sorted(set(r["origin_index"] for r in origins))
        if len(oset) != len(origins) or len(origins) != proto["actual_origin_counts"][tf]:
            print("FATAL: origin count/dup mismatch", tf)
            sys.exit(2)

        rows = []
        for r in origins:
            o = r["origin_index"]
            if r["origin_timestamp"] != s.timestamp[o] or o + p["H"] >= n:
                print("FATAL: origin timestamp/bound mismatch", tf, o)
                sys.exit(2)
            if r["target_start"] != o + 1 or r["target_end"] != o + p["H"] or r["training_end"] != o:
                print("FATAL: origin window mismatch", tf, o)
                sys.exit(2)

            recs = {
                "naive": naive.forecast(s, o, ctx, c_hash, dataset_hash),
                "drift": drift.forecast(s, o, ctx, c_hash, dataset_hash),
                "arima": arima_model.forecast(s, o, ctx, c_hash, dataset_hash),
                "sax": fast_sax_record(s, o, ctx, sax_cfg, sym_idx, mid),
            }
            recs["hybrid"] = make_hybrid_record(recs["arima"], recs["sax"])

            actual = float(close[o + p["H"]] - close[o])
            ato = float(atr[o]) if not np.isnan(atr[o]) and atr[o] > 0 else None
            actual_atr = actual / float(atr[o]) if not np.isnan(atr[o]) and atr[o] > 0 else 0.0
            for rec in recs.values():
                rec.actual_forward_return = actual
                rec.actual_forward_return_ATR = actual_atr
                rec.prediction_timestamp = s.timestamp[o]
                rec.actual_outcome_timestamp = s.timestamp[o + p["H"]]
                rec.train_boundary = o

            for m in MODELS:
                rec = recs[m]
                rows.append({
                    "experiment_id": "R1-R", "timeframe": tf,
                    "origin_index": o, "origin_timestamp": r["origin_timestamp"],
                    "training_start": r["training_start"], "training_end": r["training_end"],
                    "target_start": r["target_start"], "target_end": r["target_end"],
                    "horizon": p["H"], "model": m,
                    "prediction": fnum(rec.normalized_expected_return),
                    "actual": fnum(actual_atr), "actual_abs": fnum(actual),
                    "ATR": fnum(atr[o]) if not np.isnan(atr[o]) else "",
                    "expected_return": fnum(rec.expected_return),
                    "forecast_direction": rec.forecast_direction,
                    "model_state": model_state(m, rec),
                    "arima_state": rec.arima_state, "sax_state": rec.sax_state,
                    "hybrid_state": rec.hybrid_state, "sax_word": rec.sax_word,
                    "sax_analog_count": rec.sax_analog_count,
                    "sax_median": fnum(rec.sax_median), "sax_P25": fnum(rec.sax_P25),
                    "sax_P75": fnum(rec.sax_P75), "sax_up_rate": fnum(rec.sax_up_rate),
                    "arima_p": rec.arima_p, "arima_d": rec.arima_d, "arima_q": rec.arima_q,
                    "arima_lower": fnum(rec.arima_lower), "arima_upper": fnum(rec.arima_upper),
                    "config_hash": c_hash, "dataset_hash": dataset_hash, "code_commit": code_commit,
                })

        csv_path = os.path.join(ORIG_DIR, f"predictions_{tf}.csv")
        with open(csv_path, "w", encoding="utf-8", newline="") as fh:
            fh.write(",".join(col_list) + "\n")
            for r in rows:
                fh.write(",".join(str(r[c]) for c in col_list) + "\n")
        row_counts[tf] = len(rows)
        pred_hashes[tf] = sha256(csv_path)

        checks = [mutation_check(s, ctx, sax_cfg, arima_cfg, sym_idx, mid, o)
                  for o in (oset[0], oset[len(oset) // 2], oset[-1])]
        no_lookahead[tf] = checks

    manifest = {
        "experiment": "R1-R", "stage": "execution",
        "execution_commit": code_commit, "protocol_commit": proto["code_commit"],
        "models": MODELS,
        "target": {"definition": "Forward Return / ATR(20)", "atr_period": ATR_PERIOD},
        "dataset": {"close_hashes": close_hash, "M1_raw_json_sha256": raw_json_hash,
                    "n_bars": {"M1": len(m1), "M5": len(m5), "M15": len(m15)}},
        "config_hash": c_hash,
        "origin_list_hashes": origin_hashes,
        "prediction_file_hashes": pred_hashes,
        "row_counts": row_counts,
        "expected_row_count_per_res": {tf: proto["actual_origin_counts"][tf] * len(MODELS)
                                       for tf in ("M1", "M5", "M15")},
        "sax_fast_equivalence_checked": validate_result,
        "no_lookahead_mutation_checks": no_lookahead,
    }
    with open(os.path.join(ORIG_DIR, "R1_R_EXECUTION_MANIFEST.json"), "w",
              encoding="utf-8", newline="") as fh:
        json.dump(manifest, fh, indent=2)

    print("R1-R execution complete. code_commit =", code_commit)
    for tf in ("M1", "M5", "M15"):
        print(f"  {tf}: rows={row_counts[tf]} expected={proto['actual_origin_counts'][tf]*len(MODELS)} "
              f"pred_hash={pred_hashes[tf][:12]} sax_eq={validate_result[tf]}")


if __name__ == "__main__":
    main()