"""
⚡ ULTRON Systems // Deterministic Tier-1 Scientific Audit & Verification Engine
100% Reproducible, Zero-Hallucination Empirical Pipeline for NVIDIA A100-SXM4-80GB.

Computes:
1. Exact Layer-by-Layer Parameter Breakdown of the 1,189,135,740 Parameter Conformer.
2. Exact Frobenius Weight Drift & Cosine Similarity across Epoch 1, Epoch 2, Epoch 3, WF-2024, WF-2025.
3. Multi-Epoch Progression (Epoch 1 vs Epoch 2 vs Epoch 3) & Walk-Forward OOS Benchmark.
4. Four Deterministic Quantitative Baselines (Always-Long Drift, 1m Momentum, 5m Reversal, EMA-12/26).
5. Full 7-Horizon Mechanistic Inception Kernel Ablation (K=3,7,15,31,63,127 + Micro/Meso/Macro Groups)
   measuring Directional Accuracy, Balanced Accuracy, ROC-AUC, Rank IC, and Continuous Logit L2 Sensitivity.
6. Strictly Causal, Chronological No-Look-Ahead Execution Backtest across Horizons & Confidence Hurdles.
"""
import os
import sys
import gc
import json
import time
import math
import numpy as np
import scipy.stats as stats
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from huggingface_hub import hf_hub_download, HfApi

from config import DatasetConfig
from config_1b import Model1BConfig, Training1BConfig
from dataset import StreamingShardManager
from model_1b import FinancialConformer1B

HORIZON_NAMES = ["1m", "3m", "5m", "15m", "30m", "60m", "1d"]
HORIZON_MINUTES = [1, 3, 5, 15, 30, 60, 390]

# Fix all random seeds for 100% deterministic reproducibility
torch.manual_seed(42)
np.random.seed(42)

def fast_roc_auc(y_true_bin: np.ndarray, y_score: np.ndarray) -> float:
    """Deterministic Mann-Whitney U statistic based ROC-AUC calculation."""
    pos = y_score[y_true_bin == 1]
    neg = y_score[y_true_bin == 0]
    if len(pos) == 0 or len(neg) == 0:
        return 0.5
    ranks = stats.rankdata(np.concatenate([pos, neg]))
    pos_ranks = ranks[:len(pos)]
    auc = (np.sum(pos_ranks) - len(pos) * (len(pos) + 1) / 2.0) / (len(pos) * len(neg))
    return float(auc)

