"""Frozen policy library: a nearest-neighbour replay of *recorded* LLM decisions.

Every online LLM decision is logged with a feature vector of its observation; the library maps a new
observation to the recorded decision of the nearest neighbour in feature space (standardized Euclidean
distance) and rescales the quantity to the current account. This is NOT an adaptive agent: it is a
frozen, non-reasoning approximation used only (i) to scale populations to 1,000+ traders at zero API
cost and (ii) as the 'cached' arm of the cached-vs-online comparison. Runs using it are labelled
provider = 'library'."""
from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import pandas as pd

from ..agents.base import Observation

FEATURES = ["mispricing", "momentum", "signal", "inv_ratio", "cash_ratio", "spread_bp", "has_open"]


def features(obs: Observation, asset: int = 0) -> np.ndarray:
    b = obs.books[asset]
    tc = obs.tick_cents[asset]
    mid = max(b.mid, 1.0)
    h = obs.mid_history[asset]
    mom = (h[-1] / h[0] - 1.0) if len(h) > 1 and h[0] > 0 else 0.0
    pv = obs.public_value[asset] / tc
    sig = obs.private_signals.get(asset, 0.0)
    pos_val = obs.positions.get(asset, 0) * mid * tc
    eq = max(obs.equity, 1.0)
    spread = ((b.best_ask - b.best_bid) / mid * 1e4) if (b.best_ask is not None and b.best_bid is not None) else 50.0
    return np.array([(pv - mid) / mid, mom, sig / 10.0, pos_val / eq, obs.free_cash / eq, spread / 10.0, float(bool(obs.open_orders))])


class DecisionLibrary:
    def __init__(self, X: np.ndarray, decisions: list[dict]) -> None:
        self.X = X
        self.mu = X.mean(axis=0)
        self.sd = X.std(axis=0) + 1e-9
        self.Z = (X - self.mu) / self.sd
        self.decisions = decisions

    @classmethod
    def from_runs(cls, run_dirs: list[str | Path]) -> "DecisionLibrary":
        X, D = [], []
        for rd in run_dirs:
            p = Path(rd) / "llm_decisions.parquet"
            if not p.exists():
                continue
            df = pd.read_parquet(p)
            df = df[df["valid"] & df["features"].notna()] if "features" in df else df.iloc[0:0]
            for _, r in df.iterrows():
                X.append(json.loads(r["features"]))
                D.append({"action": r["action"], "side": r["side"], "order_type": r["order_type"], "quantity": int(r["quantity"]),
                          "rel_price": float(r["rel_price"]) if "rel_price" in r and pd.notna(r["rel_price"]) else 0.0,
                          "rationale": "library replay"})
        if not X:
            raise ValueError("no recorded decisions with features in the given runs")
        return cls(np.array(X, dtype=float), D)

    def save(self, path: str | Path) -> None:
        np.savez(path, X=self.X, decisions=json.dumps(self.decisions))

    @classmethod
    def load(cls, path: str | Path) -> "DecisionLibrary":
        z = np.load(path, allow_pickle=False)
        return cls(z["X"], json.loads(str(z["decisions"])))

    def lookup(self, f: np.ndarray, rng: np.random.Generator, k: int = 5) -> dict:
        z = (f - self.mu) / self.sd
        d = ((self.Z - z) ** 2).sum(axis=1)
        idx = np.argpartition(d, min(k, len(d) - 1))[:k]
        return dict(self.decisions[int(rng.choice(idx))])

    def __len__(self) -> int:
        return len(self.decisions)


class LibraryProvider:
    """Provider-shaped wrapper so the LLMPolicy pipeline (validation, logging) is reused unchanged."""
    name = "library"

    def __init__(self, library_path: str, seed: int = 0, **_: object) -> None:
        self.lib = DecisionLibrary.load(library_path)
        self.rng = np.random.Generator(np.random.PCG64(seed))

    def complete_obs(self, obs: Observation, mid_units: float) -> str:
        dec = self.lib.lookup(features(obs), self.rng)
        if dec["action"] == "place" and dec["order_type"] == "limit":
            dec["limit_price"] = round(mid_units * (1.0 + dec.get("rel_price", 0.0)), 2)
        else:
            dec["limit_price"] = None
        dec.setdefault("cancel_order_ids", [])
        dec.setdefault("confidence", None)
        return json.dumps(dec)