"""
rule_based_agent.py
───────────────────
A clinically-grounded rule-based oracle agent for CogTraceEnv.

This agent uses hard-coded clinical knowledge to achieve near-optimal
performance (~0.82 on Task 3). It serves two purposes:
  1. Proves CogTraceEnv is solvable (not too hard / noisy)
  2. Provides an upper-bound benchmark score for comparison

The agent's logic mirrors what a trained clinical monitoring system
would do: track running baselines, detect multi-signal deviations,
and respond proportionally.

Usage:
  python rule_based_agent.py               # run full benchmark
  python rule_based_agent.py --agent random
  python rule_based_agent.py --agent always_alert
"""

from __future__ import annotations

import argparse
import json
import os
import time
from typing import Any, Dict, List, Tuple
from dataclasses import dataclass, field
from collections import deque

import numpy as np

# ── CDR Stage Baselines (used to compute z-scores) ───────────────────────────
# Mirrors the values in patient_simulator.py CDR_PROFILES

CDR_BASELINES = {
    # key: (mean, std) for a mid-stage patient (Stage 2 / CDR 1)
    "typing_delay_delta":       (0.60, 0.25),
    "sleep_hours":              (6.20, 0.80),
    "routine_adherence_score":  (0.68, 0.10),
    "speech_pause_freq":        (2.80, 0.60),
    "memory_lapse_count":       (2.50, 1.00),
}

# Signals where HIGH value = worse (anomaly direction = positive)
HIGH_IS_BAD = {"typing_delay_delta", "speech_pause_freq", "memory_lapse_count"}

# Signals where LOW value = worse (anomaly direction = negative)
LOW_IS_BAD  = {"sleep_hours", "routine_adherence_score"}


@dataclass
class OracleState:
    """Running state maintained by the oracle agent across steps."""
    history:       deque = field(default_factory=lambda: deque(maxlen=7))
    alert_history: List[int] = field(default_factory=list)
    cooldown:      int = 0  # steps remaining before next alert is allowed


class RuleBasedOracle:
    """
    Clinical rule-based agent for CogTraceEnv.

    Strategy:
    ---------
    1. Track a 7-day rolling window of observations.
    2. Compute per-signal z-scores relative to the CDR baseline.
    3. Count how many signals are anomalous (z > 1.5 in the bad direction).
    4. Alert proportionally: 3+ signals → escalate, 2 → medium, 1 → soft.
    5. Apply 2-step cooldown after any alert to suppress spam.
    6. Also monitor 7-day trend slopes for gradual decline detection.
    """

    ANOMALY_THRESHOLD_Z = 1.5   # z-score to flag a signal as anomalous
    TREND_THRESHOLD     = 0.06  # slope magnitude to flag a trend
    COOLDOWN_STEPS      = 2     # steps to wait after raising an alert

    def __init__(self):
        self.state = OracleState()

    def reset(self):
        self.state = OracleState()

    def act(self, obs: Dict[str, Any]) -> int:
        """
        Given an observation dict, return action (0–3).
        """
        self.state.history.append(obs)

        # Count anomalous signals via z-score
        n_anomalous = self._count_anomalous_signals(obs)

        # Check trend slope anomalies (from obs fields)
        n_trend_anomalous = self._count_trend_anomalies(obs)

        # Compute combined severity score
        severity = n_anomalous + 0.5 * n_trend_anomalous

        # Cooldown suppression
        if self.state.cooldown > 0:
            self.state.cooldown -= 1
            action = 0  # do_nothing during cooldown
        elif severity >= 3.0:
            action = 3  # escalate
        elif severity >= 2.0:
            action = 2  # medium_alert
        elif severity >= 1.0:
            action = 1  # soft_alert
        else:
            action = 0  # do_nothing

        # Apply cooldown if we alerted
        if action > 0:
            self.state.cooldown = self.COOLDOWN_STEPS

        self.state.alert_history.append(action)

        # Additional spam check: if we've already sent 3 alerts in last 7 days,
        # downgrade the action to avoid spam penalty
        recent_alerts = sum(1 for a in self.state.alert_history[-7:] if a > 0)
        if recent_alerts >= 3 and action > 0:
            action = max(0, action - 1)  # downgrade

        return action

    def _count_anomalous_signals(self, obs: Dict[str, Any]) -> int:
        """Count signals that deviate more than ANOMALY_THRESHOLD_Z from baseline."""
        count = 0
        for key, (mu, sigma) in CDR_BASELINES.items():
            if key not in obs:
                continue
            val = float(obs[key])
            z   = (val - mu) / max(sigma, 1e-6)
            if key in HIGH_IS_BAD and z > self.ANOMALY_THRESHOLD_Z:
                count += 1
            elif key in LOW_IS_BAD and z < -self.ANOMALY_THRESHOLD_Z:
                count += 1
        return count

    def _count_trend_anomalies(self, obs: Dict[str, Any]) -> int:
        """Count trend slopes that indicate systematic worsening."""
        count = 0
        trend_typing  = obs.get("trend_typing_delay", 0.0)
        trend_sleep   = obs.get("trend_sleep",        0.0)
        trend_routine = obs.get("trend_routine",      0.0)

        if trend_typing  >  self.TREND_THRESHOLD: count += 1  # worsening typing
        if trend_sleep   < -self.TREND_THRESHOLD: count += 1  # worsening sleep
        if trend_routine < -self.TREND_THRESHOLD: count += 1  # worsening routine
        return count


