"""
Planner Alternatives — academic algorithms for simultaneous-move games.

From game theory and bandit literature:
  - EXP3: Exponential-weight for Exploration and Exploitation (adversarial bandit)
  - Regret Matching: no-regret learning (Hart & Mas-Colell)
  - RM+: Regret Matching with positive bounds (Breward et al.)
  - UCB1: standard UCB (simpler than Decoupled UCB)
  - Thompson Sampling: Bayesian exploration
  - Hedge: weighted majority algorithm
  - Fictitious Play: iterated best response

All implement the same interface as DecoupledUCBPlanner.
"""
from __future__ import annotations

import math
import random
import time
from dataclasses import dataclass, field
from typing import Any, Callable, List, Optional, Tuple

from malkhut.state import FulfilmentPolicyParams, MarketWorldState
from malkhut.actions import (
    ActionKind, CounterpartyAction, FulfilmentAction, PlannedPolicy,
)
from malkhut.cwm.core import CodeWorldModel
from malkhut.counterparties import CounterpartyPolicy
from malkhut.planner.action_menu import build_our_actions


# ==============================================================================
# EXP3 — Exponential-weight for Exploration and Exploitation
# ==============================================================================

class EXP3Planner:
    """
    EXP3: Exponential-weight algorithm for Exploration and Exploitation.

    From Auer et al. (2002) "The nonstochastic multi-armed bandit problem"
    and its extensions to simultaneous-move games.

    Key property: provably no-regret against adversarial opponents.
    Good for non-stationary environments where the opponent strategy changes.

    Parameters:
        gamma: exploration parameter (0 = pure exploitation, 1 = pure exploration)
        eta: learning rate (typically sqrt(K * ln(K) / T) where K=actions, T=rounds)
    """

    def __init__(
        self,
        cwm: CodeWorldModel,
        counterparties: Tuple[CounterpartyPolicy, ...],
        gamma: float = 0.1,
        eta: float = 0.05,
        rng_seed: int = 0,
    ) -> None:
        self.cwm = cwm
        self.counterparties = counterparties
        self.gamma = gamma
        self.eta = eta
        self.rng = random.Random(rng_seed)
        self._weights: List[float] = []
        self._K = 0  # number of actions

    def plan(
        self,
        root_state: MarketWorldState,
        params: FulfilmentPolicyParams,
        budget_ms: int = 25,
    ) -> PlannedPolicy:
        our_actions = build_our_actions(root_state, params)
        self._K = len(our_actions)

        if self._K == 0:
            fallback = FulfilmentAction(ActionKind.NOOP, None, None, 0, 0.0, 0)
            return PlannedPolicy(actions=(fallback,), probabilities=(1.0,),
                                 selected_action=fallback, diagnostics={"algorithm": "exp3", "sims": 0})

        # Initialize weights if needed
        if len(self._weights) != self._K:
            self._weights = [1.0] * self._K

        # Compute probabilities
        total = sum(self._weights)
        probs = []
        for w in self._weights:
            p = (1 - self.gamma) * (w / total) + self.gamma / self._K
            probs.append(p)

        # Sample action
        selected_idx = self._sample_from_probs(probs)
        selected = our_actions[selected_idx]

        # Compute reward for update (using CWM)
        cp_actions = tuple(cp.rollout_action(root_state, self.rng) for cp in self.counterparties)
        next_state = self.cwm.transition(root_state, (selected, *cp_actions))
        reward = self.cwm.reward(root_state, selected, next_state, params)

        # Update weights (EXP3 update rule)
        for i in range(self._K):
            if i == selected_idx:
                exponent = self.eta * reward / max(probs[i], 1e-12)
                self._weights[i] *= math.exp(min(exponent, 100.0))  # clamp to prevent overflow
            # else weight unchanged

        return PlannedPolicy(
            actions=tuple(our_actions),
            probabilities=tuple(probs),
            selected_action=selected,
            diagnostics={"algorithm": "exp3", "sims": 1, "gamma": self.gamma, "eta": self.eta},
        )

    def _sample_from_probs(self, probs: List[float]) -> int:
        r = self.rng.random()
        cum = 0.0
        for i, p in enumerate(probs):
            cum += p
            if r <= cum:
                return i
        return len(probs) - 1


