from dataclasses import dataclass
import math
from typing import Dict, Tuple

from einops import rearrange
import torch
import torch.nn as nn

from .models.convnet import FrameEncoder, FrameDecoder
from .data import Batch
from .models.tokenizer.quantizer import Quantizer, QuantizerOutput
from .models.utils import init_weights, LossWithIntermediateLosses

class Tokenizer(nn.Module):
    def __init__(self, config: dict) -> None:
        super().__init__()
        self.config = config

        self.latent_res = config["image_size"] // 2 ** sum(config["encoder_config"]["down"])
        self.tokens_grid_res = int(math.sqrt(config["num_tokens"]))
        self.token_res = self.latent_res // self.tokens_grid_res

        self.encoder_act_emb = nn.Embedding(config["num_actions"], config["image_size"] ** 2)
        self.decoder_act_emb = nn.Embedding(config["num_actions"], config["decoder_act_channels"] * self.latent_res ** 2)

        self.quantizer = Quantizer(
            config["codebook_size"], config["codebook_dim"],
            input_dim=config["encoder_config"]["latent_dim"] * self.token_res ** 2,
            max_codebook_updates_with_revival=config["max_codebook_updates_with_revival"]
        )

        self.encoder = FrameEncoder(config["encoder_config"])
        self.decoder = FrameDecoder(config["decoder_config"])
        self.frame_cnn = FrameEncoder(config["frame_cnn_config"])

        self.apply(init_weights)

    def __repr__(self) -> str:
        return "tokenizer"

    def forward(self, x1: torch.FloatTensor, a: torch.LongTensor, x2: torch.FloatTensor) -> QuantizerOutput:
        z = self.encode(x1, a, x2)
        z = rearrange(z, 'b t c (h k) (w l) -> b t (h w) (k l c)', h=self.tokens_grid_res, w=self.tokens_grid_res)

        return self.quantizer(z)

    def compute_loss(self, batch: Batch, **kwargs) -> Tuple[LossWithIntermediateLosses, Dict]:
        x1 = batch.observations[:, :-1]
        a = batch.actions[:, :-1]
        x2 = batch.observations[:, 1:]

        quantizer_outputs = self(x1, a, x2)

        r = self.decode(x1, a, rearrange(quantizer_outputs.q, 'b t (h w) (k l e) -> b t e (h k) (w l)', h=self.tokens_grid_res, k=self.token_res, l=self.token_res))
        delta = (x2 - r)
        delta = delta[torch.logical_and(batch.mask_padding[:, 1:], batch.mask_padding[:, :-1])]

        losses = {
            **quantizer_outputs.loss,
            'reconstruction_loss_l1': 0.1 * torch.abs(delta).mean(),
            'reconstruction_loss_l2': delta.pow(2).mean(),
            'reconstruction_loss_l2_worst_pixel': 0.01 * rearrange(delta, 'b c h w -> b (c h w)').pow(2).max(dim=-1)[0].mean(),
        }

        return LossWithIntermediateLosses(**losses), quantizer_outputs.metrics

    def encode(self, x1: torch.FloatTensor, a: torch.LongTensor, x2: torch.FloatTensor) -> torch.FloatTensor:
        a_emb = rearrange(self.encoder_act_emb(a), 'b t (h w) -> b t 1 h w', h=x1.size(3))
        encoder_input = torch.cat((x1, a_emb, x2), dim=2)
        z = self.encoder(encoder_input)

        return z

    def decode(self, x1: torch.FloatTensor, a: torch.LongTensor, q2: torch.FloatTensor, should_clamp: bool = False) -> torch.FloatTensor:
        x1_emb = self.frame_cnn(x1)
        a_emb = rearrange(self.decoder_act_emb(a), 'b t (c h w) -> b t c h w', c=self.config["decoder_act_channels"], h=x1_emb.size(3))

        decoder_input = torch.cat((x1_emb, a_emb, q2), dim=2)

        r = self.decoder(decoder_input)
        r = torch.clamp(r, 0, 1).mul(255).round().div(255) if should_clamp else r

        return r

    @torch.no_grad()
    def encode_decode(self, x1: torch.FloatTensor, a: torch.LongTensor, x2: torch.FloatTensor) -> torch.Tensor:
        z = self.encode(x1, a, x2)
        z = rearrange(z, 'b t c (h k) (w l) -> b t (h w) (k l c)', k=self.token_res, l=self.token_res)
        q = rearrange(self.quantizer(z).q, 'b t (h w) (k l e) -> b t e (h k) (w l)', h=self.tokens_grid_res, k=self.token_res, l=self.token_res)
        r = self.decode(x1, a, q, should_clamp=True)

        return r

    @torch.no_grad()
    def burn_in(self, obs: torch.FloatTensor, act: torch.LongTensor) -> torch.LongTensor: 
        assert obs.size(1) == act.size(1) + 1
        quantizer_output = self(obs[:, :-1], act, obs[:, 1:])

        return quantizer_output.tokens