"""
Agent-reviewed momentum (US): the overnight_recovery score picks 50
candidates, an LLM agent scores each one from its filings, fundamentals and
business logic, and a fixed blend of the two picks 10.

Daily workflow (``agent_review.py`` at the repo root, skill
``.claude/skills/agent-stock-review``):
  1. ``agent_review.py candidates`` writes the candidate pack for the next
     session to data/agent/us/candidates_<as_of>.json; ``as_of`` is the last
     completed session, i.e. the bars the next decision uses.
  2. The agent researches the candidates and writes a scores JSON.
  3. ``agent_review.py submit <file>`` validates it and stores it as
     data/agent/us/scores_<as_of>_<HHMMSS>.json with a ``created_at`` stamp.
     Files are never overwritten, so a replay sees what was known at the time.

Decision rule, fixed on 2026-09-30 before any forward data:
  final = pct + AGENT_WEIGHT * agent_score
    pct          the name's place among the 50 candidates by the
                 overnight_recovery score, 1 for the best down to 0
    agent_score  integer in [-2, 2]; 0 for a name the file does not score
  Vetoed names are dropped. The top 10 by ``final`` are held, with the base
  strategy's hysteresis (a held name stays while its final rank is below
  top_n * (1 + buffer) = 30) and on the same rebalance rows as
  overnight_recovery_top10 (every 10 rows). Between rebalances, a held name
  vetoed in a fresh file (scored on the latest bars) is sold and its weight
  goes to the best candidate not held.

Point in time: a scores file may steer the weights traded at the close of
session s only if its as_of is before s, it is at most MAX_AGE_DAYS calendar
days old, and it was stamped before 16:00 New York time on s. Without such a
file the strategy uses the quant order alone and holds exactly what
overnight_recovery_top10 holds, so the difference between the two books
measures the agent.
"""

from __future__ import annotations

import glob
import json
import os
from dataclasses import dataclass, field

import numpy as np
import pandas as pd

from strategies.base import BarStrategy
from strategies.robust_momentum import SIGNALS, overnight_momentum, recovery

SCORE_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "data", "agent", "us")
AGENT_WEIGHT = 0.25
MAX_AGE_DAYS = 7
SCORE_RANGE = (-2, 2)
VETO_REASONS = ("pending_cash_acquisition", "fraud_or_restatement", "regulatory_action",
                "large_equity_offering", "binary_event")
_NY = "America/New_York"


@dataclass
class AgentScores:
    as_of: pd.Timestamp                    # last session the agent's data covered
    created_at: pd.Timestamp               # UTC stamp from ``submit``
    scores: dict[str, int] = field(default_factory=dict)
    vetoes: frozenset[str] = frozenset()


def load_scores(score_dir: str = SCORE_DIR) -> list[AgentScores]:
    out = []
    for path in sorted(glob.glob(os.path.join(score_dir, "scores_*.json"))):
        with open(path) as fh:
            d = json.load(fh)
        rows = d.get("scores", {})
        out.append(AgentScores(
            as_of=pd.Timestamp(d["as_of"]),
            created_at=pd.Timestamp(d["created_at"]).tz_convert("UTC"),
            scores={t: int(v.get("score", 0)) for t, v in rows.items()},
            vetoes=frozenset(t for t, v in rows.items() if v.get("veto")),
        ))
    return out


def session_deadline(session: pd.Timestamp) -> pd.Timestamp:
    """16:00 New York time on ``session``, in UTC."""
    return (pd.Timestamp(session.date()).tz_localize(_NY) + pd.Timedelta(hours=16)).tz_convert("UTC")


def usable_scores(files: list[AgentScores], session: pd.Timestamp) -> AgentScores | None:
    """The latest file that may steer the close of ``session``."""
    deadline = session_deadline(session)
    ok = [f for f in files if f.as_of < session and f.created_at <= deadline
          and (session - f.as_of).days <= MAX_AGE_DAYS]
    return max(ok, key=lambda f: (f.as_of, f.created_at)) if ok else None


def blended_order(score_row: np.ndarray, pool: int, agent: AgentScores | None,
                  columns: list[str]) -> list[int]:
    """Candidate column indices ordered by ``final`` (best first), vetoes removed."""
    valid = np.where(~np.isnan(score_row))[0]
    cand = valid[np.argsort(-score_row[valid], kind="stable")][:pool]
    k = len(cand)
    pct = 1.0 - np.arange(k) / max(k - 1, 1)
    if agent is None:
        return list(cand)
    final = np.array([pct[j] + AGENT_WEIGHT * agent.scores.get(columns[i], 0) for j, i in enumerate(cand)])
    order = [cand[j] for j in np.argsort(-final, kind="stable")]
    return [i for i in order if columns[i] not in agent.vetoes]


