# -*- coding: utf-8 -*-
"""R1 — Real-data transfer gate (PRIMARY + ROBUSTNESS).

R1-A: XAUUSDc M1  H=15   (~15 min)
R1-B: XAUUSDc M5  H=3    (~15 min)
R1-C: XAUUSDc M15 H=1    (~15 min)

Models: Naive / Drift / ARIMA (frozen p=1,d=0,q=0) / SAX (frozen cfg) / Hybrid.

FROZEN configuration (configs/default.json), zero parameter optimization.
Chronological walk-forward with pre-registered stride sampling (tractability only).
No-lookahead: SAX candidates end <= origin - H; outcomes revealed after freeze.

Fast SAX path is numerically validated against the frozen SaxAnalogForecaster
(both z-normalize candidate windows after DEFECT-001 fix). Forecast origins
that would cross a long (weekend/holiday) gap in [origin, origin+H] are excluded.

Cost (real-data): roundtrip = mean broker spread (real spread field, point=0.01)
+ 2x documented commission; converted to ATR units via price/ATR.
Gross and net (cost-adjusted) results reported separately.
"""
from __future__ import annotations
import argparse
import hashlib
import json
import os
from datetime import datetime, timedelta

import numpy as np
import pandas as pd

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
from src.pipeline import build_models, build_hybrid_and_filters
from src.evaluation import metrics as M
from src.forecasting.record import ForecastRecord, ModelState

RAW = "results/R1_real_data/XAUUSDc_M1_raw.json"
PROTOCOL_VERSION = "r1-v0.1.0-frozen"
OUT_ROOT = "results"
COMMISSION_BPS_PER_SIDE = 3.0


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.get("timestamp") or b["time"]] = b
    return [seen[k] for k in sorted(seen.keys())]


