from __future__ import annotations

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

import numpy as np

from .environment import ChanceState, History, LifeGridWorld


def _logsumexp_weighted(values: np.ndarray, weights: np.ndarray, tau: float) -> float:
    values = np.asarray(values, dtype=np.float64)
    weights = np.asarray(weights, dtype=np.float64)
    positive = weights > 0
    if not np.any(positive):
        raise ValueError("weighted log-sum-exp requires positive total weight")
    z = values[positive] / tau + np.log(weights[positive])
    m = np.max(z)
    return float(tau * (m + np.log(np.exp(z - m).sum())))


@dataclass
class DPResult:
    kind: str
    V: Dict[History, float]
    W: Dict[ChanceState, float]
    Q: Dict[History, np.ndarray]
    policy: Dict[History, np.ndarray]
    event_posterior: Dict[ChanceState, np.ndarray]


def solve(
    world: LifeGridWorld,
    kind: str,
    event_kernels: Mapping[ChanceState, np.ndarray] | None = None,
    action_rewards: Mapping[tuple[History, int], float] | None = None,
    event_rewards: Mapping[tuple[History, int, int], float] | None = None,
    gate: Mapping[History, Sequence[int]] | None = None,
) -> DPResult:
    if kind not in {"control", "generative"}:
        raise ValueError("kind must be 'control' or 'generative'")

    V: Dict[History, float] = {}
    W: Dict[ChanceState, float] = {}
    Q: Dict[History, np.ndarray] = {}
    policy: Dict[History, np.ndarray] = {}
    event_posterior: Dict[ChanceState, np.ndarray] = {}

    terminal_states = [s for s in world.decision_states() if len(s) == world.horizon]
    for state in terminal_states:
        V[state] = 0.0

    for depth in range(world.horizon - 1, -1, -1):
        states = [s for s in world.decision_states(include_terminal=False) if len(s) == depth]
        for state in states:
            q = np.empty(world.n_actions, dtype=np.float64)
            for action in range(world.n_actions):
                mu = (
                    world.event_kernel(state, action)
                    if event_kernels is None
                    else np.asarray(event_kernels[(state, action)], dtype=np.float64)
                )
                z = np.empty(world.n_events, dtype=np.float64)
                for event in range(world.n_events):
                    child = world.next_state(state, action, event)
                    r_event = (
                        world.event_reward(state, action, event)
                        if event_rewards is None
                        else event_rewards[(state, action, event)]
                    )
                    z[event] = r_event + V[child]

                if kind == "control":
                    w = float(np.dot(mu, z))
                    posterior = mu.copy()
                else:
                    w = _logsumexp_weighted(z, mu, world.tau)
                    logits = np.log(mu) + z / world.tau
                    logits -= logits.max()
                    posterior = np.exp(logits)
                    posterior /= posterior.sum()
                W[(state, action)] = w
                event_posterior[(state, action)] = posterior
                r_action = (
                    world.action_reward(state, action)
                    if action_rewards is None
                    else action_rewards[(state, action)]
                )
                q[action] = r_action + w

            allowed = np.arange(world.n_actions) if gate is None else np.asarray(gate[state], dtype=int)
            weights = np.zeros(world.n_actions, dtype=np.float64)
            weights[allowed] = world.pi_ref[allowed]
            V[state] = _logsumexp_weighted(q, weights, world.tau)
            logits = np.full(world.n_actions, -np.inf, dtype=np.float64)
            logits[allowed] = np.log(world.pi_ref[allowed]) + q[allowed] / world.tau
            m = np.max(logits[allowed])
            p = np.zeros(world.n_actions, dtype=np.float64)
            p[allowed] = np.exp(logits[allowed] - m)
            p /= p.sum()
            Q[state] = q
            policy[state] = p

    return DPResult(kind=kind, V=V, W=W, Q=Q, policy=policy, event_posterior=event_posterior)


def evaluate_policy(
    world: LifeGridWorld,
    policy: Mapping[History, np.ndarray],
    event_kernels: Mapping[ChanceState, np.ndarray] | None = None,
    action_rewards: Mapping[tuple[History, int], float] | None = None,
    event_rewards: Mapping[tuple[History, int, int], float] | None = None,
) -> Dict[History, float]:
    value: Dict[History, float] = {
        state: 0.0 for state in world.decision_states() if len(state) == world.horizon
    }
    for depth in range(world.horizon - 1, -1, -1):
        states = [s for s in world.decision_states(include_terminal=False) if len(s) == depth]
        for state in states:
            pi = np.asarray(policy[state], dtype=np.float64)
            action_values = np.zeros(world.n_actions, dtype=np.float64)
            for action in range(world.n_actions):
                mu = (
                    world.event_kernel(state, action)
                    if event_kernels is None
                    else event_kernels[(state, action)]
                )
                continuation = 0.0
                for event in range(world.n_events):
                    child = world.next_state(state, action, event)
                    r_event = (
                        world.event_reward(state, action, event)
                        if event_rewards is None
                        else event_rewards[(state, action, event)]
                    )
                    continuation += mu[event] * (r_event + value[child])
                r_action = (
                    world.action_reward(state, action)
                    if action_rewards is None
                    else action_rewards[(state, action)]
                )
                action_values[action] = r_action + continuation

            positive = pi > 0
            kl = float(np.sum(pi[positive] * np.log(pi[positive] / world.pi_ref[positive])))
            value[state] = float(np.dot(pi, action_values) - world.tau * kl)
    return value


def policy_occupancy(
    world: LifeGridWorld,
    policy: Mapping[History, np.ndarray],
    event_kernels: Mapping[ChanceState, np.ndarray] | None = None,
) -> Dict[History, float]:
    occupancy: Dict[History, float] = {world.start: 1.0}
    for depth in range(world.horizon):
        states = [s for s in world.decision_states(include_terminal=False) if len(s) == depth]
        for state in states:
            mass = occupancy.get(state, 0.0)
            if mass == 0.0:
                continue
            for action in range(world.n_actions):
                mu = (
                    world.event_kernel(state, action)
                    if event_kernels is None
                    else event_kernels[(state, action)]
                )
                for event in range(world.n_events):
                    child = world.next_state(state, action, event)
                    occupancy[child] = occupancy.get(child, 0.0) + mass * policy[state][action] * mu[event]
    return occupancy