def compute_rich_horizon_metrics(preds: np.ndarray, logits: np.ndarray, targets: np.ndarray) -> dict:
    """
    Computes rigorous statistical metrics across all 7 horizons:
    - Raw Directional Accuracy (%)
    - Non-Zero Directional Accuracy (%) (strictly excluding y_true == 0)
    - Balanced Accuracy (%) = 0.5 * (TPR + TNR)
    - Matthews Correlation Coefficient (MCC)
    - ROC-AUC (Rank-based)
    - Spearman Rank IC & two-sided p-value
    - Pearson r
    - Huber Loss (delta=0.01) & BCE Loss
    - Binomial z-statistic & p-value vs 50% null
    - Class balance (positive %, negative %, flat %)
    """
    num_h = targets.shape[1]
    per_h = {}
    raw_accs, nz_accs, bal_accs, aucs, ics = [], [], [], [], []

    for h in range(num_h):
        yt = targets[:, h]
        yp = preds[:, h]
        yl = logits[:, h]

        valid = ~(np.isnan(yt) | np.isnan(yp) | np.isnan(yl) | np.isinf(yt) | np.isinf(yp) | np.isinf(yl))
        yt, yp, yl = yt[valid], yp[valid], yl[valid]
        n_total = len(yt)
        if n_total < 10:
            continue

        # Class distribution
        pos_mask = yt > 0.0
        neg_mask = yt < 0.0
        flat_mask = yt == 0.0
        pos_pct = float(np.mean(pos_mask) * 100.0)
        neg_pct = float(np.mean(neg_mask) * 100.0)
        flat_pct = float(np.mean(flat_mask) * 100.0)

        pred_pos = yl > 0.0

        # 1. Raw Directional Accuracy (standard (yl > 0) == (yt > 0))
        raw_correct = (pred_pos == pos_mask).astype(np.float64)
        raw_acc = float(np.mean(raw_correct) * 100.0)
        raw_accs.append(raw_acc)

        # Binomial z-stat vs 50%
        k_succ = int(np.sum(raw_correct))
        z_stat = (k_succ - 0.5 * n_total) / math.sqrt(0.25 * n_total)
        p_val_50 = float(2.0 * (1.0 - stats.norm.cdf(abs(z_stat))))

        # 2. Strict Non-Zero Directional Accuracy (excluding flat bars yt == 0)
        nz_mask = ~flat_mask
        if np.sum(nz_mask) > 10:
            yt_nz = yt[nz_mask] > 0.0
            yl_nz = yl[nz_mask] > 0.0
            nz_acc = float(np.mean(yl_nz == yt_nz) * 100.0)
        else:
            nz_acc = raw_acc
        nz_accs.append(nz_acc)

        # 3. Confusion Matrix, Balanced Accuracy, MCC
        tp = float(np.sum(pred_pos & pos_mask))
        tn = float(np.sum((~pred_pos) & (~pos_mask)))
        fp = float(np.sum(pred_pos & (~pos_mask)))
        fn = float(np.sum((~pred_pos) & pos_mask))

        tpr = tp / max(1.0, tp + fn)
        tnr = tn / max(1.0, tn + fp)
        bal_acc = float(0.5 * (tpr + tnr) * 100.0)
        bal_accs.append(bal_acc)

        denom = math.sqrt(max(1e-12, (tp + fp) * (tp + fn) * (tn + fp) * (tn + fn)))
        mcc = float((tp * tn - fp * fn) / denom) if denom > 1e-6 else 0.0

        # 4. ROC-AUC
        auc = fast_roc_auc(pos_mask.astype(np.int32), yl)
        aucs.append(auc)

        # 5. Spearman Rank IC & Pearson r
        if np.std(yp) > 1e-12 and np.std(yt) > 1e-12:
            sp_rho, sp_pval = stats.spearmanr(yp, yt)
            pr_r, _ = stats.pearsonr(yp, yt)
            sp_rho = 0.0 if np.isnan(sp_rho) else float(sp_rho)
            sp_pval = 1.0 if np.isnan(sp_pval) else float(sp_pval)
            pr_r = 0.0 if np.isnan(pr_r) else float(pr_r)
        else:
            sp_rho, sp_pval, pr_r = 0.0, 1.0, 0.0
        ics.append(sp_rho)

        # 6. Huber & BCE Loss
        err = np.abs(yp - yt)
        delta = 0.01
        huber = float(np.mean(np.where(err <= delta, 0.5 * (err ** 2), delta * (err - 0.5 * delta))))
        # Numerically stable BCE with logits
        y_bin = pos_mask.astype(np.float64)
        yl_64 = np.clip(yl.astype(np.float64), -30.0, 30.0)
        bce = float(np.mean(np.maximum(yl_64, 0) - yl_64 * y_bin + np.log1p(np.exp(-np.abs(yl_64)))))

        # Mean absolute return in bps for this horizon
        mean_abs_ret_bps = float(np.mean(np.abs(yt)) * 10000.0)

        h_name = HORIZON_NAMES[h]
        per_h[h_name] = {
            "n_samples": int(n_total),
            "pos_pct": round(pos_pct, 2),
            "neg_pct": round(neg_pct, 2),
            "flat_pct": round(flat_pct, 2),
            "mean_abs_move_bps": round(mean_abs_ret_bps, 2),
            "raw_dir_acc": round(raw_acc, 2),
            "nonzero_dir_acc": round(nz_acc, 2),
            "balanced_acc": round(bal_acc, 2),
            "mcc": round(mcc, 4),
            "roc_auc": round(auc, 4),
            "spearman_ic": round(sp_rho, 4),
            "spearman_pval": float(f"{sp_pval:.4e}"),
            "pearson_r": round(pr_r, 4),
            "huber_loss": float(f"{huber:.6e}"),
            "bce_loss": round(bce, 4),
            "z_stat_vs_50": round(z_stat, 2),
            "p_val_vs_50": float(f"{p_val_50:.4e}")
        }

    return {
        "overall_raw_dir_acc": round(float(np.mean(raw_accs)), 2),
        "overall_nonzero_dir_acc": round(float(np.mean(nz_accs)), 2),
        "overall_balanced_acc": round(float(np.mean(bal_accs)), 2),
        "overall_roc_auc": round(float(np.mean(aucs)), 4),
        "overall_mean_ic": round(float(np.mean(ics)), 4),
        "per_horizon": per_h
    }


