from __future__ import annotations

from dataclasses import dataclass
from typing import Optional, Tuple

import torch
import torch.nn as nn
from torch.distributions import Normal

from transformers import PreTrainedModel
from transformers.utils import ModelOutput

from .configuration_mrbalance import MrBalanceConfig


@dataclass
class MrBalanceOutput(ModelOutput):
    """
    Outputs returned by MrBalance.
    """

    action: Optional[torch.FloatTensor] = None
    action_mean: Optional[torch.FloatTensor] = None
    action_std: Optional[torch.FloatTensor] = None
    value: Optional[torch.FloatTensor] = None
    log_prob: Optional[torch.FloatTensor] = None


class MrBalanceMLPForRL(PreTrainedModel):
    """
    Hugging Face-compatible MrBalance Actor-Critic network.

    Exact architecture:

        64 -> 128 -> 128 -> 64

    with:

        actor:  64 -> 2
        critic: 64 -> 1

    and:

        learned log_std: 2
    """

    config_class = MrBalanceConfig
    base_model_prefix = "mrbalance"
    main_input_name = "observation"

    def __init__(self, config: MrBalanceConfig):
        super().__init__(config)

        self.backbone = nn.Sequential(
            nn.Linear(
                config.observation_size,
                config.hidden_size,
            ),
            nn.SiLU(),

            nn.Linear(
                config.hidden_size,
                config.intermediate_size,
            ),
            nn.SiLU(),

            nn.Linear(
                config.intermediate_size,
                config.bottleneck_size,
            ),
            nn.SiLU(),
        )

        self.actor = nn.Linear(
            config.bottleneck_size,
            config.action_size,
        )

        self.critic = nn.Linear(
            config.bottleneck_size,
            1,
        )

        self.log_std = nn.Parameter(
            torch.full(
                (config.action_size,),
                config.actor_log_std_init,
            )
        )

        self.post_init()

    def _init_weights(self, module: nn.Module) -> None:
        """
        Match the original MrBalance initialization.

        These values are only relevant for a newly initialized model.
        During export, the learned checkpoint weights overwrite them.
        """

        if not isinstance(module, nn.Linear):
            return

        if module is self.actor:
            gain = 0.01
        elif module is self.critic:
            gain = 1.0
        else:
            gain = 2.0 ** 0.5

        nn.init.orthogonal_(
            module.weight,
            gain=gain,
        )

        if module.bias is not None:
            nn.init.zeros_(module.bias)

    def _stats(
        self,
        observation: torch.Tensor,
    ) -> Tuple[
        torch.Tensor,
        torch.Tensor,
        torch.Tensor,
    ]:
        hidden = self.backbone(observation)

        mean = self.actor(hidden)

        value = self.critic(hidden).squeeze(-1)

        std = self.log_std.exp().expand_as(mean)

        return mean, std, value

    def get_value(
        self,
        observation: torch.Tensor,
    ) -> torch.Tensor:
        _, _, value = self._stats(observation)
        return value

    def get_action_and_value(
        self,
        observation: torch.Tensor,
        raw_action: Optional[torch.Tensor] = None,
        deterministic: bool = False,
    ):
        mean, std, value = self._stats(observation)

        dist = Normal(mean, std)

        if raw_action is None:
            raw_action = (
                mean
                if deterministic
                else dist.sample()
            )

        action = torch.tanh(raw_action)

        # Exact same Tanh log-probability correction
        # as the original PPO implementation.
        log_prob = dist.log_prob(
            raw_action
        ).sum(-1)

        log_prob -= torch.log(
            torch.clamp(
                1.0 - action.pow(2),
                min=1e-6,
            )
        ).sum(-1)

        entropy = dist.entropy().sum(-1)

        return (
            action,
            log_prob,
            entropy,
            value,
            raw_action,
        )

    def forward(
        self,
        observation: torch.Tensor,
        deterministic: bool = False,
        raw_action: Optional[torch.Tensor] = None,
        return_dict: bool = True,
        **kwargs,
    ):
        action, log_prob, _, value, _ = (
            self.get_action_and_value(
                observation=observation,
                raw_action=raw_action,
                deterministic=deterministic,
            )
        )

        mean, std, _ = self._stats(observation)

        if not return_dict:
            return (
                action,
                mean,
                std,
                value,
                log_prob,
            )

        return MrBalanceOutput(
            action=action,
            action_mean=mean,
            action_std=std,
            value=value,
            log_prob=log_prob,
        )

    @torch.no_grad()
    def predict_action(
        self,
        observation: torch.Tensor,
        deterministic: bool = True,
    ) -> torch.Tensor:
        output = self.forward(
            observation,
            deterministic=deterministic,
        )
        return output.action


MrBalanceModel = MrBalanceMLPForRL