import torch
import torch.nn as nn
import torch.nn.functional as F

class VQ(nn.Module):
    def __init__(self, num_codes, code_dim, beta=0.25):
        super().__init__()
        self.codebook = nn.Embedding(num_codes, code_dim)
        nn.init.uniform_(self.codebook.weight, -1.0 / num_codes, 1.0 / num_codes)
        self.beta = beta

    def forward(self, z_e):
        B, C, H, W = z_e.shape
        z = z_e.permute(0,2,3,1).contiguous().view(B*H*W, C)
        codes = self.codebook.weight
        d = (z.pow(2).sum(1, keepdim=True)
             - 2*z@codes.t()
             + codes.pow(2).sum(1, keepdim=True).t())
        idx = d.argmin(1)
        z_q = self.codebook(idx).view(B, H, W, C).permute(0,3,1,2).contiguous()
        commit = self.beta * F.mse_loss(z_e.detach(), z_q)
        codebook_loss = F.mse_loss(z_e, z_q.detach())
        z_q = z_e + (z_q - z_e).detach()
        return z_q, idx.view(B, H, W), commit + codebook_loss

class Tokenizer(nn.Module):
    def __init__(self, in_ch=9, hidden=128, code_dim=128, num_codes=1024):
        super().__init__()
        self.enc = nn.Sequential(
            nn.Conv2d(in_ch, 64, 4, 2, 1), nn.ReLU(),
            nn.Conv2d(64, 128, 4, 2, 1), nn.ReLU(),
            nn.Conv2d(128, code_dim, 3, 1, 1)
        )
        self.vq = VQ(num_codes, code_dim)
        self.dec = nn.Sequential(
            nn.ConvTranspose2d(code_dim, 128, 4, 2, 1), nn.ReLU(),
            nn.ConvTranspose2d(128, 64, 4, 2, 1), nn.ReLU(),
            nn.Conv2d(64, in_ch, 3, 1, 1), nn.Sigmoid()
        )
        self.num_codes = num_codes
        self.code_dim = code_dim

    @torch.no_grad()
    def encode_indices(self, x):
        z_e = self.enc(x)
        _, idx, _ = self.vq(z_e)
        return idx

    @torch.no_grad()
    def decode_indices(self, idx):
        B, H, W = idx.shape
        z_q = self.vq.codebook(idx.view(B*H*W)).view(B, H, W, self.vq.codebook.embedding_dim)
        z_q = z_q.permute(0,3,1,2).contiguous()
        x_hat = self.dec(z_q)
        return x_hat

    def forward(self, x):
        z_e = self.enc(x)
        z_q, idx, vq_loss = self.vq(z_e)
        x_hat = self.dec(z_q)
        rec = F.mse_loss(x_hat, x)
        loss = rec + vq_loss
        return x_hat, idx, loss

    @staticmethod
    def tokens_hw(inp_h=84, inp_w=84, levels=2):
        h = inp_h // (2**levels)
        w = inp_w // (2**levels)
        return h, w
