from __future__ import annotations
import numpy as np
from dataclasses import dataclass, field
from typing import Optional, Dict, Any


@dataclass
class SessionState:
    """
    Tracks player state over a session.
    Models fatigue, tilt, and focus using Ornstein-Uhlenbeck processes.
    """
    fatigue: float = 0.0      # 0 = fresh, 1 = exhausted
    tilt: float = 0.0         # -1 = tilted, 0 = neutral, 1 = confident
    focus: float = 1.0        # 0 = distracted, 1 = locked in
    time_elapsed_ms: float = 0.0
    
    def to_dict(self) -> Dict[str, float]:
        return {
            "fatigue": self.fatigue,
            "tilt": self.tilt,
            "focus": self.focus,
            "time_elapsed_ms": self.time_elapsed_ms,
        }
    
    @classmethod
    def from_dict(cls, d: Dict[str, float]) -> SessionState:
        return cls(
            fatigue=d.get("fatigue", 0.0),
            tilt=d.get("tilt", 0.0),
            focus=d.get("focus", 1.0),
            time_elapsed_ms=d.get("time_elapsed_ms", 0.0),
        )


@dataclass
class SessionSimulator:
    """
    Simulates player state evolution over a gaming session.
    
    Uses Ornstein-Uhlenbeck processes for realistic temporal dynamics:
    - Fatigue: Slowly increases, slowly recovers
    - Tilt: Affected by wins/losses, mean-reverts
    - Focus: Affected by round importance
    """
    # OU process parameters
    fatigue_reversion: float = 0.001   # Slow recovery
    fatigue_drift: float = 0.0001      # Gradual increase
    fatigue_volatility: float = 0.001
    
    tilt_reversion: float = 0.01       # Faster emotional recovery
    tilt_volatility: float = 0.01
    
    focus_reversion: float = 0.005
    focus_volatility: float = 0.005
    
    # Mental resilience affects tilt response (0-100)
    mental_resilience: float = 50.0
    
    # Current state
    state: SessionState = field(default_factory=SessionState)
    
    # Random generator
    _rng: Optional[np.random.Generator] = field(default=None, repr=False)
    
    def __post_init__(self):
        if self._rng is None:
            self._rng = np.random.default_rng()
    
    @classmethod
    def from_skill(
        cls,
        mental_resilience: float = 50.0,
        seed: Optional[int] = None,
    ) -> SessionSimulator:
        """Create simulator with skill-based parameters."""
        return cls(
            mental_resilience=mental_resilience,
            _rng=np.random.default_rng(seed),
        )
    
    def update(self, event: str, dt_ms: float) -> None:
        """
        Update session state based on game event.
        
        Args:
            event: One of "round_win", "round_loss", "kill", "death", 
                   "clutch_situation", "idle"
            dt_ms: Time elapsed in milliseconds
        """
        self.state.time_elapsed_ms += dt_ms
        
        # Fatigue: OU process with drift (slowly increases)
        # dF = θ(0 - F)dt + drift*dt + σ*dW
        fatigue_noise = self._rng.normal(0, self.fatigue_volatility * np.sqrt(dt_ms))
        self.state.fatigue += (
            self.fatigue_reversion * (0 - self.state.fatigue) * dt_ms +
            self.fatigue_drift * dt_ms +
            fatigue_noise
        )
        self.state.fatigue = np.clip(self.state.fatigue, 0.0, 1.0)
        
        # Tilt: affected by outcomes
        tilt_impact = 0.0
        if event == "round_loss":
            # More impact if low mental resilience
            tilt_impact = -0.1 * (1.0 - self.mental_resilience / 100.0)
        elif event == "round_win":
            tilt_impact = 0.05
        elif event == "death":
            tilt_impact = -0.02 * (1.0 - self.mental_resilience / 100.0)
        elif event == "kill":
            tilt_impact = 0.01
        
        self.state.tilt += tilt_impact
        
        # Tilt mean reversion (OU process)
        tilt_noise = self._rng.normal(0, self.tilt_volatility * np.sqrt(dt_ms))
        self.state.tilt += (
            self.tilt_reversion * (0 - self.state.tilt) * dt_ms +
            tilt_noise
        )
        self.state.tilt = np.clip(self.state.tilt, -1.0, 1.0)
        
        # Focus: affected by round importance
        if event == "clutch_situation":
            # Focus increases in clutch (if mentally strong)
            focus_boost = 0.2 * (0.5 + self.mental_resilience / 200.0)
            self.state.focus = min(1.0, self.state.focus + focus_boost)
        
        # Focus mean reversion
        focus_noise = self._rng.normal(0, self.focus_volatility * np.sqrt(dt_ms))
        self.state.focus += (
            self.focus_reversion * (1.0 - self.state.focus) * dt_ms +
            focus_noise
        )
        self.state.focus = np.clip(self.state.focus, 0.0, 1.0)
    
    def get_modifiers(self) -> Dict[str, float]:
        """
        Get performance modifiers based on current state.
        
        Returns:
            Dict with multipliers for different stats
        """
        return {
            # Fatigue slows reactions and reduces accuracy
            "reaction_time_mult": 1.0 + self.state.fatigue * 0.2,  # Up to 20% slower
            "accuracy_mult": 1.0 - self.state.fatigue * 0.15,      # Up to 15% less accurate
            
            # Tilt affects consistency
            "consistency_mult": (
                1.0 - abs(self.state.tilt) * 0.3 if self.state.tilt < 0 
                else 1.0 + self.state.tilt * 0.1
            ),
            
            # Tilt affects game sense when negative
            "game_sense_mult": 1.0 - max(0, -self.state.tilt) * 0.2,
            
            # Focus improves reaction time
            "focus_reaction_mult": 1.0 - (self.state.focus - 0.5) * 0.1,
        }
    
    def apply_to_stats(self, base_stats: Dict[str, float]) -> Dict[str, float]:
        """
        Apply session modifiers to base stats.
        
        Args:
            base_stats: Dict with keys like "reaction_time", "accuracy", etc.
            
        Returns:
            Modified stats dict
        """
        mods = self.get_modifiers()
        modified = base_stats.copy()
        
        if "reaction_time" in modified:
            modified["reaction_time"] *= mods["reaction_time_mult"] * mods["focus_reaction_mult"]
        
        if "accuracy" in modified:
            modified["accuracy"] *= mods["accuracy_mult"]
        
        if "consistency" in modified:
            modified["consistency"] *= mods["consistency_mult"]
        
        if "game_sense" in modified:
            modified["game_sense"] *= mods["game_sense_mult"]
        
        return modified
    
    def reset(self) -> None:
        """Reset session state to fresh."""
        self.state = SessionState()