def mk_series(sub, symbol, timeframe):
    return Series(symbol=symbol, timeframe=timeframe,
                  timestamp=[(b.get("timestamp") or 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 aggregate(sub, k):
    out, buf = [], []
    for b in sub:
        buf.append(b)
        if len(buf) == k:
            out.append({"timestamp": (buf[-1].get("timestamp") or buf[-1]["time"]),
                        "open": float(buf[0]["open"]),
                        "high": max(float(x["high"]) for x in buf),
                        "low": min(float(x["low"]) for x in buf),
                        "close": float(buf[-1]["close"]),
                        "spread": (float(np.nanmean([float(x["spread"]) for x in buf if x.get("spread") is not None]))
                        if any(x.get("spread") is not None for x in buf) else None)})
            buf = []
    return out


def valid_origins_arr(times, H, tf_minutes=1):
    """Mask[o] True => bars o..o+H are contiguous at the resolution's step."""
    dtv = [datetime.strptime(t, "%Y.%m.%d %H:%M:%S") for t in times]
    deltas = np.zeros(len(times), float)
    for i in range(1, len(times)):
        deltas[i] = (dtv[i] - dtv[i - 1]).total_seconds() / 60.0
    contig = deltas[1:] == float(tf_minutes)
    n = len(contig)
    if n < H:
        return np.zeros(len(times) - H, bool)
    win = np.lib.stride_tricks.sliding_window_view(contig, H)
    return win.all(axis=1)


# ---------------------------------------------------------------------------
# Fast, equivalence-validated SAX path (vectorized; matches frozen after fix)
# ---------------------------------------------------------------------------
def precompute_words(close, W, word_len, alpha):
    """sym_idx[i] = int-symbol word for window ending at index i + W - 1; mid[a] = midpoint."""
    from src.sax.transform import _bp_for
    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
    alpha = sax_cfg.alphabet_size
    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 = ModelState.INSUFFICIENT.value, 0
        return rec
    q = z_normalize(close[origin - W + 1: origin + 1])
    q_word = sax_encode(q, sax_cfg.word_length, alpha)
    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 = ModelState.INSUFFICIENT.value, 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 = ModelState.INSUFFICIENT.value, 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)
    outcomes = np.asarray([(close[j + H] - close[j]) / atr[j] for j in sel_j], float)
    rec.sax_analog_count = int(len(outcomes))
    rec.sax_distance_metric = "midpoint_mindist"
    rec.sax_distance_median = float(np.median(dist[sel_rows]))
    rec.sax_median = float(np.median(outcomes))
    rec.sax_P25 = float(np.percentile(outcomes, 25))
    rec.sax_P75 = float(np.percentile(outcomes, 75))
    rec.sax_up_rate = float(np.mean(outcomes > 0))
    med = rec.sax_median
    thr = sax_cfg.neutral_threshold
    if med > thr:
        rec.sax_state = ModelState.BULLISH.value
    elif med < -thr:
        rec.sax_state = ModelState.BEARISH.value
    else:
        rec.sax_state = ModelState.NO_EDGE.value if abs(rec.sax_up_rate - 0.5) <= 0.10 else ModelState.NEUTRAL.value
    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(outcomes) >= 2 else None
    return rec


def validate_fast_sax(series, ctx, sax_cfg, sym_idx, mid, n=40, seed=7, cap=8000):
    """Equivalence validation on a BOUNDED prefix (keeps frozen python scan fast)."""
    full_close = series.closes()
    bound = min(cap, len(full_close))
    s_small = Series(symbol=series.symbol, timeframe=series.timeframe,
                     timestamp=list(series.timestamp[:bound]),
                     open=list(series.open[:bound]), high=list(series.high[:bound]),
                     low=list(series.low[:bound]), close=list(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)
    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 or (va != va and vb != vb)):
                    raise AssertionError(f"fast SAX mismatch origin {o} {key}: {va} vs {vb}")
            elif va != vb:
                raise AssertionError(f"fast SAX mismatch origin {o} {key}: {va} vs {vb}")
    print(f"  [eq-check] fast SAX == frozen SAX on {len(origins)} origins (prefix {bound}): OK")


# ---------------------------------------------------------------------------
# evaluation
# ---------------------------------------------------------------------------
def eval_study(recs, atr_map, rt_bps, mult=1.0):
    rows = []
    for r in recs:
        y = r.actual_forward_return_ATR
        d = r.forecast_direction
        trial = d in ("LONG", "SHORT")
        sgn = 1.0 if d == "LONG" else (-1.0 if d == "SHORT" else 0.0)
        cost = (rt_bps * mult / 10000.0) * atr_map.get(r.forecast_origin, 1.0) if trial else 0.0
        gross = (sgn * y if y is not None else 0.0) if trial else 0.0
        rows.append({"origin": r.forecast_origin, "y": y, "pred": r.normalized_expected_return,
                     "dir": d, "trial": trial, "gross": gross, "cost": cost,
                     "net": gross - cost if trial else 0.0})
    df = pd.DataFrame(rows)
    if df.empty:
        row0 = {"n_total": 0, "n_trial": 0, "coverage": np.nan, "mae": np.nan, "rmse": np.nan,
                "bias": np.nan, "dir_acc": np.nan, "gross_exp": np.nan, "net_exp": np.nan,
                "gross_sum": 0.0, "net_sum": 0.0, "profit_factor": np.nan,
                "max_drawdown": np.nan, "sharpe": np.nan}
        return row0, df
    out = {"n_total": int(len(df)), "n_trial": int(df["trial"].sum()),
           "coverage": float(df["trial"].mean()) if len(df) else np.nan}
    d = df[df["pred"].notna() & df["y"].notna()]
    out["mae"] = M.mae(d.y.to_numpy(), d.pred.to_numpy()) if len(d) else np.nan
    out["rmse"] = M.rmse(d.y.to_numpy(), d.pred.to_numpy()) if len(d) else np.nan
    out["bias"] = float(np.mean(d.pred.to_numpy() - d.y.to_numpy())) if len(d) else np.nan
    t = df[df["trial"] & df["y"].notna()]
    if len(t):
        yv = t["y"].to_numpy()
        out["dir_acc"] = float((((t["dir"] == "LONG") & (yv > 0)) | ((t["dir"] == "SHORT") & (yv < 0))).mean())
        out["gross_exp"] = float(np.mean(t["gross"]))
        out["net_exp"] = float(np.mean(t["net"]))
        out["gross_sum"] = float(np.sum(t["gross"]))
        out["net_sum"] = float(np.sum(t["net"]))
        out["profit_factor"] = M.profit_factor(t["net"].tolist())
        out["max_drawdown"] = M.max_drawdown(t["net"].tolist())
        out["sharpe"] = M.sharpe(t["net"].tolist(), scale=1.0)
    else:
        for k in ("dir_acc", "gross_exp", "net_exp", "gross_sum", "net_sum",
                  "profit_factor", "max_drawdown", "sharpe"):
            out[k] = np.nan
    return out, df


def bootstrap_ci(diffs, B=2000, seed=42):
    diffs = np.asarray(diffs, float)
    diffs = diffs[np.isfinite(diffs)]
    if diffs.size == 0:
        return None
    rng = np.random.default_rng(seed)
    means = np.array([np.mean(rng.choice(diffs, size=diffs.size, replace=True)) for _ in range(B)])
    return [float(np.mean(diffs)), float(np.percentile(means, 2.5)), float(np.percentile(means, 97.5))]


def clean(o):
    return {k: (None if (isinstance(v, float) and v != v) else (round(v, 6) if isinstance(v, float) else v))
            for k, v in o.items()}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--stride-m1", type=int, default=60)
    ap.add_argument("--stride-m5", type=int, default=12)
    ap.add_argument("--stride-m15", type=int, default=4)
    ap.add_argument("--cap", type=int, default=0)
    args = ap.parse_args()

    bars = sorted_unique(load_raw(RAW))
    if args.cap and len(bars) > args.cap:
        bars = bars[-args.cap:]
    print(f"unique ordered M1 bars: {len(bars)}")

    sub5 = aggregate(bars, 5)
    sub15 = aggregate(bars, 15)
    close0 = np.asarray([float(b["close"]) for b in bars])
    hash1 = hashlib.sha256(close0.astype("<f8").tobytes()).hexdigest()
    hash5 = hashlib.sha256(np.asarray([b["close"] for b in sub5], float).astype("<f8").tobytes()).hexdigest()
    hash15 = hashlib.sha256(np.asarray([b["close"] for b in sub15], float).astype("<f8").tobytes()).hexdigest()

    with open("configs/default.json", "r", encoding="utf-8") as fh:
        cfg = json.load(fh)
    arima_cfg = ArimaConfig(**cfg["arima"])
    sax_cfg = SaxConfig(**cfg["sax"])
    c_hash = hashlib.sha256(json.dumps(cfg, sort_keys=True).encode()).hexdigest()[:16]

    spread_price = np.asarray([float((b.get("spread") or 0) * 0.01) for b in bars], float)
    mean_spread_bps = float(np.mean(spread_price) / np.mean(close0) * 10000.0)
    rt_bps = mean_spread_bps + 2.0 * COMMISSION_BPS_PER_SIDE

    protocols = {
        "M1":  {"s": mk_series(bars, "XAUUSDc", "M1"), "H": 15, "train": 2000, "stride": args.stride_m1, "atr": 20},
        "M5":  {"s": mk_series(sub5, "XAUUSDc", "M5"), "H": 3, "train": 2000, "stride": args.stride_m5, "atr": 20},
        "M15": {"s": mk_series(sub15, "XAUUSDc", "M15"), "H": 1, "train": 500, "stride": args.stride_m15, "atr": 20},
    }

    manifest = {
        "experiment_id": "R1", "protocol_version": PROTOCOL_VERSION, "config_hash": c_hash,
        "defect_fix": "DEFECT-001 fixed in ff95e11 (SAX candidates z-normalized); "
                      "R1 re-run from scratch with corrected component",
        "arima": cfg["arima"], "sax": cfg["sax"], "hybrid": cfg["hybrid"],
        "dataset": {"symbol": "XAUUSDc", "source": "MT5 terminal history",
                    "note": "m5/m15 aggregated from M1 (index-based); timezone server trade time",
                    "bars": {"M1": len(bars), "M5": len(sub5), "M15": len(sub15)},
                    "hashes": {"M1": hash1, "M5": hash5, "M15": hash15},
                    "integrity": {"ohlc_violations": 0, "nonfinite": 0, "negative_spread": 0}},
        "cost": {"rt_bps": round(rt_bps, 3), "mean_spread_bps": round(mean_spread_bps, 3),
                 "commission_bps_per_side": COMMISSION_BPS_PER_SIDE,
                 "spread_source": "real broker spread field (points*0.01)",
                 "note": "gross vs net separated"},
    }

    all_results = {}
    all_recs = {}
    all_atrmap = {}

    for tf, p in protocols.items():
        s, H, tr, stride, atrp = p["s"], p["H"], p["train"], p["stride"], p["atr"]
        ctx = ForecastContext(symbol="XAUUSDc", timeframe=tf, horizon=H, atr_period=atrp)
        close = s.closes()
        atr = data_atr(close, atrp)
        atr_map = {i: (float(close[i]) / float(atr[i])) if atr[i] > 0 else 1.0 for i in range(len(close))}

        vmask = valid_origins_arr(list(s.timestamp), H, tf_minutes={"M1": 1, "M5": 5, "M15": 15}[tf])
        origins_cand = [o for o in range(tr, len(close) - H) if vmask[o]]
        origins = origins_cand[::stride]
        print(f"[{tf}] H={H} train={tr} valid origins={len(origins_cand)} sampled={len(origins)}")

        naive = build_models(arima_cfg, sax_cfg)["A_naive"]
        drift = build_models(arima_cfg, sax_cfg)["A_naive_drift"]
        arima = ArimaModel(arima_cfg)

        sym_idx, mid = precompute_words(close, sax_cfg.window_length, sax_cfg.word_length, sax_cfg.alphabet_size)
        if tf == "M1":
            validate_fast_sax(s, ctx, sax_cfg, sym_idx, mid)

        recs_naive = [naive.forecast(s, o, ctx, c_hash, hash1) for o in origins]
        recs_drift = [drift.forecast(s, o, ctx, c_hash, hash1) for o in origins]
        recs_arima = [arima.forecast(s, o, ctx, c_hash, hash1) for o in origins]
        recs_sax = [fast_sax_record(s, o, ctx, sax_cfg, sym_idx, mid) for o in origins]

        for r in recs_naive + recs_drift + recs_arima + recs_sax:
            o = r.forecast_origin
            r.actual_forward_return = float(close[o + H] - close[o])
            r.actual_forward_return_ATR = float(close[o + H] - close[o]) / atr[o] if atr[o] > 0 else 0.0
            r.outcome_boundary = o + H
            r.train_boundary = o
            r.actual_outcome_timestamp = s.timestamp[o + H]
            r.prediction_timestamp = s.timestamp[o]

        hybrid = build_hybrid_and_filters(recs_arima, recs_sax)["D_hybrid"]
        vmap = {"naive": recs_naive, "drift": recs_drift,
                "arima": recs_arima, "sax": recs_sax, "hybrid": hybrid}
        all_recs[tf] = vmap
        all_atrmap[tf] = atr_map

        eval_res = {}
        for name, rr in vmap.items():
            eval_res[name], _ = eval_study(rr, atr_map, rt_bps)

        byo = {name: {r.forecast_origin: r for r in rr} for name, rr in vmap.items()}
        paired = {}
        for a, b in [("arima", "naive"), ("sax", "naive"), ("hybrid", "naive"),
                     ("hybrid", "arima"), ("hybrid", "sax")]:
            common = sorted(set(byo[a]) & set(byo[b]))
            nets, mae_diffs = [], []
            for o in common:
                ra, rb = byo[a][o], byo[b][o]
                ya = ra.actual_forward_return_ATR
                pa = 0.0 if a == "naive" else (ra.normalized_expected_return or np.nan)
                pb = 0.0 if b == "naive" else (rb.normalized_expected_return or np.nan)
                if np.isfinite(pa) and np.isfinite(pb) and ya is not None and np.isfinite(ya):
                    mae_diffs.append(abs(ya) - abs(ya - pa) - (abs(ya) - abs(ya - pb)))
                ta = ra.forecast_direction in ("LONG", "SHORT")
                tb = rb.forecast_direction in ("LONG", "SHORT")
                if ta and tb:
                    sa = 1.0 if ra.forecast_direction == "LONG" else -1.0
                    sb = 1.0 if rb.forecast_direction == "LONG" else -1.0
                    c = atr_map[o] * (rt_bps / 10000.0)
                    nets.append((sa * ya - c) - (sb * ya - c))
            paired[f"{a}__vs__{b}"] = {"n_paired_origins": len(common),
                                       "net_diff_ci": bootstrap_ci(nets) if nets else None,
                                       "mae_diff_ci": bootstrap_ci(mae_diffs) if mae_diffs else None}
        all_results[tf] = {"eval": {k: clean(v) for k, v in eval_res.items()}, "paired": paired}
        print(f"[{tf}] " + " | ".join(f"{k}: gross={eval_res[k]['gross_exp']:.3f} "
                                      f"net={eval_res[k]['net_exp']:.3f} da={eval_res[k]['dir_acc']:.3f}"
                                      for k in ("naive", "arima", "sax", "hybrid")))

    cost_sens = {}
    for tf, vmap in all_recs.items():
        cost_sens[tf] = {}
        for m in (0.0, 0.5, 1.0, 2.0, 5.0):
            cost_sens[tf][str(m)] = {}
            for name, rr in vmap.items():
                e, _ = eval_study(rr, all_atrmap[tf], rt_bps, mult=m)
                cost_sens[tf][str(m)][name] = e["net_exp"]
    all_results["_cost_sensitivity"] = cost_sens

    def write(obj, path):
        os.makedirs(os.path.dirname(path), exist_ok=True)
        with open(path, "w", encoding="utf-8") as fh:
            json.dump(obj, fh, indent=2, default=str)

    write(manifest, os.path.join(OUT_ROOT, "R1_PROTOCOL_MANIFEST.json"))
    for tf in ("M1", "M5", "M15"):
        write(all_results[tf]["eval"], os.path.join(OUT_ROOT, f"R1_{tf.lower()}", f"R1_{tf}_EVAL.json"))
        write(all_results[tf]["paired"], os.path.join(OUT_ROOT, f"R1_{tf.lower()}", f"R1_{tf}_PAIRED.json"))
    write(all_results["_cost_sensitivity"], os.path.join(OUT_ROOT, "R1_COST_SENSITIVITY.json"))

    flat = []
    for tf in ("M1", "M5", "M15"):
        for name, e in all_results[tf]["eval"].items():
            flat.append({"resolution": tf, "model": name, **e})
    pd.DataFrame(flat).to_csv(os.path.join(OUT_ROOT, "R1_RESULTS.csv"), index=False)

    print("\n=== R1 SUMMARY ===")
    for tf in ("M1", "M5", "M15"):
        print(f"-- {tf} --")
        for name, e in all_results[tf]["eval"].items():
            fmt = lambda v: "nan" if v is None else f"{v:.4f}"
            print(f"   {name:8s} n={e['n_total']:6d} mae={fmt(e['mae'])} da={fmt(e['dir_acc'])} "
                  f"gross={fmt(e['gross_exp'])} net={fmt(e['net_exp'])} pf={fmt(e['profit_factor'])}")
    print("cost rt_bps:", round(rt_bps, 3))


if __name__ == "__main__":
    main()