from __future__ import annotations

from dataclasses import dataclass
from typing import Dict, Iterable, List, Mapping, Sequence, Tuple

import numpy as np

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


@dataclass
class NeuralConditionals:
    pi_hat: Dict[History, np.ndarray]
    mu_hat: Dict[ChanceState, np.ndarray]


def encode_trajectory(trajectory: Sequence[tuple[int, int]], n_actions: int) -> List[int]:
    tokens: List[int] = []
    for action, event in trajectory:
        tokens.append(1 + int(action))
        tokens.append(1 + n_actions + int(event))
    return tokens


def encode_history(state: History, n_actions: int) -> List[int]:
    return encode_trajectory(state, n_actions)


def sample_token_sequences(
    world: LifeGridWorld, n_trajectories: int, seed: int
) -> np.ndarray:
    rng = np.random.default_rng(seed)
    sequences = np.empty((n_trajectories, 2 * world.horizon), dtype=np.int64)
    for i in range(n_trajectories):
        state: History = ()
        cursor = 0
        for _depth in range(world.horizon):
            action = int(rng.choice(world.n_actions, p=world.behavior_policy(state)))
            event = int(rng.choice(world.n_events, p=world.event_kernel(state, action)))
            sequences[i, cursor] = 1 + action
            sequences[i, cursor + 1] = 1 + world.n_actions + event
            cursor += 2
            state = world.next_state(state, action, event)
    return sequences


def make_model(world: LifeGridWorld, seed: int = 0):
    import torch
    from torch import nn

    torch.manual_seed(seed)
    vocab_size = 2 + world.n_actions + world.n_events  # BOS + typed tokens + PAD
    pad_id = vocab_size - 1

    class TinyCausalTransformer(nn.Module):
        def __init__(self):
            super().__init__()
            width = 32
            self.token_embedding = nn.Embedding(vocab_size, width, padding_idx=pad_id)
            self.position_embedding = nn.Embedding(2 * world.horizon + 1, width)
            layer = nn.TransformerEncoderLayer(
                d_model=width,
                nhead=4,
                dim_feedforward=128,
                dropout=0.0,
                activation="gelu",
                batch_first=True,
                norm_first=True,
            )
            self.encoder = nn.TransformerEncoder(layer, num_layers=2, enable_nested_tensor=False)
            self.output = nn.Linear(width, vocab_size)

        def forward(self, tokens, padding_mask=None):
            batch, length = tokens.shape
            positions = torch.arange(length, device=tokens.device).unsqueeze(0)
            h = self.token_embedding(tokens) + self.position_embedding(positions)
            causal = torch.triu(
                torch.ones(length, length, device=tokens.device, dtype=torch.bool), diagonal=1
            )
            h = self.encoder(h, mask=causal, src_key_padding_mask=padding_mask)
            return self.output(h)

    return TinyCausalTransformer(), pad_id


def train_transformer(
    world: LifeGridWorld,
    sequences: np.ndarray,
    seed: int,
    steps: int = 1200,
    batch_size: int = 512,
    learning_rate: float = 3e-3,
    weight_decay: float = 1e-4,
    device: str = "cpu",
    num_threads: int = 8,
) -> tuple[object, int, list[dict[str, float]]]:
    import torch
    from torch import nn

    torch.manual_seed(seed)
    np_rng = np.random.default_rng(seed)
    torch.set_num_threads(max(1, min(num_threads, torch.get_num_threads())))
    model, pad_id = make_model(world, seed)
    model.to(device)
    optimizer = torch.optim.AdamW(
        model.parameters(), lr=learning_rate, weight_decay=weight_decay
    )
    total = sequences.shape[0]
    bos = np.zeros((total, 1), dtype=np.int64)
    inputs = np.concatenate([bos, sequences[:, :-1]], axis=1)
    targets = sequences
    history: list[dict[str, float]] = []
    model.train()
    for step in range(1, steps + 1):
        idx = np_rng.integers(0, total, size=min(batch_size, total))
        x = torch.as_tensor(inputs[idx], dtype=torch.long, device=device)
        y = torch.as_tensor(targets[idx], dtype=torch.long, device=device)
        logits = model(x)

        valid = torch.zeros_like(logits, dtype=torch.bool)
        action_ids = slice(1, 1 + world.n_actions)
        event_ids = slice(1 + world.n_actions, 1 + world.n_actions + world.n_events)
        valid[:, 0::2, action_ids] = True
        valid[:, 1::2, event_ids] = True
        logits = logits.masked_fill(~valid, -1e9)
        loss = nn.functional.cross_entropy(logits.reshape(-1, logits.shape[-1]), y.reshape(-1))
        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()

        if step == 1 or step % 50 == 0 or step == steps:
            history.append({"step": float(step), "loss": float(loss.detach().cpu())})
    return model, pad_id, history


