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

from .models.convnet import FrameEncoder, FrameDecoder
from .models.quantizer import Quantizer

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"])

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

    def forward(self, x1: torch.FloatTensor, a: torch.LongTensor, x2: torch.FloatTensor) -> dict:
        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 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

    def embed_tokens(self, tokens):
        q = self.quantizer.embed_tokens(tokens)
        b, t, hw, kle = q.shape

        h = self.tokens_grid_res
        w = self.tokens_grid_res
        k = self.token_res
        l = self.token_res
        e = kle // (k * l)

        q = q.reshape(b, t, h, w, k, l, e)
        q = q.permute(0, 1, 6, 2, 4, 3, 5)
        q = q.reshape(b, t, e, h * k, w * l)

        return q

    @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