import torch
from typing import OrderedDict
from attend.asym import ConvND, ConvTransposeND, AsymptoticSphereNorm, GELUTanh, Quantize, VitBlockND
from attend.sep import SpatiotemporallySeparableConvND
from attend.monarch import FactorizedConvND


class ImageCodec(torch.nn.Module):
    def __init__(self, dim, analysis_channels, ps, latent_dim, hidden_dim, decoder_depth, head_dim, near_lossless, expand_ratio=3):
        super().__init__()
        self.layers = torch.nn.Sequential(
            OrderedDict(
                [
                    ("analysis_transform", torch.nn.Sequential(
                        OrderedDict(
                            [
                                ("spatiotempoal", SpatiotemporallySeparableConvND(dim=dim, ch=analysis_channels, kernel_size=ps, groups=analysis_channels[0], bias=True)),
                                ("channels", FactorizedConvND(dim, in_chs=analysis_channels[-1], out_chs=latent_dim, kernel_size=1, bias=True)),
                            ]
                        )
                    )),
                    ("norm", AsymptoticSphereNorm()),
                    ("quantize", Quantize(k=near_lossless)),
                    ("synthesis_transform", torch.nn.Sequential(
                        OrderedDict(
                            [
                                ("nn", torch.nn.Sequential(
                                    ConvND(dim, latent_dim, hidden_dim, kernel_size=1, stride=1, padding=0),
                                    *[VitBlockND(dim, in_channels=hidden_dim, norm_layer=AsymptoticSphereNorm, act_layer=GELUTanh, quant_layer=torch.nn.Identity, head_dim=head_dim, expand_ratio=expand_ratio, drop_path=0.0)
                                      for _ in range(decoder_depth)],
                                    ConvTransposeND(dim, hidden_dim, analysis_channels[0], kernel_size=ps, stride=ps, padding=0),
                                )),
                            ]
                        )
                    )),
                ]
            )
        )

    def forward(self, x):
        return self.layers(x)

    def forward_rate(self, x):
        z = self.layers[:3](x)
        rate = z.std().log2()
        xhat = self.layers[-1](z)
        return xhat, rate

    def forward_hard_quant(self, x):
        with torch.no_grad():
            z = self.layers[:2](x).round()
        xhat = self.layers[-1](z)
        return xhat