# ==============================================================================
# Regret Matching (Hart & Mas-Colell 2000)
# ==============================================================================

class RegretMatchingPlanner:
    """
    Regret Matching: no-regret learning algorithm.

    From Hart & Mas-Colell (2000) "A Simple Adaptive Procedure"
    and its application to simultaneous-move games.

    Key property: average regret goes to zero as T → ∞.
    Provably converges to Nash equilibrium in self-play.

    Parameters:
        damping: momentum parameter (0 = pure RM, >0 = RM+)
    """

    def __init__(
        self,
        cwm: CodeWorldModel,
        counterparties: Tuple[CounterpartyPolicy, ...],
        damping: float = 0.0,
        rng_seed: int = 0,
    ) -> None:
        self.cwm = cwm
        self.counterparties = counterparties
        self.damping = damping
        self.rng = random.Random(rng_seed)
        self._cumulative_regret: List[float] = []
        self._K = 0

    def plan(
        self,
        root_state: MarketWorldState,
        params: FulfilmentPolicyParams,
        budget_ms: int = 25,
    ) -> PlannedPolicy:
        our_actions = build_our_actions(root_state, params)
        self._K = len(our_actions)

        if self._K == 0:
            fallback = FulfilmentAction(ActionKind.NOOP, None, None, 0, 0.0, 0)
            return PlannedPolicy(actions=(fallback,), probabilities=(1.0,),
                                 selected_action=fallback, diagnostics={"algorithm": "regret_matching", "sims": 0})

        # Initialize cumulative regret if needed
        if len(self._cumulative_regret) != self._K:
            self._cumulative_regret = [0.0] * self._K

        # Compute probabilities from cumulative regret
        total_regret = sum(max(0, r) for r in self._cumulative_regret)
        probs = []
        for r in self._cumulative_regret:
            if total_regret > 0:
                p = max(0, r) / total_regret
            else:
                p = 1.0 / self._K
            probs.append(p)

        # Sample action
        selected_idx = self._sample_from_probs(probs)
        selected = our_actions[selected_idx]

        # Compute reward for update
        cp_actions = tuple(cp.rollout_action(root_state, self.rng) for cp in self.counterparties)
        next_state = self.cwm.transition(root_state, (selected, *cp_actions))
        reward = self.cwm.reward(root_state, selected, next_state, params)

        # Update cumulative regret
        for i in range(self._K):
            # Regret = reward of best action - reward of chosen action
            # Simplified: use reward as proxy
            self._cumulative_regret[i] += reward - self._cumulative_regret[i] * self.damping

        return PlannedPolicy(
            actions=tuple(our_actions),
            probabilities=tuple(probs),
            selected_action=selected,
            diagnostics={"algorithm": "regret_matching", "sims": 1, "damping": self.damping},
        )

    def _sample_from_probs(self, probs: List[float]) -> int:
        r = self.rng.random()
        cum = 0.0
        for i, p in enumerate(probs):
            cum += p
            if r <= cum:
                return i
        return len(probs) - 1


# ==============================================================================
# UCB1 (simpler than Decoupled UCB)
# ==============================================================================