def run_causal_chronological_backtest(preds: np.ndarray, logits: np.ndarray, targets: np.ndarray) -> dict:
    """
    Strictly Causal, Chronological No-Look-Ahead Execution Backtest.
    Evaluates:
    1. Horizon sweep (1m, 5m, 15m, 60m, 1d) at confidence hurdles tau in {0.0, 0.5, 1.0, 1.5, 2.0} sigma.
       Rolling historical mean/std of signal is computed over past W=256 windows (strictly causal).
    2. Slippage grid: 0.0, 1.0, 2.0, 5.0, 10.0 bps round-trip cost.
    3. Breakeven fee capacitance (in bps) = Gross Mean Return per Trade / 2.
    """
    results = {}
    warmup = 256
    # Annualization factor: non-overlapping intervals per year (252 trading days * 390 mins)
    mins_per_year = 252.0 * 390.0

    for h_idx in [0, 2, 3, 5, 6]:  # 1m, 5m, 15m, 60m, 1d
        h_name = HORIZON_NAMES[h_idx]
        h_min = HORIZON_MINUTES[h_idx]
        periods_per_year = mins_per_year / float(h_min)

        yt = targets[:, h_idx]
        yp = preds[:, h_idx]
        yl = logits[:, h_idx]
        valid = ~(np.isnan(yt) | np.isnan(yp) | np.isnan(yl))
        yt, yp, yl = yt[valid], yp[valid], yl[valid]

        if len(yt) <= warmup + 50:
            continue

        # Composite causal signal: standardized logit + standardized predicted return
        # Compute strictly expanding/rolling historical z-score (no look-ahead!)
        n = len(yt)
        z_sig = np.zeros(n, dtype=np.float64)
        # Vectorized causal rolling window (past 256 samples)
        csum_p = np.cumsum(yp)
        csum_p2 = np.cumsum(yp ** 2)
        csum_l = np.cumsum(yl)
        csum_l2 = np.cumsum(yl ** 2)

        idx = np.arange(warmup, n)
        cnt = float(warmup)
        mean_p = (csum_p[idx - 1] - np.where(idx - 1 - warmup >= 0, csum_p[idx - 1 - warmup], 0.0)) / cnt
        var_p = np.maximum(1e-12, (csum_p2[idx - 1] - np.where(idx - 1 - warmup >= 0, csum_p2[idx - 1 - warmup], 0.0)) / cnt - mean_p ** 2)
        zp = (yp[idx] - mean_p) / np.sqrt(var_p)

        mean_l = (csum_l[idx - 1] - np.where(idx - 1 - warmup >= 0, csum_l[idx - 1 - warmup], 0.0)) / cnt
        var_l = np.maximum(1e-12, (csum_l2[idx - 1] - np.where(idx - 1 - warmup >= 0, csum_l2[idx - 1 - warmup], 0.0)) / cnt - mean_l ** 2)
        zl = (yl[idx] - mean_l) / np.sqrt(var_l)

        z_sig[idx] = 0.5 * (zp + zl)
        yt_eval = yt[idx]
        z_eval = z_sig[idx]

        results[h_name] = {}
        for tau in [0.0, 0.5, 1.0, 1.5, 2.0]:
            active_mask = np.abs(z_eval) > tau if tau > 0.0 else np.ones_like(z_eval, dtype=bool)
            n_trades = int(np.sum(active_mask))
            if n_trades < 20:
                continue

            positions = np.sign(z_eval[active_mask])
            # If tau == 0.0, use raw sign of logit
            if tau == 0.0:
                positions = np.where(yl[idx] > 0.0, 1.0, -1.0)

            gross_trade_rets = positions * yt_eval[active_mask]
            gross_mean_bps = float(np.mean(gross_trade_rets) * 10000.0)
            breakeven_one_way_bps = round(gross_mean_bps / 2.0, 2)
            turnover_Saved_pct = round((1.0 - float(n_trades) / float(len(yt_eval))) * 100.0, 2)

            slippage_dict = {}
            for bps in [0.0, 1.0, 2.0, 5.0, 10.0]:
                rt_cost = (bps / 10000.0) * 2.0  # Round-trip cost
                net_rets = gross_trade_rets - rt_cost
                mean_r = float(np.mean(net_rets))
                std_r = float(np.std(net_rets)) + 1e-12

                # Scale annualization by trade participation rate
                participation = float(n_trades) / float(len(yt_eval))
                eff_periods = max(1.0, periods_per_year * participation)
                sharpe = (mean_r / std_r) * math.sqrt(min(eff_periods, 252.0 * 24.0))

                neg_r = net_rets[net_rets < 0.0]
                down_std = float(np.std(neg_r)) + 1e-12 if len(neg_r) > 1 else std_r
                sortino = (mean_r / down_std) * math.sqrt(min(eff_periods, 252.0 * 24.0))

                # Strictly chronological cumulative return & max drawdown
                cum_curve = np.cumsum(net_rets)
                peak = np.maximum.accumulate(np.insert(cum_curve, 0, 0.0))[1:]
                dd = peak - cum_curve
                max_dd = float(np.max(dd)) if len(dd) > 0 else 0.0
                win_rate = float(np.mean(net_rets > 0.0) * 100.0)

                slippage_dict[f"{bps}bps"] = {
                    "sharpe": round(sharpe, 2),
                    "sortino": round(sortino, 2),
                    "win_rate_pct": round(win_rate, 2),
                    "mean_net_ret_bps": round(mean_r * 10000.0, 2),
                    "chronological_max_dd": round(max_dd, 4)
                }

            results[h_name][f"tau_{tau}sigma"] = {
                "n_trades": n_trades,
                "total_bars": int(len(yt_eval)),
                "turnover_reduction_pct": turnover_Saved_pct,
                "gross_mean_ret_bps": round(gross_mean_bps, 2),
                "breakeven_fee_bps": breakeven_one_way_bps,
                "slippage_grid": slippage_dict
            }

    return results


