from __future__ import annotations

import math

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
from torch.nn.parameter import Parameter

from .metrics import masked_mae


class TemporalDecay(nn.Module):
    def __init__(
        self,
        input_size: int,
        output_size: int,
        diag: bool = False,
        use_self_attention: bool = False,
        attention_size: int | None = None,
    ) -> None:
        super().__init__()
        self.diag = diag
        self.use_self_attention = use_self_attention
        self.W = Parameter(torch.Tensor(output_size, input_size))
        self.b = Parameter(torch.Tensor(output_size))
        if diag:
            if input_size != output_size:
                raise ValueError("input_size must equal output_size when diag=True.")
            m = torch.eye(input_size, input_size)
            self.register_buffer("m", m)
        if use_self_attention:
            attn_size = attention_size or max(1, input_size)
            self.query = nn.Linear(1, attn_size)
            self.key = nn.Linear(1, attn_size)
            self.value = nn.Linear(1, attn_size)
            self.attn_out = nn.Linear(attn_size, 1)
            self.attn_projection = nn.Linear(input_size, output_size)
        self.reset_parameters()

    def reset_parameters(self) -> None:
        stdv = 1.0 / math.sqrt(self.W.size(0))
        self.W.data.uniform_(-stdv, stdv)
        self.b.data.uniform_(-stdv, stdv)

    def forward(self, deltas: torch.Tensor) -> torch.Tensor:
        if self.diag:
            gamma = F.relu(F.linear(deltas, self.W * Variable(self.m), self.b))
        else:
            gamma = F.relu(F.linear(deltas, self.W, self.b))
        if self.use_self_attention:
            delta_tokens = deltas.unsqueeze(-1)
            query = self.query(delta_tokens)
            key = self.key(delta_tokens)
            value = self.value(delta_tokens)

            scores = torch.matmul(query, key.transpose(-1, -2)) / math.sqrt(query.size(-1))
            weights = torch.softmax(scores, dim=-1)
            attended = torch.matmul(weights, value)
            attended = self.attn_out(attended).squeeze(-1)
            gamma = gamma + F.relu(self.attn_projection(attended))
        return torch.exp(-gamma)


class ITINBlock(nn.Module):
    def __init__(self, n_steps: int, n_features: int, rnn_hidden_size: int) -> None:
        super().__init__()
        self.n_steps = n_steps
        self.n_features = n_features
        self.rnn_hidden_size = rnn_hidden_size
        self.rnn_cell = nn.GRUCell(n_features * 2, rnn_hidden_size)
        self.temp_decay_h = TemporalDecay(n_features, rnn_hidden_size, diag=False)
        self.temp_decay_x = TemporalDecay(
            n_features,
            n_features,
            diag=True,
            use_self_attention=True,
        )
        self.hist_reg = nn.Linear(rnn_hidden_size, n_features)
        self.input_reg = nn.Linear(n_features, n_features)
        self.combining_weight = nn.Linear(n_features * 2, n_features)

    def impute(self, inputs: dict, direction: str):
        values = inputs[direction]["X"]
        masks = inputs[direction]["missing_mask"]
        deltas = inputs[direction]["deltas"]
        hidden_states = torch.zeros((values.size(0), self.rnn_hidden_size), device=values.device)

        estimations = []
        reconstruction_loss = torch.tensor(0.0, device=values.device)
        for t in range(self.n_steps):
            x = values[:, t, :]
            m = masks[:, t, :]
            d = deltas[:, t, :]

            gamma_h = self.temp_decay_h(d)
            gamma_x = self.temp_decay_x(d)
            hidden_states = hidden_states * gamma_h

            x_h = self.hist_reg(hidden_states)
            reconstruction_loss += masked_mae(x_h, x, m)

            x_c = m * x + (1 - m) * x_h
            z_h = self.input_reg(x_c)
            reconstruction_loss += masked_mae(z_h, x, m)

            beta = torch.sigmoid(self.combining_weight(torch.cat([gamma_x, m], dim=1)))
            c_h = beta * x_h + (1 - beta) * z_h
            reconstruction_loss += masked_mae(c_h, x, m)

            c_c = m * x + (1 - m) * c_h
            estimations.append(c_h.unsqueeze(1))
            hidden_states = self.rnn_cell(torch.cat([c_c, m], dim=1), hidden_states)

        estimations = torch.cat(estimations, dim=1)
        imputed_data = masks * values + (1 - masks) * estimations
        return imputed_data, estimations, reconstruction_loss / (self.n_steps * 3)

    def forward(self, inputs: dict, direction: str = "forward") -> dict:
        imputed_data, estimations, reconstruction_loss = self.impute(inputs, direction)
        return {
            "consistency_loss": torch.tensor(0.0, device=imputed_data.device),
            "reconstruction_loss": reconstruction_loss,
            "imputed_data": imputed_data,
            "reconstructed_data": estimations,
        }


class ITINCore(nn.Module):
    def __init__(self, n_steps: int, n_features: int, rnn_hidden_size: int) -> None:
        super().__init__()
        self.itin_f = ITINBlock(n_steps, n_features, rnn_hidden_size)
        self.itin_b = ITINBlock(n_steps, n_features, rnn_hidden_size)

    @staticmethod
    def _reverse(ret: dict) -> dict:
        def reverse_tensor(tensor_: torch.Tensor) -> torch.Tensor:
            if tensor_.dim() <= 1:
                return tensor_
            indices = torch.arange(tensor_.size(1) - 1, -1, -1, device=tensor_.device)
            return tensor_.index_select(1, indices)

        return {key: reverse_tensor(value) for key, value in ret.items()}

    @staticmethod
    def _consistency_loss(pred_f: torch.Tensor, pred_b: torch.Tensor) -> torch.Tensor:
        return torch.abs(pred_f - pred_b).mean() * 1e-1

    def forward(self, inputs: dict, training: bool = True) -> dict:
        ret_f = self.itin_f(inputs, "forward")
        ret_b = self._reverse(self.itin_b(inputs, "backward"))
        imputed_data = (ret_f["imputed_data"] + ret_b["imputed_data"]) / 2

        results = {"imputed_data": imputed_data}
        if training:
            consistency_loss = self._consistency_loss(ret_f["imputed_data"], ret_b["imputed_data"])
            loss = consistency_loss + ret_f["reconstruction_loss"] + ret_b["reconstruction_loss"]
            results["consistency_loss"] = consistency_loss
            results["loss"] = loss
        return results