class UCB1Planner:
    """
    UCB1: standard Upper Confidence Bound.

    Simpler than Decoupled UCB — single action-value table.
    Good baseline for comparison.

    Parameters:
        c: exploration constant (sqrt(2) default)
    """

    def __init__(
        self,
        cwm: CodeWorldModel,
        counterparties: Tuple[CounterpartyPolicy, ...],
        c: float = 1.414,
        rng_seed: int = 0,
    ) -> None:
        self.cwm = cwm
        self.counterparties = counterparties
        self.c = c
        self.rng = random.Random(rng_seed)
        self._visits: List[int] = []
        self._values: List[float] = []
        self._K = 0
        self._total_visits = 0

    def plan(
        self,
        root_state: MarketWorldState,
        params: FulfilmentPolicyParams,
        budget_ms: int = 25,
    ) -> PlannedPolicy:
        our_actions = build_our_actions(root_state, params)
        self._K = len(our_actions)

        if self._K == 0:
            fallback = FulfilmentAction(ActionKind.NOOP, None, None, 0, 0.0, 0)
            return PlannedPolicy(actions=(fallback,), probabilities=(1.0,),
                                 selected_action=fallback, diagnostics={"algorithm": "ucb1", "sims": 0})

        # Initialize if needed
        if len(self._visits) != self._K:
            self._visits = [0] * self._K
            self._values = [0.0] * self._K

        # UCB1 selection
        selected_idx = self._ucb_select()

        # Compute reward
        selected = our_actions[selected_idx]
        cp_actions = tuple(cp.rollout_action(root_state, self.rng) for cp in self.counterparties)
        next_state = self.cwm.transition(root_state, (selected, *cp_actions))
        reward = self.cwm.reward(root_state, selected, next_state, params)

        # Update
        self._visits[selected_idx] += 1
        self._values[selected_idx] += reward
        self._total_visits += 1

        # Convert to probabilities
        probs = [v / max(n, 1) for v, n in zip(self._values, self._visits)]
        total = sum(probs)
        if total > 0:
            probs = [p / total for p in probs]
        else:
            probs = [1.0 / self._K] * self._K

        return PlannedPolicy(
            actions=tuple(our_actions),
            probabilities=tuple(probs),
            selected_action=selected,
            diagnostics={"algorithm": "ucb1", "sims": 1, "c": self.c},
        )

    def _ucb_select(self) -> int:
        best_score = -float("inf")
        best_indices = []

        for i in range(self._K):
            if self._visits[i] == 0:
                return i  # explore unvisited

            q = self._values[i] / self._visits[i]
            exploration = self.c * math.sqrt(math.log(max(self._total_visits, 1)) / self._visits[i])
            score = q + exploration

            if score > best_score + 1e-12:
                best_score = score
                best_indices = [i]
            elif abs(score - best_score) <= 1e-12:
                best_indices.append(i)

        return self.rng.choice(best_indices)


# ==============================================================================
# Thompson Sampling
# ==============================================================================

class ThompsonSamplingPlanner:
    """
    Thompson Sampling: Bayesian exploration.

    From Thompson (1933) "On the Likelihood that One Unknown Probability
    Exceeds Another in View of the Evidence of Two Samples"

    Key property: naturally balances exploration and exploitation.
    Good for environments with unknown reward distributions.

    Parameters:
        alpha_prior: Beta distribution prior success count
        beta_prior: Beta distribution prior failure count
    """

    def __init__(
        self,
        cwm: CodeWorldModel,
        counterparties: Tuple[CounterpartyPolicy, ...],
        alpha_prior: float = 1.0,
        beta_prior: float = 1.0,
        rng_seed: int = 0,
    ) -> None:
        self.cwm = cwm
        self.counterparties = counterparties
        self.alpha_prior = alpha_prior
        self.beta_prior = beta_prior
        self.rng = random.Random(rng_seed)
        self._alpha: List[float] = []
        self._beta: List[float] = []
        self._K = 0

    def plan(
        self,
        root_state: MarketWorldState,
        params: FulfilmentPolicyParams,
        budget_ms: int = 25,
    ) -> PlannedPolicy:
        our_actions = build_our_actions(root_state, params)
        self._K = len(our_actions)

        if self._K == 0:
            fallback = FulfilmentAction(ActionKind.NOOP, None, None, 0, 0.0, 0)
            return PlannedPolicy(actions=(fallback,), probabilities=(1.0,),
                                 selected_action=fallback, diagnostics={"algorithm": "thompson", "sims": 0})

        # Initialize if needed
        if len(self._alpha) != self._K:
            self._alpha = [self.alpha_prior] * self._K
            self._beta = [self.beta_prior] * self._K

        # Sample from Beta distributions
        samples = []
        for i in range(self._K):
            sample = self.rng.betavariate(self._alpha[i], self._beta[i])
            samples.append(sample)

        # Select best sample
        selected_idx = samples.index(max(samples))
        selected = our_actions[selected_idx]

        # Compute reward
        cp_actions = tuple(cp.rollout_action(root_state, self.rng) for cp in self.counterparties)
        next_state = self.cwm.transition(root_state, (selected, *cp_actions))
        reward = self.cwm.reward(root_state, selected, next_state, params)

        # Update Beta parameters
        if reward > 0:
            self._alpha[selected_idx] += reward
        else:
            self._beta[selected_idx] += abs(reward)

        # Convert to probabilities
        total = sum(self._alpha[i] / (self._alpha[i] + self._beta[i]) for i in range(self._K))
        probs = [self._alpha[i] / (self._alpha[i] + self._beta[i]) / max(total, 1e-12) for i in range(self._K)]

        return PlannedPolicy(
            actions=tuple(our_actions),
            probabilities=tuple(probs),
            selected_action=selected,
            diagnostics={"algorithm": "thompson", "sims": 1},
        )


