"""
-----------------------------------------------------------------------------
Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.

NVIDIA CORPORATION and its licensors retain all intellectual property
and proprietary rights in and to this software, related documentation
and any modifications thereto. Any use, reproduction, disclosure or
distribution of this software and related documentation without an express
license agreement from NVIDIA CORPORATION is strictly prohibited.
-----------------------------------------------------------------------------
"""

import importlib

from transformers import Dinov2Model
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import tqdm
from torchvision import transforms

from flow.configs.schema import ModelConfig
from flow.flow_matching import FlowMatchingScheduler
from flow.modules.dit import DiT
from vae.model import Model as VAE
from vae.utils import sync_timer


class Model(nn.Module):
    def __init__(self, config: ModelConfig) -> None:
        super().__init__()
        self.config = config
        self.precision = torch.bfloat16

        # image condition model (dinov2)
        if self.config.dino_model == "dinov2_vitg14":
            self.dino = Dinov2Model.from_pretrained("facebook/dinov2-giant")
        elif self.config.dino_model == "dinov2_vitl14_reg":
            self.dino = Dinov2Model.from_pretrained("facebook/dinov2-with-registers-large")
        else:
            raise ValueError(f"DINOv2 model {self.config.dino_model} not supported")
    
        # hack to match our implementation
        self.dino.layernorm = torch.nn.Identity()

        self.dino.eval().to(dtype=self.precision)
        self.dino.requires_grad_(False)

        cond_dim = 1024 if self.config.dino_model == "dinov2_vitl14_reg" else 1536
        assert cond_dim == config.hidden_dim, "DINOv2 dim must match backbone dim"

        self.preprocess_cond_image = transforms.Compose(
            [
                transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
            ]
        )

        # vae encoder
        vae_config = importlib.import_module(config.vae_conf).make_config()
        self.vae = VAE(vae_config).eval().to(dtype=self.precision)
        self.vae.requires_grad_(False)

        # load vae
        if self.config.preload_vae:
            try:
                vae_ckpt = torch.load(self.config.vae_ckpt_path, weights_only=True)  # local path
                if "model" in vae_ckpt:
                    vae_ckpt = vae_ckpt["model"]
                self.vae.load_state_dict(vae_ckpt, strict=True)
                del vae_ckpt
                print(f"Loaded VAE from {self.config.vae_ckpt_path}")
            except Exception as e:
                print(
                    f"Failed to load VAE from {self.config.vae_ckpt_path}: {e}, make sure you resumed from a valid checkpoint!"
                )

        # load info from vae config
        if config.latent_size is None:
            config.latent_size = self.vae.config.latent_size
        if config.latent_dim is None:
            config.latent_dim = self.vae.config.latent_dim

        # dit
        self.dit = DiT(
            hidden_dim=config.hidden_dim,
            num_heads=config.num_heads,
            num_layers=config.num_layers,
            latent_size=config.latent_size,
            latent_dim=config.latent_dim,
            qknorm=config.qknorm,
            qknorm_type=config.qknorm_type,
            use_pos_embed=config.use_pos_embed,
            use_parts=config.use_parts,
            part_embed_mode=config.part_embed_mode,
        )

        # num_part condition
        if self.config.use_num_parts_cond:
            assert self.config.use_parts, "use_num_parts_cond requires use_parts"
            self.num_part_embed = nn.Embedding(5, config.hidden_dim)

        # preload from a checkpoint (NOTE: this happens BEFORE checkpointer loading latest checkpoint!)
        if self.config.pretrain_path is not None:
            try:
                ckpt = torch.load(self.config.pretrain_path)  # local path
                self.load_state_dict(ckpt["model"], strict=True)
                del ckpt
                print(f"Loaded DiT from {self.config.pretrain_path}")
            except Exception as e:
                print(
                    f"Failed to load DiT from {self.config.pretrain_path}: {e}, make sure you resumed from a valid checkpoint!"
                )

        # sampler
        self.scheduler = FlowMatchingScheduler(shift=config.flow_shift)

        n_params = 0
        for p in self.dit.parameters():
            n_params += p.numel()
        print(f"Number of parameters in DiT: {n_params/1e6:.2f}M")

    # override state_dict to exclude vae and dino, so we only save the trainable params.
    def state_dict(self, *args, **kwargs):
        state_dict = super().state_dict(*args, **kwargs)

        keys_to_del = []
        for k in state_dict.keys():
            if "vae" in k or "dino" in k:
                keys_to_del.append(k)

        for k in keys_to_del:
            del state_dict[k]

        return state_dict

    # override to support tolerant loading (only load matched shape)
    def load_state_dict(self, state_dict, strict=True, assign=False):
        local_state_dict = self.state_dict()
        seen_keys = {k: False for k in local_state_dict.keys()}
        for k, v in state_dict.items():
            if k in local_state_dict:
                seen_keys[k] = True
                if local_state_dict[k].shape == v.shape:
                    local_state_dict[k].copy_(v)
                else:
                    print(f"mismatching shape for key {k}: loaded {local_state_dict[k].shape} but model has {v.shape}")
            else:
                print(f"unexpected key {k} in loaded state dict")
        for k in seen_keys:
            if not seen_keys[k]:
                print(f"missing key {k} in loaded state dict")

    # this happens before checkpointer loading old models !!!
    def on_train_start(self, memory_format: torch.memory_format = torch.preserve_format) -> None:
        super().on_train_start(memory_format=memory_format)
        device = next(self.dit.parameters()).device

        self.dit.to(dtype=self.precision)

        if self.config.use_num_parts_cond:
            self.num_part_embed.to(dtype=self.precision)

        # cast scheduler to device
        self.scheduler.to(device)

    def get_cond(self, cond_image, num_part=None):
        # image condition
        cond_image = cond_image.to(dtype=self.precision)
        with torch.no_grad():
            cond = self.dino(cond_image).last_hidden_state
        cond = F.layer_norm(cond.float(), cond.shape[-1:]).to(dtype=self.precision)  # [B, L, C]

        # num_part condition
        if self.config.use_num_parts_cond:
            if num_part is None:
                # use a default value (2-10 parts)
                num_part_coarse = torch.ones(cond.shape[0], dtype=torch.int64, device=cond.device) * 2
            else:
                # coarse range
                num_part_coarse = torch.ones(cond.shape[0], dtype=torch.int64, device=cond.device)
                num_part_coarse[num_part == 2] = 1
                num_part_coarse[(num_part > 2) & (num_part <= 10)] = 2
                num_part_coarse[(num_part > 10) & (num_part <= 100)] = 3
                num_part_coarse[num_part > 100] = 4
            num_part_embed = self.num_part_embed(num_part_coarse).unsqueeze(1)  # [B, 1, C]
            cond = torch.cat([cond, num_part_embed], dim=1)  # [B, L+1, C]

        return cond

    def training_step(
        self,
        data: dict[str, torch.Tensor],
        iteration: int,
    ) -> tuple[dict[str, torch.Tensor], torch.Tensor]:
        output = {}
        loss = 0

        cond_images = self.preprocess_cond_image(
            data["cond_images"]
        )  # [B, N, 3, 518, 518], we may load multiple (N) cond images for the same shape
        B, N, C, H, W = cond_images.shape

        if self.config.use_num_parts_cond:
            cond_num_part = data["num_part"].repeat_interleave(N, dim=0)
        else:
            cond_num_part = None

        cond = self.get_cond(cond_images.view(-1, C, H, W), cond_num_part)  # [B*N, L, C]

        # random CFG dropout
        if self.training:
            mask = torch.rand((B * N, 1, 1), device=cond.device, dtype=cond.dtype) >= 0.1
            cond = cond * mask

        with torch.no_grad():
            # encode latent
            if self.config.use_parts:
                # encode two parts and concat latent
                part0_data = {k.replace("_part0", ""): v for k, v in data.items() if "_part0" in k}
                part1_data = {k.replace("_part1", ""): v for k, v in data.items() if "_part1" in k}
                posterior0 = self.vae.encode(part0_data)
                posterior1 = self.vae.encode(part1_data)
                if self.training and self.config.shuffle_parts:
                    if np.random.rand() < 0.5:
                        posterior0, posterior1 = posterior1, posterior0
                latent = torch.cat(
                    [
                        posterior0.mode().float().nan_to_num_(0),
                        posterior1.mode().float().nan_to_num_(0),
                    ],
                    dim=1,
                )  # [B, 2L, C]
            else:
                posterior = self.vae.encode(data)
                latent = posterior.mode().float().nan_to_num_(0)  # use mean as the latent, [B, L, C]

            # repeat latent for each cond image
            if N != 1:
                latent = latent.repeat_interleave(N, dim=0)

            # random sample timesteps and add noise
            noisy_latent, noise, timesteps = self.scheduler.add_noise(
                latent, self.config.logitnorm_mean, self.config.logitnorm_std
            )

        noisy_latent = noisy_latent.to(dtype=self.precision)
        model_pred = self.dit(noisy_latent, cond, timesteps)

        # flow-matching loss
        target = noise - latent
        loss = F.mse_loss(model_pred.float(), target.float())

        # metrics
        with torch.no_grad():
            output["scalar"] = {}  # for wandb logging
            output["scalar"]["loss_mse"] = loss.detach()

        return output, loss

    @torch.no_grad()
    def validation_step(
        self,
        data: dict[str, torch.Tensor],
        iteration: int,
    ) -> tuple[dict[str, torch.Tensor], torch.Tensor]:
        return self.training_step(data, iteration)

    @torch.inference_mode()
    @sync_timer("flow forward")
    def forward(
        self,
        data: dict[str, torch.Tensor],
        num_steps: int = 30,
        cfg_scale: float = 7.0,
        verbose: bool = True,
        generator: torch.Generator | None = None,
    ) -> dict[str, torch.Tensor]:
        # the inference sampling
        cond_images = self.preprocess_cond_image(data["cond_images"])  # [B, 3, 518, 518]
        B = cond_images.shape[0]
        assert B == 1, "Only support batch size 1 for now."

        # num_part condition
        if self.config.use_num_parts_cond and "num_part" in data:
            cond_num_part = data["num_part"]  # [B,], int
        else:
            cond_num_part = None

        cond = self.get_cond(cond_images, cond_num_part)

        if self.config.use_parts:
            x = torch.randn(
                B,
                self.config.latent_size * 2,
                self.config.latent_dim,
                device=cond.device,
                dtype=torch.float32,
                generator=generator,
            )
        else:
            x = torch.randn(
                B,
                self.config.latent_size,
                self.config.latent_dim,
                device=cond.device,
                dtype=torch.float32,
                generator=generator,
            )

        cond_input = torch.cat([cond, torch.zeros_like(cond)], dim=0)

        # flow-matching
        sigmas = np.linspace(1, 0, num_steps + 1)
        sigmas = self.scheduler.shift * sigmas / (1 + (self.scheduler.shift - 1) * sigmas)
        sigmas_pair = list((sigmas[i], sigmas[i + 1]) for i in range(num_steps))

        for sigma, sigma_prev in tqdm.tqdm(sigmas_pair, desc="Flow Sampling", disable=not verbose):
            # classifier-free guidance
            timesteps = torch.tensor([1000 * sigma] * B * 2, device=x.device, dtype=x.dtype)
            x_input = torch.cat([x, x], dim=0)

            # predict v
            x_input = x_input.to(dtype=self.precision)
            pred = self.dit(x_input, cond_input, timesteps).float()
            cond_v, uncond_v = pred.chunk(2, dim=0)
            pred_v = uncond_v + (cond_v - uncond_v) * cfg_scale

            # scheduler step
            x = x - (sigma - sigma_prev) * pred_v

        output = {}
        output["latent"] = x

        # leave mesh extraction to vae
        return output