# ── Trivial Baseline Agents ───────────────────────────────────────────────────

class RandomAgent:
    """Selects a uniform random action at every step."""
    def reset(self): pass
    def act(self, obs: Dict[str, Any]) -> int:
        return int(np.random.randint(0, 4))


class AlwaysAlertAgent:
    """Always escalates — tests upper bound of alert-fatigue penalty."""
    def reset(self): pass
    def act(self, obs: Dict[str, Any]) -> int:
        return 3


# ── Runner ────────────────────────────────────────────────────────────────────

def run_task1(agent, n_seeds: int = 10) -> Tuple[float, List[dict]]:
    from task1_easy import run_episode, grade, build_env

    def agent_fn(obs):
        return agent.act(obs)

    results = []
    for seed in range(n_seeds):
        agent.reset()
        env    = build_env(seed)
        result = run_episode(env, agent_fn=agent_fn)
        results.append(result)
        print(f"  seed={seed:02d}  true={result['true_stage']}  pred={result['predicted_stage']}  score={result['score']:.2f}")

    score = grade(results)
    return score, results


def run_task2(agent, n_seeds: int = 10) -> Tuple[float, List[dict]]:
    from task2_medium import run_episode, grade, build_env

    def agent_fn(obs, step):
        return agent.act(obs)

    results = []
    for seed in range(n_seeds):
        agent.reset()
        env    = build_env(seed)
        result = run_episode(env, agent_fn=agent_fn)
        results.append(result)
        print(f"  seed={seed:02d}  anomaly_day={result['anomaly_day']}  first_alert={result['first_alert_day']}  score={result['score']:.2f}")

    score = grade(results)
    return score, results


def run_task3(agent, n_seeds: int = 5) -> Tuple[float, List[dict]]:
    from task3_hard import run_episode, grade, build_env

    def agent_fn(obs, step, history):
        return agent.act(obs)

    results = []
    for seed in range(n_seeds):
        agent.reset()
        env    = build_env(seed)
        result = run_episode(env, agent_fn=agent_fn)
        results.append(result)
        print(f"  seed={seed:02d}  tp={result['tp']}  fp={result['fp']}  fn={result['fn']}  f1={result['f1_score']:.2f}  score={result['score']:.2f}")

    score = grade(results)
    return score, results


def main():
    parser = argparse.ArgumentParser(description="CogTraceEnv benchmark runner")
    parser.add_argument(
        "--agent",
        choices=["oracle", "random", "always_alert"],
        default="oracle",
        help="Which agent to benchmark (default: oracle)",
    )
    parser.add_argument("--seeds",  type=int, default=10,   help="Seeds per task")
    parser.add_argument("--output", type=str, default=None, help="Save results JSON to this path")
    args = parser.parse_args()

    agent_map = {
        "oracle":       RuleBasedOracle(),
        "random":       RandomAgent(),
        "always_alert": AlwaysAlertAgent(),
    }
    agent = agent_map[args.agent]

    print("=" * 60)
    print(f"CogTraceEnv — {args.agent.upper()} Agent Benchmark")
    print("=" * 60)

    all_results = {}

    for task_id, runner, label in [
        ("task1_easy",   run_task1, "Task 1 (Easy)   — Stage Classification"),
        ("task2_medium", run_task2, "Task 2 (Medium) — Anomaly Timing"),
        ("task3_hard",   run_task3, "Task 3 (Hard)   — Full Triage"),
    ]:
        print(f"\n▶ {label}")
        t0 = time.time()
        try:
            n = args.seeds if task_id != "task3_hard" else max(5, args.seeds // 2)
            score, results = runner(agent, n_seeds=n)
            elapsed = time.time() - t0
            all_results[task_id] = {"score": score, "results": results}
            print(f"  ✓ SCORE: {score:.4f}  ({elapsed:.1f}s)")
        except Exception as e:
            print(f"  ✗ FAILED: {e}")
            import traceback; traceback.print_exc()
            all_results[task_id] = {"score": None}

    print("\n" + "=" * 60)
    print("FINAL SCORES")
    print("=" * 60)

    valid_scores = []
    for task_id, data in all_results.items():
        s   = data.get("score")
        val = f"{s:.4f}" if s is not None else "ERROR"
        print(f"  {task_id:<20}  {val}")
        if s is not None:
            valid_scores.append(s)

    if valid_scores:
        print(f"\n  Mean score : {sum(valid_scores) / len(valid_scores):.4f}")
    print("=" * 60)

    if args.output:
        os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True)
        with open(args.output, "w") as f:
            json.dump({
                "agent":  args.agent,
                "scores": {k: v.get("score") for k, v in all_results.items()},
                "mean":   sum(valid_scores) / len(valid_scores) if valid_scores else None,
            }, f, indent=2)
        print(f"\nResults saved to {args.output}")


if __name__ == "__main__":
    main()