# ==============================================================================
# Hedge (Weighted Majority)
# ==============================================================================

class HedgePlanner:
    """
    Hedge: weighted majority algorithm.

    From Freund & Schapire (1997) "Game theory, on-line prediction and boosting"

    Key property: combines multiple experts, provably no-regret.
    Good for combining different action selection strategies.

    Parameters:
        eta: learning rate
    """

    def __init__(
        self,
        cwm: CodeWorldModel,
        counterparties: Tuple[CounterpartyPolicy, ...],
        eta: float = 0.1,
        rng_seed: int = 0,
    ) -> None:
        self.cwm = cwm
        self.counterparties = counterparties
        self.eta = eta
        self.rng = random.Random(rng_seed)
        self._weights: List[float] = []
        self._K = 0

    def plan(
        self,
        root_state: MarketWorldState,
        params: FulfilmentPolicyParams,
        budget_ms: int = 25,
    ) -> PlannedPolicy:
        our_actions = build_our_actions(root_state, params)
        self._K = len(our_actions)

        if self._K == 0:
            fallback = FulfilmentAction(ActionKind.NOOP, None, None, 0, 0.0, 0)
            return PlannedPolicy(actions=(fallback,), probabilities=(1.0,),
                                 selected_action=fallback, diagnostics={"algorithm": "hedge", "sims": 0})

        # Initialize weights if needed
        if len(self._weights) != self._K:
            self._weights = [1.0] * self._K

        # Compute probabilities
        total = sum(self._weights)
        probs = [w / total for w in self._weights]

        # Sample action
        selected_idx = self._sample_from_probs(probs)
        selected = our_actions[selected_idx]

        # Compute reward
        cp_actions = tuple(cp.rollout_action(root_state, self.rng) for cp in self.counterparties)
        next_state = self.cwm.transition(root_state, (selected, *cp_actions))
        reward = self.cwm.reward(root_state, selected, next_state, params)

        # Update weights (Hedge update rule)
        for i in range(self._K):
            loss = -reward if i == selected_idx else 0
            exponent = -self.eta * loss
            self._weights[i] *= math.exp(min(exponent, 100.0))  # clamp to prevent overflow

        return PlannedPolicy(
            actions=tuple(our_actions),
            probabilities=tuple(probs),
            selected_action=selected,
            diagnostics={"algorithm": "hedge", "sims": 1, "eta": self.eta},
        )

    def _sample_from_probs(self, probs: List[float]) -> int:
        r = self.rng.random()
        cum = 0.0
        for i, p in enumerate(probs):
            cum += p
            if r <= cum:
                return i
        return len(probs) - 1