def simulate_round_sequence(
    n_rounds: int,
    win_probability: float = 0.5,
    seed: Optional[int] = None,
) -> list[str]:
    """
    Simulate a sequence of round outcomes.
    
    Args:
        n_rounds: Number of rounds
        win_probability: Base probability of winning each round
        seed: Random seed
        
    Returns:
        List of events ("round_win" or "round_loss")
    """
    rng = np.random.default_rng(seed)
    events = []
    
    for _ in range(n_rounds):
        if rng.random() < win_probability:
            events.append("round_win")
        else:
            events.append("round_loss")
    
    return events


def generate_session_trace(
    n_rounds: int = 24,
    mental_resilience: float = 50.0,
    win_probability: float = 0.5,
    round_duration_ms: float = 90000.0,  # 1.5 minutes per round
    seed: Optional[int] = None,
) -> list[Dict[str, Any]]:
    """
    Generate complete session trace with states at each round.
    
    Args:
        n_rounds: Number of rounds
        mental_resilience: Player's mental resilience (0-100)
        win_probability: Base win probability
        round_duration_ms: Average round duration
        seed: Random seed
        
    Returns:
        List of dicts with round number, event, state, and modifiers
    """
    rng = np.random.default_rng(seed)
    sim = SessionSimulator.from_skill(mental_resilience, seed)
    
    events = simulate_round_sequence(n_rounds, win_probability, seed)
    trace = []
    
    for round_num, event in enumerate(events):
        # Add some kills/deaths during round
        n_kills = rng.poisson(1.0)
        n_deaths = rng.poisson(0.5)
        
        for _ in range(n_kills):
            sim.update("kill", round_duration_ms / (n_kills + n_deaths + 1))
        for _ in range(n_deaths):
            sim.update("death", round_duration_ms / (n_kills + n_deaths + 1))
        
        # Round outcome
        sim.update(event, round_duration_ms / 3)
        
        trace.append({
            "round": round_num,
            "event": event,
            "state": sim.state.to_dict(),
            "modifiers": sim.get_modifiers(),
        })
    
    return trace