"""Competition submission entrypoint — streaming global-norm GRU ensemble (CPU/ONNX).

Reads ensemble_config.json (next to this file):
  {
    "global_mean": [...32...], "global_std": [...32...],
    "models": [{"onnx": "m1.onnx", "scale": [s0, s1]}, ...]
  }
Each model is a stateful one-step GRU exported by export_gru_onnx.py. Predictions
are scale-normalized per model and averaged (weighted-Pearson is shift/scale
invariant, so only relative scale matters). Recurrent state resets per sequence.
"""
from __future__ import annotations

import json
import os
from pathlib import Path
from typing import List, Optional

import numpy as np

os.environ.setdefault("OMP_NUM_THREADS", "1")
os.environ.setdefault("MKL_NUM_THREADS", "1")
os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")

import onnxruntime as ort

from utils import DataPoint

RAW_FEATURE_SIZE = 32
# Superset of rolling windows; 160-feat models use the first 5 blocks (cols 0:160),
# 224-feat models use all 7 blocks. Column order: norm, lag1, delta, roll5, roll10,
# roll20, roll40 -> slicing the prefix is exact.
ROLLING_WINDOWS = (5, 10, 20, 40)


class _OnlineGlobalFeatures:
    def __init__(self, mean, std):
        self.mean = np.asarray(mean, dtype=np.float32)
        self.std = np.asarray(std, dtype=np.float32)
        self.maxw = max(ROLLING_WINDOWS)
        self.reset()

    def reset(self):
        self._hist: List[np.ndarray] = []
        self._prev: Optional[np.ndarray] = None

    def update(self, raw_state):
        x = np.asarray(raw_state, dtype=np.float32).reshape(-1)
        norm = (x - self.mean) / self.std
        if self._prev is None:
            lag1 = np.zeros(RAW_FEATURE_SIZE, dtype=np.float32)
            delta = np.zeros(RAW_FEATURE_SIZE, dtype=np.float32)
        else:
            lag1 = self._prev
            delta = norm - self._prev
        self._hist.append(norm)
        if len(self._hist) > self.maxw:
            self._hist = self._hist[-self.maxw:]
        parts = [norm, lag1, delta]
        hist = np.asarray(self._hist, dtype=np.float32)
        for W in ROLLING_WINDOWS:
            parts.append(hist[-W:].mean(axis=0).astype(np.float32))
        self._prev = norm
        return np.concatenate(parts).astype(np.float32)


class _Model:
    def __init__(self, onnx_path: Path):
        so = ort.SessionOptions()
        so.intra_op_num_threads = 1
        so.inter_op_num_threads = 1
        so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        self.sess = ort.InferenceSession(str(onnx_path), so, providers=["CPUExecutionProvider"])
        ins = self.sess.get_inputs()
        self.in_feat = ins[0].name
        self.in_h = ins[1].name
        self.in_size = int(ins[0].shape[2])     # 160 or 224
        hshape = ins[1].shape  # (L, 1, H)
        self.L = int(hshape[0]); self.H = int(hshape[2])
        self.out_names = [o.name for o in self.sess.get_outputs()]
        self.reset()

    def reset(self):
        self.h = np.zeros((self.L, 1, self.H), dtype=np.float32)

    def step(self, feat_1x1xF: np.ndarray) -> np.ndarray:
        x = feat_1x1xF[:, :, : self.in_size]
        pred, h1 = self.sess.run(self.out_names, {self.in_feat: x, self.in_h: self.h})
        self.h = h1.astype(np.float32)
        return pred[0].astype(np.float32)


class PredictionModel:
    def __init__(self, config_path: Optional[str] = None):
        base = Path(__file__).resolve().parent
        cfg = json.loads((base / (config_path or "ensemble_config.json")).read_text())
        self.feats = _OnlineGlobalFeatures(cfg["global_mean"], cfg["global_std"])
        self.models = [_Model(base / m["onnx"]) for m in cfg["models"]]
        # per-model, per-target weights (M, 2): t0 fitted, t1 uniform (see blend.py)
        self.weights = np.asarray([m["weight"] for m in cfg["models"]], dtype=np.float32)
        self.current_seq_ix: Optional[int] = None

    def _reset(self):
        self.feats.reset()
        for m in self.models:
            m.reset()

    def predict(self, data_point: DataPoint):
        if self.current_seq_ix != data_point.seq_ix:
            self._reset()
            self.current_seq_ix = data_point.seq_ix
        feat = self.feats.update(data_point.state)
        x = feat.reshape(1, 1, -1).astype(np.float32)
        # advance every model's recurrent state every step (warmup included)
        acc = np.zeros(2, dtype=np.float32)
        for i, m in enumerate(self.models):
            acc += self.weights[i] * m.step(x)
        if not data_point.need_prediction:
            return None
        return np.clip(acc, -6.0, 6.0).astype(np.float32)


if __name__ == "__main__":
    from utils import ScorerStepByStep
    data_path = Path(__file__).resolve().parent / "data" / "valid.parquet"
    model = PredictionModel()
    scorer = ScorerStepByStep(str(data_path))
    print(scorer.score(model))