def _predict_group(model, prefixes: list[list[int]], allowed_ids: Sequence[int], device: str):
    import torch

    if not prefixes:
        return np.empty((0, len(allowed_ids)))
    tokens = torch.as_tensor(prefixes, dtype=torch.long, device=device)
    model.eval()
    with torch.no_grad():
        logits = model(tokens)[:, -1, allowed_ids]
        probs = torch.softmax(logits, dim=-1).cpu().numpy()
    return probs


def extract_conditionals(
    world: LifeGridWorld, model, device: str = "cpu"
) -> NeuralConditionals:
    pi_hat: Dict[History, np.ndarray] = {}
    mu_hat: Dict[ChanceState, np.ndarray] = {}
    action_ids = list(range(1, 1 + world.n_actions))
    event_ids = list(
        range(1 + world.n_actions, 1 + world.n_actions + world.n_events)
    )
    states = list(world.decision_states(include_terminal=False))

    for depth in range(world.horizon):
        group = [s for s in states if len(s) == depth]
        prefixes = [[0] + encode_history(s, world.n_actions) for s in group]
        probs = _predict_group(model, prefixes, action_ids, device)
        for state, p in zip(group, probs):
            pi_hat[state] = p.astype(np.float64)

        chance_states = [(s, a) for s in group for a in range(world.n_actions)]
        chance_prefixes = [
            [0] + encode_history(s, world.n_actions) + [1 + a]
            for s, a in chance_states
        ]
        event_probs = _predict_group(model, chance_prefixes, event_ids, device)
        for key, p in zip(chance_states, event_probs):
            mu_hat[key] = p.astype(np.float64)
    return NeuralConditionals(pi_hat=pi_hat, mu_hat=mu_hat)


def neural_exact_trajectory_kl(
    world: LifeGridWorld, conditionals: NeuralConditionals
) -> float:
    kl = 0.0
    for trajectory in world.trajectories():
        state: History = ()
        p = 1.0
        q = 1.0
        for action, event in trajectory:
            true_pi = world.behavior_policy(state)
            true_mu = world.event_kernel(state, action)
            p *= true_pi[action] * true_mu[event]
            q *= conditionals.pi_hat[state][action] * conditionals.mu_hat[(state, action)][event]
            state = world.next_state(state, action, event)
        kl += p * np.log(p / max(q, 1e-300))
    return float(kl)


def neural_conditional_metrics(
    world: LifeGridWorld, conditionals: NeuralConditionals
) -> dict[str, float]:
    behavior = {
        state: world.behavior_policy(state)
        for state in world.decision_states(include_terminal=False)
    }
    occupancy = policy_occupancy(world, behavior)
    action_weighted = 0.0
    action_mass = 0.0
    event_weighted = 0.0
    event_mass = 0.0
    max_event_tv = 0.0
    for state in world.decision_states(include_terminal=False):
        mass = occupancy.get(state, 0.0)
        action_weighted += mass * 0.5 * np.abs(
            behavior[state] - conditionals.pi_hat[state]
        ).sum()
        action_mass += mass
        for action in range(world.n_actions):
            chance_mass = mass * behavior[state][action]
            tv = 0.5 * np.abs(
                world.event_kernel(state, action) - conditionals.mu_hat[(state, action)]
            ).sum()
            event_weighted += chance_mass * tv
            event_mass += chance_mass
            max_event_tv = max(max_event_tv, float(tv))
    return {
        "action_weighted_tv": float(action_weighted / action_mass),
        "event_weighted_tv": float(event_weighted / event_mass),
        "max_event_tv": float(max_event_tv),
    }
