from __future__ import annotations

from dataclasses import dataclass
from itertools import product
from typing import Dict, Iterator, Mapping, Sequence, Tuple

import numpy as np

History = Tuple[Tuple[int, int], ...]
ChanceState = Tuple[History, int]
Trajectory = Tuple[Tuple[int, int], ...]


def _softmax(x: np.ndarray) -> np.ndarray:
    z = np.asarray(x, dtype=np.float64)
    z = z - np.max(z)
    p = np.exp(z)
    return p / p.sum()


@dataclass(frozen=True)
class LifeGridWorld:
    horizon: int
    n_actions: int
    n_events: int
    tau: float
    pi_ref: np.ndarray
    action_base: np.ndarray
    action_event_effect: np.ndarray
    event_base: np.ndarray
    event_effect: np.ndarray
    transition_logits: np.ndarray
    behavior_logits: np.ndarray
    transition_temperature: float
    dispersion_scale: float
    behavior_floor: float = 0.04

    @property
    def start(self) -> History:
        return ()

    @property
    def terminal_depth(self) -> int:
        return self.horizon

    @property
    def reward_bound(self) -> float:
        values = []
        for state in self.decision_states(include_terminal=False):
            for action in range(self.n_actions):
                values.append(abs(self.action_reward(state, action)))
                for event in range(self.n_events):
                    values.append(abs(self.event_reward(state, action, event)))
        return float(max(values, default=0.0))

    def previous_event_index(self, state: History) -> int:
        return self.n_events if not state else state[-1][1]

    def action_reward(self, state: History, action: int) -> float:
        depth = len(state)
        prev = self.previous_event_index(state)
        value = self.action_base[depth, action]
        value += self.dispersion_scale * self.action_event_effect[depth, prev, action]
        return float(value)

    def event_reward(self, state: History, action: int, event: int) -> float:
        depth = len(state)
        value = self.event_base[depth, action]
        value += self.dispersion_scale * self.event_effect[depth, action, event]
        return float(value)

    def event_kernel(self, state: History, action: int) -> np.ndarray:
        depth = len(state)
        prev = self.previous_event_index(state)
        return _softmax(self.transition_logits[depth, prev, action] / self.transition_temperature)

    def behavior_policy(self, state: History) -> np.ndarray:
        depth = len(state)
        prev = self.previous_event_index(state)
        p = _softmax(self.behavior_logits[depth, prev])
        floor = self.behavior_floor
        return (1.0 - floor) * p + floor * np.ones(self.n_actions) / self.n_actions

    def next_state(self, state: History, action: int, event: int) -> History:
        return state + ((int(action), int(event)),)

    def decision_states(self, include_terminal: bool = True) -> Iterator[History]:
        yield ()
        frontier = [()]
        for _depth in range(self.horizon):
            nxt = []
            for state in frontier:
                for action in range(self.n_actions):
                    for event in range(self.n_events):
                        child = self.next_state(state, action, event)
                        nxt.append(child)
                        if include_terminal or len(child) < self.horizon:
                            yield child
            frontier = nxt

    def chance_states(self) -> Iterator[ChanceState]:
        for state in self.decision_states(include_terminal=False):
            for action in range(self.n_actions):
                yield (state, action)

    def trajectories(self) -> Iterator[Trajectory]:
        pairs = tuple(product(range(self.n_actions), range(self.n_events)))
        yield from product(pairs, repeat=self.horizon)

    def trajectory_reward(self, trajectory: Trajectory) -> float:
        state: History = ()
        total = 0.0
        for action, event in trajectory:
            total += self.action_reward(state, action)
            total += self.event_reward(state, action, event)
            state = self.next_state(state, action, event)
        return float(total)

    def trajectory_probability(
        self,
        trajectory: Trajectory,
        policy: Mapping[History, np.ndarray] | None = None,
        event_kernels: Mapping[ChanceState, np.ndarray] | None = None,
    ) -> float:
        state: History = ()
        prob = 1.0
        for action, event in trajectory:
            pi = self.behavior_policy(state) if policy is None else policy[state]
            mu = self.event_kernel(state, action) if event_kernels is None else event_kernels[(state, action)]
            prob *= float(pi[action]) * float(mu[event])
            state = self.next_state(state, action, event)
        return float(prob)


def build_world(
    seed: int,
    entropy: str = "low",
    dispersion: str = "low",
    horizon: int = 4,
    n_actions: int = 4,
    n_events: int = 3,
    tau: float = 1.0,
) -> LifeGridWorld:
    """Build one cell of the prespecified 2 x 2 LifeGrid design.

    Reusing ``seed`` across cells reuses every base random table. Only the
    transition temperature and dispersion scale change.
    """

    entropy_temperatures = {"low": 0.35, "high": 2.50}
    dispersion_scales = {"low": 0.15, "high": 1.10}
    if entropy not in entropy_temperatures:
        raise ValueError(f"unknown entropy level: {entropy}")
    if dispersion not in dispersion_scales:
        raise ValueError(f"unknown dispersion level: {dispersion}")

    rng = np.random.default_rng(seed)
    prev_states = n_events + 1
    action_base = rng.uniform(-0.35, 0.35, size=(horizon, n_actions))
    action_event_effect = rng.normal(0.0, 0.25, size=(horizon, prev_states, n_actions))
    event_base = rng.uniform(-0.20, 0.20, size=(horizon, n_actions))
    event_effect = rng.normal(0.0, 0.80, size=(horizon, n_actions, n_events))
    event_effect -= event_effect.mean(axis=-1, keepdims=True)
    transition_logits = rng.normal(
        0.0, 1.0, size=(horizon, prev_states, n_actions, n_events)
    )
    behavior_logits = rng.normal(0.0, 0.75, size=(horizon, prev_states, n_actions))
    pi_ref = np.ones(n_actions, dtype=np.float64) / n_actions

    return LifeGridWorld(
        horizon=horizon,
        n_actions=n_actions,
        n_events=n_events,
        tau=float(tau),
        pi_ref=pi_ref,
        action_base=action_base,
        action_event_effect=action_event_effect,
        event_base=event_base,
        event_effect=event_effect,
        transition_logits=transition_logits,
        behavior_logits=behavior_logits,
        transition_temperature=entropy_temperatures[entropy],
        dispersion_scale=dispersion_scales[dispersion],
    )

