"""Local recurrent QMIX/VDN trainer without the historical offpolicy package."""

import copy

import torch
import torch.nn.functional as F

from algorithms.qmix.algorithm.q_mixer import QMixer
from algorithms.vdn.algorithm.vdn_mixer import VDNMixer


def _to_tensor_tree(value, device):
    if isinstance(value, dict):
        return {key: _to_tensor_tree(item, device) for key, item in value.items()}
    return torch.as_tensor(value, dtype=torch.float32, device=device)


def _time(tree, index):
    if isinstance(tree, dict):
        return {key: value[:, index] for key, value in tree.items()}
    return tree[:, index]


def _time_slice(tree, start, stop=None):
    if isinstance(tree, dict):
        return {key: _time_slice(value, start, stop) for key, value in tree.items()}
    return tree[:, start:stop]


class QMix:
    def __init__(self, args, policy, state_dim, num_agents, device):
        self.args = args
        self.policy = policy
        self.target_policy = copy.deepcopy(policy).to(device)
        self.device = device
        self.num_agents = int(num_agents)
        trainer_name = getattr(args, "trainer_algorithm_name", None)
        self.vdn = str(trainer_name or args.algorithm_name).lower() == "vdn"
        if self.vdn:
            self.mixer = VDNMixer().to(device)
        else:
            self.mixer = QMixer(
                state_dim, num_agents,
                args.qmix_mixer_hidden_dim,
                args.qmix_hypernet_hidden_dim,
            ).to(device)
        self.target_mixer = copy.deepcopy(self.mixer).to(device)
        parameters = list(self.policy.q_network.parameters()) + list(self.mixer.parameters())
        if args.qmix_optimizer == "rmsprop":
            self.optimizer = torch.optim.RMSprop(
                parameters,
                lr=args.lr,
                alpha=args.qmix_rmsprop_alpha,
                eps=args.opti_eps,
                weight_decay=args.weight_decay,
            )
        else:
            self.optimizer = torch.optim.Adam(
                parameters,
                lr=args.lr,
                eps=args.opti_eps,
                weight_decay=args.weight_decay,
            )
        self.parameters = parameters
        self.update_count = 0

    def prep_training(self):
        self.policy.train()
        self.mixer.train()

    def prep_rollout(self):
        self.policy.eval()
        self.mixer.eval()

    def train_batch(self, batch):
        self.prep_training()
        batch = _to_tensor_tree(batch, self.device)
        observations = batch["obs"]
        states = batch["states"]
        available = batch["available_actions"]
        actions = batch["actions"].long()
        batch_size, transitions = actions.shape[:2]
        hidden = torch.zeros(
            batch_size, self.num_agents, self.policy.hidden_size, device=self.device
        )
        target_hidden = torch.zeros_like(hidden)
        burn_in = min(int(self.args.qmix_burn_in), transitions)
        recurrent = bool(self.args.qmix_use_recurrent_agent)
        if burn_in > 0 and recurrent:
            prefix_obs = _time_slice(observations, 0, burn_in)
            prefix_available = available[:, :burn_in]
            prefix_alive = batch["alive_masks"][:, :burn_in]
            if self.policy.use_hatt:
                chunk_size = int(self.args.qmix_hatt_sequence_chunk_size)
                with torch.no_grad():
                    _, hidden = self.policy.q_network.forward_sequence(
                        prefix_obs, hidden, prefix_available, chunk_size,
                        alive_masks=prefix_alive,
                    )
                    _, target_hidden = self.target_policy.q_network.forward_sequence(
                        prefix_obs, target_hidden, prefix_available, chunk_size,
                        alive_masks=prefix_alive,
                    )
            else:
                with torch.no_grad():
                    for step in range(burn_in):
                        hidden = hidden * prefix_alive[:, step]
                        target_hidden = target_hidden * prefix_alive[:, step]
                        _, hidden = self.policy.q_network.forward_step(
                            _time(prefix_obs, step), hidden, prefix_available[:, step]
                        )
                        _, target_hidden = self.target_policy.q_network.forward_step(
                            _time(prefix_obs, step), target_hidden,
                            prefix_available[:, step]
                        )
            hidden = hidden.detach()
            target_hidden = target_hidden.detach()

        # Burn-in establishes recurrent context without retaining its graph.  The
        # same leading samples are dropped for feedforward runs to keep their
        # number of supervised transitions comparable.
        observations = _time_slice(observations, burn_in)
        states = states[:, burn_in:]
        available = available[:, burn_in:]
        for key in (
            "actions", "rewards", "dones", "held_actions", "discounts", "valid"
        ):
            batch[key] = batch[key][:, burn_in:]
        for key in ("active_masks", "alive_masks"):
            batch[key] = batch[key][:, burn_in:]
        actions = batch["actions"].long()
        transitions = actions.shape[1]
        if self.policy.use_hatt:
            chunk_size = int(self.args.qmix_hatt_sequence_chunk_size)
            online_values, hidden = self.policy.q_network.forward_sequence(
                observations, hidden, available, chunk_size,
                alive_masks=batch["alive_masks"],
            )
            with torch.no_grad():
                target_values, target_hidden = (
                    self.target_policy.q_network.forward_sequence(
                        observations, target_hidden, available, chunk_size,
                        alive_masks=batch["alive_masks"],
                    )
                )
        else:
            online_steps, target_steps = [], []
            for step in range(transitions + 1):
                hidden = hidden * batch["alive_masks"][:, step]
                target_hidden = target_hidden * batch["alive_masks"][:, step]
                q_values, hidden = self.policy.q_network.forward_step(
                    _time(observations, step), hidden, available[:, step]
                )
                with torch.no_grad():
                    target_q, target_hidden = self.target_policy.q_network.forward_step(
                        _time(observations, step), target_hidden, available[:, step]
                    )
                online_steps.append(q_values)
                target_steps.append(target_q)
            online_values = torch.stack(online_steps, dim=1)
            target_values = torch.stack(target_steps, dim=1)
        chosen = online_values[:, :-1].gather(-1, actions.unsqueeze(-1)).squeeze(-1)
        next_online = online_values[:, 1:].masked_fill(
            available[:, 1:] <= 0, torch.finfo(online_values.dtype).min
        )
        next_target = target_values[:, 1:].masked_fill(
            available[:, 1:] <= 0, torch.finfo(target_values.dtype).min
        )
        if self.args.qmix_double_q:
            greedy_actions = next_online.argmax(dim=-1)
        else:
            greedy_actions = next_target.argmax(dim=-1)
        transition_mode = str(self.args.qmix_transition_mode)
        if transition_mode == "physical_legacy":
            next_actions = greedy_actions
            current_alive = torch.ones_like(batch["alive_masks"][:, :-1, :, 0])
            next_alive = torch.ones_like(batch["alive_masks"][:, 1:, :, 0])
        else:
            next_active = batch["active_masks"][:, 1:, :, 0] > 0.5
            next_alive = batch["alive_masks"][:, 1:, :, 0]
            fixed_actions = batch["held_actions"].long()
            next_actions = torch.where(
                next_active & (next_alive > 0.5), greedy_actions, fixed_actions
            )
            current_alive = batch["alive_masks"][:, :-1, :, 0]
        next_chosen = target_values[:, 1:].gather(
            -1, next_actions.unsqueeze(-1)
        ).squeeze(-1)
        chosen = chosen * current_alive
        next_chosen = next_chosen * next_alive
        joint_q = self.mixer(chosen, states[:, :-1])
        with torch.no_grad():
            target_joint_q = self.target_mixer(next_chosen, states[:, 1:])
            targets = batch["rewards"] + batch["discounts"] * (
                1.0 - batch["dones"]
            ) * target_joint_q
        valid = batch["valid"].clone()
        td_error = joint_q - targets
        max_abs_q = torch.maximum(joint_q.detach().abs().max(), targets.abs().max())
        threshold = float(self.args.qmix_max_abs_q)
        if not torch.isfinite(td_error).all() or not torch.isfinite(max_abs_q):
            raise FloatingPointError("QMIX produced a non-finite Q value or TD error")
        if threshold > 0.0 and float(max_abs_q) > threshold:
            raise FloatingPointError(
                f"QMIX Q magnitude {float(max_abs_q):.3g} exceeded "
                f"qmix_max_abs_q={threshold:.3g}"
            )
        if self.args.use_huber_loss:
            losses = F.smooth_l1_loss(joint_q, targets, reduction="none")
        else:
            losses = td_error.square()
        loss = (losses * valid).sum() / valid.sum().clamp_min(1.0)
        self.optimizer.zero_grad(set_to_none=True)
        loss.backward()
        grad_norm = torch.nn.utils.clip_grad_norm_(
            self.parameters, self.args.max_grad_norm
        )
        self.optimizer.step()
        self.update_count += 1
        info = {
            "loss": float(loss.detach()),
            "grad_norm": float(grad_norm),
            "Q_tot": float((joint_q * valid).sum().detach() / valid.sum().clamp_min(1.0)),
            "td_error": float((td_error.abs() * valid).sum().detach() / valid.sum().clamp_min(1.0)),
            "max_abs_q": float(max_abs_q),
        }
        if transition_mode != "physical_legacy":
            constrained = (batch["active_masks"][:, 1:, :, 0] <= 0.5) & (
                next_alive > 0.5
            )
            denom = constrained.float().sum().clamp_min(1.0)
            info["inactive_argmax_mismatch"] = float(
                (((greedy_actions != batch["held_actions"].long()) & constrained)
                 .float().sum() / denom).detach()
            )
            with torch.no_grad():
                unconstrained_utilities = target_values[:, 1:].gather(
                    -1, greedy_actions.unsqueeze(-1)
                ).squeeze(-1) * next_alive
                unconstrained_joint = self.target_mixer(
                    unconstrained_utilities, states[:, 1:]
                )
                info["constrained_target_gap"] = float(
                    ((unconstrained_joint - target_joint_q).abs() * valid).sum()
                    / valid.sum().clamp_min(1.0)
                )
        return info

    def hard_target_updates(self):
        self.target_policy.q_network.load_state_dict(self.policy.q_network.state_dict())
        self.target_mixer.load_state_dict(self.mixer.state_dict())

    def state_dict(self):
        return {
            "policy": self.policy.get_actor_checkpoint(),
            "mixer": self.mixer.state_dict(),
            "target_policy": self.target_policy.q_network.state_dict(),
            "target_mixer": self.target_mixer.state_dict(),
            "optimizer": self.optimizer.state_dict(),
            "update_count": self.update_count,
        }