class AgentReviewStrategy(BarStrategy):
    def __init__(self, signal: str = "overnight_recovery", top_n: int = 10, pool: int = 50,
                 rebal_freq: int = 10, buffer: float = 2.0, warmup: int = 252,
                 score_dir: str = SCORE_DIR):
        self.fields, self._score = SIGNALS[signal]
        self.top_n, self.pool, self.rebal_freq = top_n, pool, rebal_freq
        self.buffer, self.warmup, self.score_dir = buffer, warmup, score_dir

    def quant_score(self, bars: dict[str, pd.DataFrame]) -> pd.DataFrame:
        """The base score, NaN for names without a price (or outside the index
        when a backtest passes ``eligible``); row t uses bars through t."""
        close = bars["close"]
        panel = {f: bars[f].reindex(index=close.index, columns=close.columns) for f in self.fields}
        eligible = close.notna()
        if "eligible" in bars:
            eligible &= bars["eligible"].reindex_like(close).fillna(False).astype(bool)
        panel["eligible"] = eligible
        return self._score(panel).where(eligible)

    def generate_weights(self, bars: dict[str, pd.DataFrame]) -> pd.DataFrame:
        close = bars["close"]
        S = self.quant_score(bars).to_numpy(dtype=float)
        C = close.ffill().to_numpy(dtype=float)
        dates, cols = close.index, list(close.columns)
        files = load_scores(self.score_dir)
        n, m = S.shape
        W = np.zeros((n, m))
        w = np.zeros(m)
        for k in range(self.warmup, n):
            if w.any():                      # drift with prices since the previous row
                g = C[k] / C[k - 1]
                w = w * np.where(np.isfinite(g), g, 1.0)
                w = w / w.sum() if w.sum() > 0 else w
            # weights of pre-shift row k trade at the close of the next session
            session = dates[k + 1] if k + 1 < n else dates[k] + pd.offsets.BDay(1)
            agent = usable_scores(files, session) if files else None
            if (k - self.warmup) % self.rebal_freq == 0:
                w = self._rebalance(S[k], w > 0, agent, cols)
            elif agent is not None and agent.as_of == dates[k] and w.any():
                w = self._replace_vetoed(S[k], w, agent, cols)
            W[k] = w
        weights = pd.DataFrame(W, index=dates, columns=close.columns)
        return weights.shift(1).fillna(0.0)

    def _rebalance(self, row, held, agent, cols) -> np.ndarray:
        w = np.zeros(len(cols))
        if (~np.isnan(row)).sum() < self.top_n:
            return w
        order = blended_order(row, self.pool, agent, cols)
        rank = {i: r for r, i in enumerate(order)}
        keep = sorted((i for i in np.where(held)[0] if rank.get(i, np.inf) < self.top_n * (1 + self.buffer)),
                      key=lambda i: rank[i])[:self.top_n]
        chosen = list(keep)
        for i in order:
            if len(chosen) >= self.top_n:
                break
            if i not in chosen:
                chosen.append(i)
        if chosen:
            w[chosen] = 1.0 / len(chosen)
        return w

    def _replace_vetoed(self, row, w, agent, cols) -> np.ndarray:
        vetoed = [i for i in np.where(w > 0)[0] if cols[i] in agent.vetoes]
        if not vetoed:
            return w
        w = w.copy()
        spare = [i for i in blended_order(row, self.pool, agent, cols) if w[i] == 0]
        for i in vetoed:
            if spare:
                w[spare.pop(0)] = w[i]
            w[i] = 0.0
        return w / w.sum() if w.sum() > 0 else w

    def candidates(self, bars: dict[str, pd.DataFrame]) -> pd.DataFrame:
        """The ``pool`` candidates on the latest bars (what the next session's
        decision ranks), best first, with the base score and its factors."""
        close = bars["close"]
        score = self.quant_score(bars).iloc[-1].dropna().sort_values(ascending=False, kind="stable")
        top = score.head(self.pool)
        on = overnight_momentum(bars["open"], close).iloc[-1]
        rec = recovery(close).iloc[-1]
        return pd.DataFrame({
            "rank": np.arange(1, len(top) + 1),
            "quant_pct": 1.0 - np.arange(len(top)) / max(len(top) - 1, 1),
            "overnight_mom_12m": on.reindex(top.index).round(4),
            "recovery_126d": rec.reindex(top.index).round(4),
        }, index=top.index)