def audit_architecture_parameters(model: nn.Module) -> dict:
    """Computes exact parameter counts per module category."""
    groups = {
        "revin": 0,
        "inception_tokenizer_branches": 0,
        "inception_tokenizer_proj_norm": 0,
        "conformer_gqa_attention": 0,
        "conformer_depthwise_conv": 0,
        "conformer_swiglu_ffn": 0,
        "conformer_rmsnorms": 0,
        "forecast_heads_and_final_norm": 0
    }
    for name, p in model.named_parameters():
        n = p.numel()
        if name.startswith("revin"):
            groups["revin"] += n
        elif name.startswith("tokenizer.branches"):
            groups["inception_tokenizer_branches"] += n
        elif name.startswith("tokenizer"):
            groups["inception_tokenizer_proj_norm"] += n
        elif ".attn." in name:
            groups["conformer_gqa_attention"] += n
        elif ".conv_module." in name:
            groups["conformer_depthwise_conv"] += n
        elif ".ffn." in name:
            groups["conformer_swiglu_ffn"] += n
        elif "layers." in name and "norm" in name:
            groups["conformer_rmsnorms"] += n
        else:
            groups["forecast_heads_and_final_norm"] += n

    groups["total_exact_parameters"] = sum(groups.values())
    return groups


def compute_weight_drift_matrix(ckpt_paths: dict, device: str = "cpu") -> dict:
    """
    Computes exact relative Frobenius norm drift ||W_B - W_A||_F / ||W_A||_F * 100%
    and cosine similarity between checkpoint pairs across module categories.
    """
    print("\n🔬 [Weight Drift Audit] Loading state dicts for exact Frobenius norm analysis...")
    states = {}
    for k, path in ckpt_paths.items():
        if os.path.exists(path):
            raw = torch.load(path, map_location="cpu", weights_only=False)
            states[k] = raw.get("model_state", raw)
            print(f"  • Loaded '{k}' from {path} ({len(states[k])} tensors)")

    pairs = [
        ("epoch1", "epoch2"),
        ("epoch2", "epoch3"),
        ("epoch1", "epoch3"),
        ("epoch3", "wf2024"),
        ("wf2024", "wf2025"),
        ("epoch3", "wf2025")
    ]

    def _categorize(param_name: str) -> str:
        if param_name.startswith("tokenizer"):
            return "inception_tokenizer"
        elif param_name.startswith("layers."):
            layer_idx = int(param_name.split(".")[1])
            if layer_idx <= 6:
                return "conformer_early_L0_L6"
            elif layer_idx <= 13:
                return "conformer_mid_L7_L13"
            else:
                return "conformer_deep_L14_L19"
        elif "return_head" in param_name:
            return "return_regression_head"
        elif "dir_head" in param_name:
            return "directional_class_head"
        else:
            return "norms_and_aux"

    drift_report = {}
    for a, b in pairs:
        if a not in states or b not in states:
            continue
        sa, sb = states[a], states[b]
        cat_diff_sq = {}
        cat_base_sq = {}
        cat_cos = {}

        total_diff_sq = 0.0
        total_base_sq = 0.0

        for k in sa.keys():
            if k not in sb:
                continue
            ta = sa[k].float()
            tb = sb[k].float()
            if ta.numel() == 0:
                continue

            diff_sq = float(torch.sum((tb - ta) ** 2).item())
            base_sq = float(torch.sum(ta ** 2).item())
            cos_sim = float(F.cosine_similarity(ta.flatten().unsqueeze(0), tb.flatten().unsqueeze(0)).item())

            cat = _categorize(k)
            cat_diff_sq[cat] = cat_diff_sq.get(cat, 0.0) + diff_sq
            cat_base_sq[cat] = cat_base_sq.get(cat, 0.0) + base_sq
            cat_cos.setdefault(cat, []).append(cos_sim)

            total_diff_sq += diff_sq
            total_base_sq += base_sq

        pair_key = f"{a}_to_{b}"
        drift_report[pair_key] = {
            "global_rel_frobenius_pct": round(math.sqrt(total_diff_sq / max(1e-12, total_base_sq)) * 100.0, 4),
            "by_module": {}
        }
        for cat in sorted(cat_diff_sq.keys()):
            rel_frob = math.sqrt(cat_diff_sq[cat] / max(1e-12, cat_base_sq[cat])) * 100.0
            mean_c = float(np.mean(cat_cos[cat]))
            drift_report[pair_key]["by_module"][cat] = {
                "rel_frobenius_pct": round(rel_frob, 4),
                "mean_cosine_sim": round(mean_c, 6)
            }
        print(f"  ✅ Drift [{pair_key}]: Global Frobenius Shift = {drift_report[pair_key]['global_rel_frobenius_pct']:.4f}%")

    # Free CPU RAM
    del states
    gc.collect()
    return drift_report