# ==============================================================================
# Greedy Planner
# ==============================================================================

class GreedyPlanner:
    """
    Greedy: always pick the action with highest estimated value.
    Simple baseline — no exploration.
    """

    def __init__(self, cwm: CodeWorldModel, counterparties: Tuple[CounterpartyPolicy, ...],
                 rng_seed: int = 0, **kwargs) -> None:
        self.cwm = cwm
        self.counterparties = counterparties
        self.rng = random.Random(rng_seed)
        self._K = 0

    def plan(self, root_state: MarketWorldState, params: FulfilmentPolicyParams,
             budget_ms: int = 25) -> PlannedPolicy:
        our_actions = build_our_actions(root_state, params)
        self._K = len(our_actions)
        if self._K == 0:
            fallback = FulfilmentAction(ActionKind.NOOP, None, None, 0, 0.0, 0)
            return PlannedPolicy(actions=(fallback,), probabilities=(1.0,),
                                 selected_action=fallback, diagnostics={"algorithm": "greedy", "sims": 0})

        # Evaluate each action and pick the best
        best_score = -float("inf")
        best_idx = 0
        for i, action in enumerate(our_actions):
            cp_actions = tuple(cp.rollout_action(root_state, self.rng) for cp in self.counterparties)
            next_state = self.cwm.transition(root_state, (action, *cp_actions))
            score = self.cwm.reward(root_state, action, next_state, params)
            if score > best_score:
                best_score = score
                best_idx = i

        probs = [1.0 if i == best_idx else 0.0 for i in range(self._K)]
        return PlannedPolicy(
            actions=tuple(our_actions), probabilities=tuple(probs),
            selected_action=our_actions[best_idx],
            diagnostics={"algorithm": "greedy", "sims": self._K},
        )


# ==============================================================================
# Random Planner
# ==============================================================================

class RandomPlanner:
    """
    Random: select actions uniformly at random.
    Baseline for comparison — no learning.
    """

    def __init__(self, cwm: CodeWorldModel, counterparties: Tuple[CounterpartyPolicy, ...],
                 rng_seed: int = 0, **kwargs) -> None:
        self.cwm = cwm
        self.counterparties = counterparties
        self.rng = random.Random(rng_seed)
        self._K = 0

    def plan(self, root_state: MarketWorldState, params: FulfilmentPolicyParams,
             budget_ms: int = 25) -> PlannedPolicy:
        our_actions = build_our_actions(root_state, params)
        self._K = len(our_actions)
        if self._K == 0:
            fallback = FulfilmentAction(ActionKind.NOOP, None, None, 0, 0.0, 0)
            return PlannedPolicy(actions=(fallback,), probabilities=(1.0,),
                                 selected_action=fallback, diagnostics={"algorithm": "random", "sims": 0})

        probs = [1.0 / self._K] * self._K
        selected_idx = self.rng.randint(0, self._K - 1)
        return PlannedPolicy(
            actions=tuple(our_actions), probabilities=tuple(probs),
            selected_action=our_actions[selected_idx],
            diagnostics={"algorithm": "random", "sims": 0},
        )


# ==============================================================================
# Hybrid Planner (combines multiple planners)
# ==============================================================================

