from __future__ import annotations

from dataclasses import dataclass
from typing import Dict, Mapping, Sequence

import numpy as np

from .dp import policy_occupancy
from .environment import ChanceState, History, LifeGridWorld


@dataclass
class TabularEstimate:
    action_counts: Dict[History, np.ndarray]
    event_counts: Dict[ChanceState, np.ndarray]
    pi_hat: Dict[History, np.ndarray]
    mu_hat: Dict[ChanceState, np.ndarray]
    action_rewards: Dict[tuple[History, int], float]
    event_rewards: Dict[tuple[History, int, int], float]


def sample_tabular(
    world: LifeGridWorld,
    n_trajectories: int,
    seed: int,
    alpha: float = 0.5,
    reward_noise: float = 0.10,
) -> TabularEstimate:
    rng = np.random.default_rng(seed)
    states = list(world.decision_states(include_terminal=False))
    action_counts = {s: np.zeros(world.n_actions, dtype=np.int64) for s in states}
    event_counts = {
        (s, a): np.zeros(world.n_events, dtype=np.int64)
        for s in states
        for a in range(world.n_actions)
    }
    action_sum = {(s, a): 0.0 for s in states for a in range(world.n_actions)}
    event_sum = {
        (s, a, e): 0.0
        for s in states
        for a in range(world.n_actions)
        for e in range(world.n_events)
    }

    for _ in range(n_trajectories):
        state: History = ()
        for _depth in range(world.horizon):
            pi = world.behavior_policy(state)
            action = int(rng.choice(world.n_actions, p=pi))
            mu = world.event_kernel(state, action)
            event = int(rng.choice(world.n_events, p=mu))
            action_counts[state][action] += 1
            event_counts[(state, action)][event] += 1
            action_sum[(state, action)] += world.action_reward(state, action) + rng.normal(
                0.0, reward_noise
            )
            event_sum[(state, action, event)] += world.event_reward(
                state, action, event
            ) + rng.normal(0.0, reward_noise)
            state = world.next_state(state, action, event)

    pi_hat = {}
    mu_hat = {}
    action_rewards = {}
    event_rewards = {}
    for state in states:
        counts = action_counts[state].astype(np.float64)
        pi_hat[state] = (counts + alpha) / (counts.sum() + alpha * world.n_actions)
        for action in range(world.n_actions):
            n_action = int(action_counts[state][action])
            action_rewards[(state, action)] = (
                action_sum[(state, action)] / n_action if n_action else 0.0
            )
            ec = event_counts[(state, action)].astype(np.float64)
            mu_hat[(state, action)] = (ec + alpha) / (
                ec.sum() + alpha * world.n_events
            )
            for event in range(world.n_events):
                n_event = int(event_counts[(state, action)][event])
                event_rewards[(state, action, event)] = (
                    event_sum[(state, action, event)] / n_event if n_event else 0.0
                )

    return TabularEstimate(
        action_counts=action_counts,
        event_counts=event_counts,
        pi_hat=pi_hat,
        mu_hat=mu_hat,
        action_rewards=action_rewards,
        event_rewards=event_rewards,
    )


def exact_trajectory_kl(world: LifeGridWorld, estimate: TabularEstimate) -> float:
    kl = 0.0
    for trajectory in world.trajectories():
        state: History = ()
        p = 1.0
        q = 1.0
        for action, event in trajectory:
            pi = world.behavior_policy(state)
            mu = world.event_kernel(state, action)
            p *= pi[action] * mu[event]
            q *= estimate.pi_hat[state][action] * estimate.mu_hat[(state, action)][event]
            state = world.next_state(state, action, event)
        if p > 0:
            kl += p * np.log(p / q)
    return float(kl)


def transition_metrics(
    world: LifeGridWorld,
    estimate: TabularEstimate,
    threshold: int,
) -> tuple[float, float, float]:
    behavior = {
        state: world.behavior_policy(state)
        for state in world.decision_states(include_terminal=False)
    }
    occupancy = policy_occupancy(world, behavior)
    weighted_sum = 0.0
    total_mass = 0.0
    supported_tvs = []
    supported_mass = 0.0
    for state in world.decision_states(include_terminal=False):
        for action in range(world.n_actions):
            mass = occupancy.get(state, 0.0) * behavior[state][action]
            true_mu = world.event_kernel(state, action)
            tv = 0.5 * float(np.abs(true_mu - estimate.mu_hat[(state, action)]).sum())
            weighted_sum += mass * tv
            total_mass += mass
            if estimate.action_counts[state][action] >= threshold:
                supported_tvs.append(tv)
                supported_mass += mass
    weighted = weighted_sum / total_mass
    max_supported = max(supported_tvs, default=float("nan"))
    coverage = supported_mass / total_mass
    return float(weighted), float(max_supported), float(coverage)


def build_support_gate(
    world: LifeGridWorld,
    estimate: TabularEstimate,
    threshold: int,
) -> Dict[History, Sequence[int]]:
    gate: Dict[History, Sequence[int]] = {}
    for state in world.decision_states(include_terminal=False):
        counts = estimate.action_counts[state]
        allowed = np.flatnonzero(counts >= threshold).tolist()
        if not allowed:
            allowed = [int(np.argmax(counts))]
        gate[state] = allowed
    return gate


def reachable_pairs_under_gate(
    world: LifeGridWorld, gate: Mapping[History, Sequence[int]]
) -> tuple[set[History], set[ChanceState]]:
    states: set[History] = {world.start}
    chance: set[ChanceState] = set()
    frontier = {world.start}
    for _depth in range(world.horizon):
        nxt: set[History] = set()
        for state in frontier:
            for action in gate[state]:
                chance.add((state, int(action)))
                for event in range(world.n_events):
                    child = world.next_state(state, int(action), event)
                    states.add(child)
                    nxt.add(child)
        frontier = nxt
    return states, chance


def model_errors_on_gate(
    world: LifeGridWorld,
    estimate: TabularEstimate,
    gate: Mapping[History, Sequence[int]],
) -> tuple[float, float, float]:
    states, chance = reachable_pairs_under_gate(world, gate)
    epsilon_r = 0.0
    epsilon_mu = 0.0
    r_bound = world.reward_bound
    for state, action in chance:
        epsilon_r = max(
            epsilon_r,
            abs(estimate.action_rewards[(state, action)] - world.action_reward(state, action)),
        )
        r_bound = max(r_bound, abs(estimate.action_rewards[(state, action)]))
        true_mu = world.event_kernel(state, action)
        epsilon_mu = max(
            epsilon_mu,
            0.5 * float(np.abs(true_mu - estimate.mu_hat[(state, action)]).sum()),
        )
        for event in range(world.n_events):
            epsilon_r = max(
                epsilon_r,
                abs(
                    estimate.event_rewards[(state, action, event)]
                    - world.event_reward(state, action, event)
                ),
            )
            r_bound = max(r_bound, abs(estimate.event_rewards[(state, action, event)]))
    return float(epsilon_r), float(epsilon_mu), float(r_bound)


def theorem_e_constant(
    world: LifeGridWorld,
    epsilon_r: float,
    epsilon_mu: float,
    reward_bound: float,
) -> float:
    bref = float(np.log(1.0 / np.min(world.pi_ref)))
    transition_sum = sum(
        2.0 * reward_bound + (j - 1) * (4.0 * reward_bound + world.tau * bref)
        for j in range(1, world.horizon + 1)
    )
    return float(2.0 * world.horizon * epsilon_r + epsilon_mu * transition_sum)