def compute_baselines_on_windows(windows: np.ndarray, targets: np.ndarray) -> dict:
    """
    Evaluates 4 deterministic quantitative baselines on the exact same [N, T, F] windows:
    Feature 0 in features.py is log_return: log(close_t / close_{t-1}).
    1. Always-Long (Bull Market Drift): predicts +1 always.
    2. 1-Bar Momentum Continuation: predicts sign(r_t) continues.
    3. 5-Bar Microstructure Mean-Reversion: predicts -sign(sum_{k=0..4} r_{t-k}).
    4. EMA(12, 26) Trend Crossover: predicts sign(EMA_12(r) - EMA_26(r)).
    """
    n_samples = len(windows)
    num_h = targets.shape[1]
    last_r = windows[:, -1, 0]  # [N]
    last5_r = np.mean(windows[:, -5:, 0], axis=1)  # [N]

    # Fast vectorized EMA12 vs EMA26 on last 64 bars of log-returns
    tail_r = windows[:, -64:, 0]  # [N, 64]
    alpha12 = 2.0 / 13.0
    alpha26 = 2.0 / 27.0
    w12 = np.array([(1.0 - alpha12) ** (63 - i) for i in range(64)], dtype=np.float32)
    w12 /= np.sum(w12)
    w26 = np.array([(1.0 - alpha26) ** (63 - i) for i in range(64)], dtype=np.float32)
    w26 /= np.sum(w26)
    macd_sig = (tail_r @ w12) - (tail_r @ w26)  # [N]

    baselines = {}

    # 1. Always-Long
    p_long = np.ones((n_samples, num_h), dtype=np.float32) * 0.001
    l_long = np.ones((n_samples, num_h), dtype=np.float32)
    baselines["Always_Long_Drift"] = compute_rich_horizon_metrics(p_long, l_long, targets)

    # 2. 1-Bar Momentum
    sig_mom = np.repeat(last_r[:, None], num_h, axis=1)
    l_mom = np.where(sig_mom >= 0.0, 1.0, -1.0).astype(np.float32)
    baselines["Momentum_1Bar"] = compute_rich_horizon_metrics(sig_mom, l_mom, targets)

    # 3. 5-Bar Mean Reversion
    sig_rev = np.repeat((-last5_r)[:, None], num_h, axis=1)
    l_rev = np.where(sig_rev > 0.0, 1.0, -1.0).astype(np.float32)
    baselines["MeanReversion_5Bar"] = compute_rich_horizon_metrics(sig_rev, l_rev, targets)

    # 4. EMA(12,26) Crossover
    sig_ema = np.repeat(macd_sig[:, None], num_h, axis=1)
    l_ema = np.where(sig_ema > 0.0, 1.0, -1.0).astype(np.float32)
    baselines["EMA_12_26_Crossover"] = compute_rich_horizon_metrics(sig_ema, l_ema, targets)

    return baselines


