import torch
import pytorch_lightning as pl
import torch.nn.functional as F
from contextlib import contextmanager

from ldm.modules.distributions.distributions import DiagonalGaussianDistribution

from ldm.util import instantiate_from_config
from ldm.modules.ema import LitEma

try:
    from modules.models import Encoder, Decoder
except:
    from modi_vae.models import Encoder, Decoder

class AutoencoderKL(pl.LightningModule):
    def __init__(self,
                 embed_dim=4,
                 ckpt_path=None,
                 ignore_keys=[],
                 image_key="image",
                 colorize_nlabels=None,
                 monitor=None,
                 ema_decay=None,
                 learn_logvar=False,
                 load_checkpoint=True
                 ):
        super().__init__()
        self.encoder = Encoder(double_z=True, z_channels=4, resolution=256, in_channels=3, out_ch=3, ch=128, ch_mult=[1,2,4,4], num_res_blocks=2, attn_resolutions=[], dropout=0.0)
        self.decoder = Decoder(double_z=True, z_channels=4, resolution=256, in_channels=3, out_ch=3, ch=128, ch_mult=[1,2,4,4], num_res_blocks=2, attn_resolutions=[], dropout=0.0)

        self.quant_conv = torch.nn.Conv2d(2*4, 2*embed_dim, 1)
        self.post_quant_conv = torch.nn.Conv2d(embed_dim, 4, 1)
        self.embed_dim = embed_dim
        if colorize_nlabels is not None:
            assert type(colorize_nlabels)==int
            self.register_buffer("colorize", torch.randn(3, colorize_nlabels, 1, 1))
        if monitor is not None:
            self.monitor = monitor
        
        if load_checkpoint:
            state_dict = torch.load('/data07/v-wenjwang/ControlNet/CIConv/models/control_sd15_ini.ckpt', map_location=torch.device("cpu"))
            new_state_dict = {}
            for s in state_dict:
                if "first_stage_model" in s:
                    new_state_dict[s.replace("first_stage_model.", "")] = state_dict[s]
            self.load_state_dict(new_state_dict, strict=False)

    def encode(self, x):
        h, hs = self.encoder(x)
        moments = self.quant_conv(h)
        posterior = DiagonalGaussianDistribution(moments)
        return posterior, hs

    def decode(self, z, hs):
        z = self.post_quant_conv(z)
        dec = self.decoder(z, hs)
        return dec

    def forward(self, input, sample_posterior=True):
        posterior, hs = self.encode(input)
        if sample_posterior:
            z = posterior.sample()
        else:
            z = posterior.mode()
        dec = self.decode(z, hs)
        return dec, posterior

if __name__ == "__main__":
    from data.laion_dataset import create_webdataset
    import torchvision

    image_dataset = create_webdataset(
        data_dir="/data06/v-wenjwang/COCO-2017/*/*.*",
    )

    import webdataset as wds
    image_dataloader = wds.WebLoader(
        dataset          =   image_dataset,
        batch_size       =   1,
        num_workers      =   8,
        pin_memory       =   True,
        prefetch_factor  =   2,
    )

    model = AutoencoderKL().cuda()

    for data in image_dataloader:
        img = data["distorted"].cuda()
        img = model(img)[0]

        torchvision.utils.save_image(img*0.5+0.5, "distorted.png")
        torchvision.utils.save_image(data["distorted"]*0.5+0.5, "original.png")

        break