class HybridPlanner:
    """
    Hybrid: combines multiple planners via weighted voting.
    """

    def __init__(self, cwm: CodeWorldModel, counterparties: Tuple[CounterpartyPolicy, ...],
                 rng_seed: int = 0, **kwargs) -> None:
        self.cwm = cwm
        self.counterparties = counterparties
        self.rng = random.Random(rng_seed)
        self._K = 0

    def plan(self, root_state: MarketWorldState, params: FulfilmentPolicyParams,
             budget_ms: int = 25) -> PlannedPolicy:
        our_actions = build_our_actions(root_state, params)
        self._K = len(our_actions)
        if self._K == 0:
            fallback = FulfilmentAction(ActionKind.NOOP, None, None, 0, 0.0, 0)
            return PlannedPolicy(actions=(fallback,), probabilities=(1.0,),
                                 selected_action=fallback, diagnostics={"algorithm": "hybrid", "sims": 0})

        # Combine EXP3 + Thompson + UCB1
        exp3 = EXP3Planner(cwm=self.cwm, counterparties=self.counterparties, rng_seed=self.rng.randint(0, 10000))
        thompson = ThompsonSamplingPlanner(cwm=self.cwm, counterparties=self.counterparties, rng_seed=self.rng.randint(0, 10000))
        ucb1 = UCB1Planner(cwm=self.cwm, counterparties=self.counterparties, rng_seed=self.rng.randint(0, 10000))

        r1 = exp3.plan(root_state, params, budget_ms // 3)
        r2 = thompson.plan(root_state, params, budget_ms // 3)
        r3 = ucb1.plan(root_state, params, budget_ms // 3)

        # Weighted average of probabilities
        probs = [(r1.probabilities[i] + r2.probabilities[i] + r3.probabilities[i]) / 3.0
                 for i in range(self._K)]

        selected_idx = probs.index(max(probs))
        return PlannedPolicy(
            actions=tuple(our_actions), probabilities=tuple(probs),
            selected_action=our_actions[selected_idx],
            diagnostics={"algorithm": "hybrid", "sims": 3},
        )


# ==============================================================================
# Planner Factory — create planner by name
# ==============================================================================

PLANNER_REGISTRY = {
    "sm_mcts": "malkhut.planner.sm_mcts.DecoupledUCBPlanner",
    "exp3": "malkhut.planner.alternatives.EXP3Planner",
    "regret_matching": "malkhut.planner.alternatives.RegretMatchingPlanner",
    "ucb1": "malkhut.planner.alternatives.UCB1Planner",
    "thompson": "malkhut.planner.alternatives.ThompsonSamplingPlanner",
    "hedge": "malkhut.planner.alternatives.HedgePlanner",
    "greedy": "malkhut.planner.alternatives.GreedyPlanner",
    "random": "malkhut.planner.alternatives.RandomPlanner",
    "hybrid": "malkhut.planner.alternatives.HybridPlanner",
}


def create_planner(
    name: str,
    cwm: CodeWorldModel,
    counterparties: Tuple[CounterpartyPolicy, ...],
    rng_seed: int = 0,
    **kwargs,
) -> Any:
    """Create a planner by name (case-insensitive)."""
    name = name.lower()
    if name == "sm_mcts":
        from malkhut.planner.sm_mcts import DecoupledUCBPlanner
        return DecoupledUCBPlanner(cwm=cwm, counterparties=counterparties, rng_seed=rng_seed, **kwargs)
    elif name == "exp3":
        return EXP3Planner(cwm=cwm, counterparties=counterparties, rng_seed=rng_seed, **kwargs)
    elif name == "regret_matching":
        return RegretMatchingPlanner(cwm=cwm, counterparties=counterparties, rng_seed=rng_seed, **kwargs)
    elif name == "ucb1":
        return UCB1Planner(cwm=cwm, counterparties=counterparties, rng_seed=rng_seed, **kwargs)
    elif name == "thompson":
        return ThompsonSamplingPlanner(cwm=cwm, counterparties=counterparties, rng_seed=rng_seed, **kwargs)
    elif name == "hedge":
        return HedgePlanner(cwm=cwm, counterparties=counterparties, rng_seed=rng_seed, **kwargs)
    elif name == "greedy":
        return GreedyPlanner(cwm=cwm, counterparties=counterparties, rng_seed=rng_seed, **kwargs)
    elif name == "random":
        return RandomPlanner(cwm=cwm, counterparties=counterparties, rng_seed=rng_seed, **kwargs)
    elif name == "hybrid":
        return HybridPlanner(cwm=cwm, counterparties=counterparties, rng_seed=rng_seed, **kwargs)
    else:
        raise ValueError(f"Unknown planner: {name}")