from __future__ import annotations

from typing import Dict, Mapping

import numpy as np

from .dp import DPResult, evaluate_policy, policy_occupancy
from .environment import History, LifeGridWorld, Trajectory


def entropic_gap(mu: np.ndarray, z: np.ndarray, tau: float) -> float:
    y = z / tau
    m = float(np.max(y))
    lse = tau * (m + np.log(np.dot(mu, np.exp(y - m))))
    return float(lse - np.dot(mu, z))


def max_gap_recursion_residual(
    world: LifeGridWorld, generative: DPResult, control: DPResult
) -> float:
    residuals = []
    for state in world.decision_states(include_terminal=False):
        delta_w = np.zeros(world.n_actions)
        for action in range(world.n_actions):
            mu = world.event_kernel(state, action)
            z_gen = np.array(
                [
                    world.event_reward(state, action, event)
                    + generative.V[world.next_state(state, action, event)]
                    for event in range(world.n_events)
                ]
            )
            propagated = sum(
                mu[event]
                * (
                    generative.V[world.next_state(state, action, event)]
                    - control.V[world.next_state(state, action, event)]
                )
                for event in range(world.n_events)
            )
            rhs = entropic_gap(mu, z_gen, world.tau) + propagated
            lhs = generative.W[(state, action)] - control.W[(state, action)]
            residuals.append(abs(lhs - rhs))
            delta_w[action] = lhs

        pi_ctrl = control.policy[state]
        positive = pi_ctrl > 0
        x = delta_w[positive] / world.tau + np.log(pi_ctrl[positive])
        m = np.max(x)
        rhs_v = world.tau * (m + np.log(np.exp(x - m).sum()))
        lhs_v = generative.V[state] - control.V[state]
        residuals.append(abs(lhs_v - rhs_v))
    return float(max(residuals, default=0.0))


def factorization_residual(world: LifeGridWorld, generative: DPResult) -> tuple[float, float]:
    total = 0.0
    max_residual = 0.0
    root_v = generative.V[world.start]
    for trajectory in world.trajectories():
        state: History = ()
        arm_prob = 1.0
        base_prob = 1.0
        for action, event in trajectory:
            arm_prob *= generative.policy[state][action]
            arm_prob *= generative.event_posterior[(state, action)][event]
            base_prob *= world.pi_ref[action] * world.event_kernel(state, action)[event]
            state = world.next_state(state, action, event)
        ebm_prob = base_prob * np.exp((world.trajectory_reward(trajectory) - root_v) / world.tau)
        max_residual = max(max_residual, abs(arm_prob - ebm_prob))
        total += arm_prob
    return float(max_residual), float(abs(total - 1.0))


def audit_reconstruction_residual(world: LifeGridWorld) -> tuple[float, float]:
    relative_residual = 0.0
    counting_residual = 0.0
    uniform_event = 1.0 / world.n_events
    for trajectory in world.trajectories():
        state: History = ()
        log_rho = 0.0
        log_ref = 0.0
        local_relative = 0.0
        local_counting = 0.0
        for action, event in trajectory:
            pi_beta = world.behavior_policy(state)
            mu = world.event_kernel(state, action)
            log_rho += np.log(pi_beta[action]) + np.log(mu[event])
            log_ref += np.log(world.pi_ref[action]) + np.log(mu[event])
            local_relative += world.tau * np.log(pi_beta[action] / world.pi_ref[action])
            local_relative += world.tau * np.log(mu[event] / mu[event])
            local_counting += world.tau * np.log(pi_beta[action] / (1.0 / world.n_actions))
            local_counting += world.tau * np.log(mu[event] / uniform_event)
            state = world.next_state(state, action, event)
        relative_residual = max(
            relative_residual, abs(local_relative - world.tau * (log_rho - log_ref))
        )
        counting_target = world.tau * (
            log_rho - world.horizon * np.log((1.0 / world.n_actions) * uniform_event)
        )
        counting_residual = max(counting_residual, abs(local_counting - counting_target))
    return float(relative_residual), float(counting_residual)


def trajectory_decomposition_residual(
    world: LifeGridWorld, control: DPResult
) -> tuple[float, float]:
    max_path = 0.0
    max_event_mean = 0.0
    for state in world.decision_states(include_terminal=False):
        for action in range(world.n_actions):
            mu = world.event_kernel(state, action)
            innovations = []
            for event in range(world.n_events):
                child = world.next_state(state, action, event)
                innovations.append(
                    world.event_reward(state, action, event)
                    + control.V[child]
                    - control.W[(state, action)]
                )
            max_event_mean = max(max_event_mean, abs(float(np.dot(mu, innovations))))

    for trajectory in world.trajectories():
        state: History = ()
        components = control.V[world.start]
        for action, event in trajectory:
            components += control.Q[state][action] - control.V[state]
            child = world.next_state(state, action, event)
            components += (
                world.event_reward(state, action, event)
                + control.V[child]
                - control.W[(state, action)]
            )
            state = child
        max_path = max(max_path, abs(world.trajectory_reward(trajectory) - components))
    return float(max_path), float(max_event_mean)


def kl_regret_identity_residual(
    world: LifeGridWorld,
    control: DPResult,
    policy: Mapping[History, np.ndarray],
) -> float:
    values = evaluate_policy(world, policy)
    occupancy = policy_occupancy(world, policy)
    total_kl = 0.0
    for state in world.decision_states(include_terminal=False):
        pi = np.asarray(policy[state])
        pi_star = control.policy[state]
        positive = pi > 0
        total_kl += occupancy.get(state, 0.0) * float(
            np.sum(pi[positive] * np.log(pi[positive] / pi_star[positive]))
        )
    lhs = control.V[world.start] - values[world.start]
    return float(abs(lhs - world.tau * total_kl))