def main():
    t_start = time.time()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    hf_token = os.environ.get("HF_TOKEN", "")
    repo_id = os.environ.get("HF_REPO_ID", "Ultron28/ohlcv-transformer")

    print("=" * 80)
    print(f"⚡ ULTRON // DETERMINISTIC TIER-1 SCIENTIFIC AUDIT ENGINE ({device})")
    print("=" * 80)

    os.makedirs("/content/ohlcv-transformer/checkpoints_1b", exist_ok=True)
    os.makedirs("/content/logs", exist_ok=True)

    # 1. Ensure all 5 checkpoints exist locally on the A100 VM
    ckpt_names = {
        "epoch1": "checkpoint_1b_epoch1.pt",
        "epoch2": "checkpoint_1b_epoch2.pt",
        "epoch3": "checkpoint_1b_epoch3.pt",
        "wf2024": "checkpoint_1b_wf2024.pt",
        "wf2025": "checkpoint_1b_wf2025.pt"
    }
    ckpt_paths = {}
    for key, fname in ckpt_names.items():
        local_p = f"/content/ohlcv-transformer/checkpoints_1b/{fname}"
        if not os.path.exists(local_p):
            print(f"☁️ Downloading {fname} from Hugging Face ({repo_id})...")
            local_p = hf_hub_download(
                repo_id=repo_id,
                filename=f"checkpoints_1b/{fname}",
                token=hf_token,
                local_dir="/content/ohlcv-transformer"
            )
        ckpt_paths[key] = local_p
        print(f"  ✅ Verified {key}: {local_p} ({os.path.getsize(local_p)/1e9:.2f} GB)")

    # 2. Instantiate model & audit exact parameter counts
    model_cfg = Model1BConfig()
    train_cfg = Training1BConfig()
    model = FinancialConformer1B(model_cfg).to(device)
    if device.type == "cuda":
        torch.backends.cuda.matmul.allow_tf32 = True
        torch.backends.cudnn.allow_tf32 = True

    param_audit = audit_architecture_parameters(model)
    print(f"\n📐 [Architecture Audit] Exact Parameter Breakdown: {json.dumps(param_audit, indent=2)}")

    # 3. Compute Exact Frobenius Weight Drift Matrix across all checkpoints
    weight_drift = compute_weight_drift_matrix(ckpt_paths)

    # 4. Prepare Representative OOS Benchmark Cohort (2024, 2025, 2026)
    # Using 3 shards from 2024 (Jan, Jun, Dec), 3 from 2025 (Jan, Jun, Dec), and all 3 from 2026 (Jan, Feb, Mar)
    # = 9 shards (~31,680 windows) for fast, exact multi-checkpoint comparison + baselines + backtests!
    eval_shards_by_year = {
        "2024": ["data/ohlcv_2024-01.parquet", "data/ohlcv_2024-06.parquet", "data/ohlcv_2024-12.parquet"],
        "2025": ["data/ohlcv_2025-01.parquet", "data/ohlcv_2025-06.parquet", "data/ohlcv_2025-12.parquet"],
        "2026": ["data/ohlcv_2026-01.parquet", "data/ohlcv_2026-02.parquet", "data/ohlcv_2026-03.parquet"]
    }

    # Cache the windowed datasets in CPU RAM (9 shards * 3520 windows * 2048 * 8 * 4 bytes ~ 2.0 GB RAM out of 85 GB RAM!)
    print("\n📦 [Data Ingestion] Loading & caching OOS evaluation cohorts into High-RAM...")
    cached_cohorts = {}
    for yr, s_list in eval_shards_by_year.items():
        data_cfg = DatasetConfig(
            seq_len=train_cfg.seq_len,
            window_stride=train_cfg.window_stride,
            pred_horizons=train_cfg.pred_horizons,
            hf_token=hf_token
        )
        manager = StreamingShardManager(data_cfg)
        manager.parquet_files = s_list
        w_list, t_list = [], []
        for s_ds in manager.stream_monthly_shards(start_shard_idx=0):
            n_take = min(len(s_ds), 3520)
            w_list.append(s_ds.windows[:n_take].numpy())
            t_list.append(s_ds.targets[:n_take].numpy())
        cached_cohorts[yr] = (np.concatenate(w_list, axis=0), np.concatenate(t_list, axis=0))
        print(f"  ✅ Cached Year {yr}: {len(cached_cohorts[yr][0])} windows")

    def _eval_loaded_model_on_array(w_arr: np.ndarray, batch_size: int = 64):
        model.eval()
        preds_l, logits_l = [], []
        n = len(w_arr)
        with torch.no_grad():
            for i in range(0, n, batch_size):
                bx = torch.from_numpy(w_arr[i:i+batch_size]).to(device, non_blocking=True)
                with torch.amp.autocast('cuda', dtype=torch.bfloat16):
                    pr, pl = model(bx)
                preds_l.append(pr.float().cpu().numpy())
                logits_l.append(pl.float().cpu().numpy())
        return np.concatenate(preds_l, axis=0), np.concatenate(logits_l, axis=0)

    # 5. Evaluate Baselines on 2024, 2025, 2026
    print("\n📊 [Baselines Audit] Computing deterministic quantitative baselines across 2024, 2025, 2026...")
    baselines_by_year = {}
    for yr, (w_arr, t_arr) in cached_cohorts.items():
        baselines_by_year[yr] = compute_baselines_on_windows(w_arr, t_arr)
        print(f"  ✅ Year {yr} Baselines computed.")

    # 6. Multi-Epoch & Walk-Forward Evaluation across 2024, 2025, 2026
    print("\n🚀 [Multi-Checkpoint OOS Benchmark] Evaluating Epoch 1, Epoch 2, Epoch 3, WF-2024, WF-2025...")
    checkpoints_eval = {}
    raw_outputs_epoch3 = {}

    for ckpt_key in ["epoch1", "epoch2", "epoch3", "wf2024", "wf2025"]:
        t_c = time.time()
        state = torch.load(ckpt_paths[ckpt_key], map_location=device, weights_only=False)
        model.load_state_dict(state.get("model_state", state), strict=False)
        model.eval()

        checkpoints_eval[ckpt_key] = {}
        for yr, (w_arr, t_arr) in cached_cohorts.items():
            pr, pl = _eval_loaded_model_on_array(w_arr, batch_size=64)
            rich_m = compute_rich_horizon_metrics(pr, pl, t_arr)
            checkpoints_eval[ckpt_key][yr] = rich_m
            if ckpt_key == "epoch3":
                raw_outputs_epoch3[yr] = (pr, pl, t_arr)

            print(f"  • [{ckpt_key} | {yr}] RawAcc={rich_m['overall_raw_dir_acc']}% | "
                  f"NonZeroAcc={rich_m['overall_nonzero_dir_acc']}% | "
                  f"BalAcc={rich_m['overall_balanced_acc']}% | "
                  f"AUC={rich_m['overall_roc_auc']} | MeanIC={rich_m['overall_mean_ic']}")
        print(f"  ⏱️ Completed {ckpt_key} in {time.time() - t_c:.1f}s")

    # 7. Strictly Chronological Causal Execution Backtest on Epoch 3 (2024, 2025, 2026)
    print("\n📈 [Chronological Causal Backtest] Running multi-horizon & confidence-hurdle simulations...")
    causal_backtests = {}
    for yr, (pr, pl, t_arr) in raw_outputs_epoch3.items():
        causal_backtests[yr] = run_causal_chronological_backtest(pr, pl, t_arr)
        print(f"  ✅ Causal chronological backtest completed for Year {yr}.")

    # 8. Full 7-Horizon Mechanistic Inception Kernel Ablation (on 2024 Cohort = 10,560 windows)
    print("\n🔬 [Mechanistic Kernel Ablation] Running full 7-horizon + scale-group ablation on 2024 OOS Cohort...")
    state_ep3 = torch.load(ckpt_paths["epoch3"], map_location=device, weights_only=False)
    model.load_state_dict(state_ep3.get("model_state", state_ep3), strict=False)
    model.eval()

    w_ab, t_ab = cached_cohorts["2024"]
    pr_base, pl_base = raw_outputs_epoch3["2024"][0], raw_outputs_epoch3["2024"][1]
    base_ab_metrics = compute_rich_horizon_metrics(pr_base, pl_base, t_ab)

    kernels = [3, 7, 15, 31, 63, 127]
    ablation_configs = [
        ("Mask_K3", [0]),
        ("Mask_K7", [1]),
        ("Mask_K15", [2]),
        ("Mask_K31", [3]),
        ("Mask_K63", [4]),
        ("Mask_K127", [5]),
        ("Mask_Micro_K3_K7", [0, 1]),
        ("Mask_Meso_K15_K31", [2, 3]),
        ("Mask_Macro_K63_K127", [4, 5])
    ]

    tokenizer = model.tokenizer
    orig_weights = [b[0].weight.data.clone() for b in tokenizer.branches]
    orig_biases = [b[0].bias.data.clone() if b[0].bias is not None else None for b in tokenizer.branches]

    kernel_ablation_full = {
        "n_windows_evaluated": int(len(w_ab)),
        "baseline": base_ab_metrics,
        "ablations": {}
    }

    for ab_name, branch_indices in ablation_configs:
        # Zero out specified branches (both weight and bias!)
        for idx in branch_indices:
            tokenizer.branches[idx][0].weight.data.zero_()
            if tokenizer.branches[idx][0].bias is not None:
                tokenizer.branches[idx][0].bias.data.zero_()

        pr_m, pl_m = _eval_loaded_model_on_array(w_ab, batch_size=64)
        m_metrics = compute_rich_horizon_metrics(pr_m, pl_m, t_ab)

        # Restore original weights and biases immediately
        for idx in branch_indices:
            tokenizer.branches[idx][0].weight.data.copy_(orig_weights[idx])
            if orig_biases[idx] is not None:
                tokenizer.branches[idx][0].bias.data.copy_(orig_biases[idx])

        # Compute exact per-horizon deltas and continuous L2 logit/prediction sensitivity
        per_h_delta = {}
        logit_l2_pcts = []
        pred_l2_pcts = []
        for h_idx, h_name in enumerate(HORIZON_NAMES):
            b_h = base_ab_metrics["per_horizon"][h_name]
            m_h = m_metrics["per_horizon"][h_name]

            l_diff_norm = float(np.linalg.norm(pl_m[:, h_idx] - pl_base[:, h_idx]))
            l_base_norm = max(1e-12, float(np.linalg.norm(pl_base[:, h_idx])))
            l_rel_pct = (l_diff_norm / l_base_norm) * 100.0
            logit_l2_pcts.append(l_rel_pct)

            p_diff_norm = float(np.linalg.norm(pr_m[:, h_idx] - pr_base[:, h_idx]))
            p_base_norm = max(1e-12, float(np.linalg.norm(pr_base[:, h_idx])))
            p_rel_pct = (p_diff_norm / p_base_norm) * 100.0
            pred_l2_pcts.append(p_rel_pct)

            per_h_delta[h_name] = {
                "raw_dir_acc": m_h["raw_dir_acc"],
                "delta_raw_dir_acc": round(m_h["raw_dir_acc"] - b_h["raw_dir_acc"], 3),
                "balanced_acc": m_h["balanced_acc"],
                "delta_balanced_acc": round(m_h["balanced_acc"] - b_h["balanced_acc"], 3),
                "roc_auc": m_h["roc_auc"],
                "delta_roc_auc": round(m_h["roc_auc"] - b_h["roc_auc"], 4),
                "spearman_ic": m_h["spearman_ic"],
                "delta_spearman_ic": round(m_h["spearman_ic"] - b_h["spearman_ic"], 4),
                "logit_l2_perturbation_pct": round(l_rel_pct, 2),
                "return_pred_l2_perturbation_pct": round(p_rel_pct, 2)
            }

        kernel_ablation_full["ablations"][ab_name] = {
            "masked_kernels": [kernels[i] for i in branch_indices],
            "overall_raw_dir_acc": m_metrics["overall_raw_dir_acc"],
            "delta_overall_raw_dir_acc": round(m_metrics["overall_raw_dir_acc"] - base_ab_metrics["overall_raw_dir_acc"], 3),
            "overall_balanced_acc": m_metrics["overall_balanced_acc"],
            "delta_overall_balanced_acc": round(m_metrics["overall_balanced_acc"] - base_ab_metrics["overall_balanced_acc"], 3),
            "overall_roc_auc": m_metrics["overall_roc_auc"],
            "delta_overall_roc_auc": round(m_metrics["overall_roc_auc"] - base_ab_metrics["overall_roc_auc"], 4),
            "overall_mean_ic": m_metrics["overall_mean_ic"],
            "delta_overall_mean_ic": round(m_metrics["overall_mean_ic"] - base_ab_metrics["overall_mean_ic"], 4),
            "mean_logit_l2_perturbation_pct": round(float(np.mean(logit_l2_pcts)), 2),
            "mean_return_l2_perturbation_pct": round(float(np.mean(pred_l2_pcts)), 2),
            "per_horizon": per_h_delta
        }
        print(f"  • [{ab_name:20s}] OverallAcc={m_metrics['overall_raw_dir_acc']}% "
              f"(Δ={kernel_ablation_full['ablations'][ab_name]['delta_overall_raw_dir_acc']:+.2f}%) | "
              f"Logit L2 Shift={kernel_ablation_full['ablations'][ab_name]['mean_logit_l2_perturbation_pct']:.2f}% | "
              f"Ret L2 Shift={kernel_ablation_full['ablations'][ab_name]['mean_return_l2_perturbation_pct']:.2f}%")

    # 9. Assemble Master Verified Audit JSON
    master_audit = {
        "audit_timestamp": time.time(),
        "total_elapsed_sec": round(time.time() - t_start, 2),
        "hardware": torch.cuda.get_device_name(0) if device.type == "cuda" else "CPU",
        "architecture_parameter_audit": param_audit,
        "weight_drift_frobenius_audit": weight_drift,
        "multi_checkpoint_oos_audit": checkpoints_eval,
        "quantitative_baselines_oos": baselines_by_year,
        "causal_chronological_backtest": causal_backtests,
        "mechanistic_kernel_ablation_7h": kernel_ablation_full
    }

    out_path = "/content/logs/tier1_verified_audit.json"
    with open(out_path, "w") as f:
        json.dump(master_audit, f, indent=2)

    print(f"\n🏆 [ULTRON SCIENTIFIC AUDIT COMPLETE] Saved to {out_path} in {time.time() - t_start:.1f}s!")

    if hf_token:
        try:
            api = HfApi(token=hf_token)
            api.upload_file(
                path_or_fileobj=out_path,
                path_in_repo="logs/tier1_verified_audit.json",
                repo_id=repo_id,
                repo_type="model"
            )
            print("☁️ Uploaded logs/tier1_verified_audit.json to Hugging Face Model Hub!")
        except Exception as e:
            print(f"⚠️ HF upload notice: {e}")

if __name__ == "__main__":
    main()