# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from contextlib import nullcontext
from dataclasses import dataclass
from typing import Callable

import torch
import torch.nn as nn
from megatron.core import parallel_state
from megatron.core.models.common.vision_module.vision_module import VisionModule
from megatron.core.tensor_parallel.layers import ColumnParallelLinear
from megatron.core.transformer.transformer_config import TransformerConfig
from torch.nn import functional as F

from nemo.collections.diffusion.models.dit.dit_layer_spec import (
    FluxSingleTransformerBlock,
    MMDiTLayer,
    get_flux_double_transformer_engine_spec,
    get_flux_single_transformer_engine_spec,
)
from nemo.collections.diffusion.models.flux.layers import EmbedND, MLPEmbedder, TimeStepEmbedder
from nemo.collections.diffusion.models.flux.model import FluxConfig, FluxModelParams, MegatronFluxModel
from nemo.collections.diffusion.models.flux_controlnet.layers import ControlNetConditioningEmbedding
from nemo.lightning import io
from nemo.utils import logging


def zero_module(module):
    """
    Initializes all parameters of the given module to zero.

    Args:
        module (nn.Module): The module whose parameters will be initialized to zero.

    Returns:
        nn.Module: The same module with zero-initialized parameters.
    """
    for p in module.parameters():
        nn.init.zeros_(p)
    return module


def flux_controlnet_data_step(dataloader_iter):
    """
    Processes a single step of data from a dataloader iterator for the Flux ControlNet.

    Args:
        dataloader_iter (Iterator): An iterator over the dataloader that provides batches of data.

    Returns:
        dict: A processed batch dictionary with an added 'loss_mask' key.
    """
    batch = next(dataloader_iter)
    if isinstance(batch, tuple) and len(batch) == 3:
        _batch = batch[0]
    else:
        _batch = batch

    _batch['loss_mask'] = torch.Tensor([1.0]).cuda(non_blocking=True)
    return _batch


@dataclass
class FluxControlNetConfig(TransformerConfig, io.IOMixin):
    '''
    Flux config inherits from TransformerConfig class.
    '''

    num_layers: int = 1  # dummy setting
    patch_size: int = 1
    in_channels: int = 64
    num_joint_layers: int = 4
    num_single_layers: int = 10
    hidden_size: int = 3072
    num_attention_heads: int = 24
    vec_in_dim: int = 768
    context_dim: int = 4096
    guidance_embed: bool = True
    num_mode: int = None
    model_channels: int = 256
    conditioning_embedding_channels: int = None
    rotary_interleaved: bool = True
    layernorm_epsilon: float = 1e-06
    hidden_dropout: float = 0
    attention_dropout: float = 0
    add_qkv_bias: bool = True
    use_cpu_initialization: bool = True

    load_from_flux_transformer: bool = True
    guidance_scale: float = 3.5

    data_step_fn: Callable = flux_controlnet_data_step


class FluxControlNet(VisionModule):
    """
    A VisionModule-based neural network designed for Flux ControlNet tasks.


    Args:
        config (FluxControlNetConfig):
        Configuration object containing model parameters such as input channels, hidden size, patch size,
            and number of transformer layers.
    """

    def __init__(self, config: FluxControlNetConfig):
        """
        Initializes the FluxControlNet model with embeddings, transformer layers, and optional conditioning blocks.

        Args:
            config (FluxControlNetConfig): Configuration object with model parameters.
        """
        super().__init__(config)
        self.out_channels = config.in_channels
        self.hidden_size = config.hidden_size
        self.patch_size = config.patch_size

        self.pos_embed = EmbedND(dim=self.hidden_size, theta=10000, axes_dim=[16, 56, 56])
        self.img_embed = nn.Linear(config.in_channels, self.hidden_size)
        self.txt_embed = nn.Linear(config.context_dim, self.hidden_size)
        self.timestep_embedding = TimeStepEmbedder(config.model_channels, self.hidden_size)
        self.vector_embedding = MLPEmbedder(in_dim=config.vec_in_dim, hidden_dim=self.hidden_size)
        if config.guidance_embed:
            self.guidance_embedding = (
                MLPEmbedder(in_dim=config.model_channels, hidden_dim=self.hidden_size)
                if config.guidance_embed
                else nn.Identity()
            )
        self.double_blocks = nn.ModuleList(
            [
                MMDiTLayer(
                    config=config,
                    submodules=get_flux_double_transformer_engine_spec().submodules,
                    layer_number=i,
                    context_pre_only=False,
                )
                for i in range(config.num_joint_layers)
            ]
        )

        self.single_blocks = nn.ModuleList(
            [
                FluxSingleTransformerBlock(
                    config=config,
                    submodules=get_flux_single_transformer_engine_spec().submodules,
                    layer_number=i,
                )
                for i in range(config.num_single_layers)
            ]
        )

        # ContolNet Blocks
        self.controlnet_double_blocks = nn.ModuleList()
        for _ in range(config.num_joint_layers):
            self.controlnet_double_blocks.append(
                zero_module(
                    ColumnParallelLinear(
                        self.hidden_size,
                        self.hidden_size,
                        config=config,
                        init_method=nn.init.normal_,
                        gather_output=True,
                    )
                )
            )

        self.controlnet_single_blocks = nn.ModuleList()
        for _ in range(config.num_single_layers):
            self.controlnet_single_blocks.append(
                zero_module(
                    ColumnParallelLinear(
                        self.hidden_size,
                        self.hidden_size,
                        config=config,
                        init_method=nn.init.normal_,
                        gather_output=True,
                    )
                )
            )

        if config.conditioning_embedding_channels is not None:
            self.input_hint_block = ControlNetConditioningEmbedding(
                conditioning_embedding_channels=config.conditioning_embedding_channels,
                block_out_channels=(16, 16, 16, 16),
            )
            self.controlnet_x_embedder = torch.nn.Linear(config.in_channels, self.hidden_size)
        else:
            self.input_hint_block = None
            self.controlnet_x_embedder = zero_module(torch.nn.Linear(config.in_channels, self.hidden_size))

    def load_from_flux_transformer(self, flux):
        """
        Loads pre-trained weights from a Flux Transformer model into the FluxControlNet.

        Args:
            flux (FluxTransformer): A pre-trained Flux Transformer model.
        """
        logging.info("Loading ControlNet layer weights from Flux...")
        self.pos_embed.load_state_dict(flux.pos_embed.state_dict())
        self.img_embed.load_state_dict(flux.img_embed.state_dict())
        self.txt_embed.load_state_dict(flux.txt_embed.state_dict())
        self.timestep_embedding.load_state_dict(flux.timestep_embedding.state_dict())
        self.vector_embedding.load_state_dict(flux.vector_embedding.state_dict())
        self.double_blocks.load_state_dict(flux.double_blocks.state_dict(), strict=False)
        self.single_blocks.load_state_dict(flux.single_blocks.state_dict(), strict=False)

    def get_fp8_context(self):
        "context manager for fp8 recipe"
        # This is first and last 2 for mamba
        if not self.config.fp8:
            fp8_context = nullcontext()
        else:
            import transformer_engine  # To keep out TE dependency when not training in fp8

            if self.config.fp8 == "e4m3":
                fp8_format = transformer_engine.common.recipe.Format.E4M3
            elif self.config.fp8 == "hybrid":
                fp8_format = transformer_engine.common.recipe.Format.HYBRID
            else:
                raise ValueError("E4M3 and HYBRID are the only supported FP8 formats.")

            fp8_recipe = transformer_engine.common.recipe.DelayedScaling(
                margin=self.config.fp8_margin,
                interval=self.config.fp8_interval,
                fp8_format=fp8_format,
                amax_compute_algo=self.config.fp8_amax_compute_algo,
                amax_history_len=self.config.fp8_amax_history_len,
                override_linear_precision=(False, False, not self.config.fp8_wgrad),
            )
            fp8_group = None
            if parallel_state.model_parallel_is_initialized():
                fp8_group = parallel_state.get_amax_reduction_group(with_context_parallel=True)
            fp8_context = transformer_engine.pytorch.fp8_autocast(
                enabled=True, fp8_recipe=fp8_recipe, fp8_group=fp8_group
            )
        return fp8_context

    def forward(
        self,
        img: torch.Tensor,
        controlnet_cond: torch.Tensor,
        txt: torch.Tensor = None,
        y: torch.Tensor = None,
        timesteps: torch.LongTensor = None,
        img_ids: torch.Tensor = None,
        txt_ids: torch.Tensor = None,
        guidance: torch.Tensor = None,
        conditioning_scale: float = 1.0,
    ):
        """
        Forward pass for the FluxControlNet model.

        Args:
            img (torch.Tensor): Input image tensor.
            controlnet_cond (torch.Tensor): Conditioning tensor for ControlNet.
            txt (torch.Tensor, optional): Text embedding tensor. Default is None.
            y (torch.Tensor, optional): Vector embedding tensor. Default is None.
            timesteps (torch.LongTensor, optional): Time step tensor. Default is None.
            img_ids (torch.Tensor, optional): Image IDs. Default is None.
            txt_ids (torch.Tensor, optional): Text IDs. Default is None.
            guidance (torch.Tensor, optional): Guidance tensor. Default is None.
            conditioning_scale (float, optional): Scaling factor for conditioning. Default is 1.0.

        Returns:
            torch.Tensor: The output of the forward pass.
        """
        hidden_states = self.img_embed(img)
        encoder_hidden_states = self.txt_embed(txt)
        if self.input_hint_block is not None:
            controlnet_cond = self.input_hint_block(controlnet_cond)
            batch_size, channels, height_pw, width_pw = controlnet_cond.shape
            height = height_pw // self.config.patch_size
            width = width_pw // self.config.patch_size
            controlnet_cond = controlnet_cond.reshape(
                batch_size, channels, height, self.patch_size, width, self.patch_size
            )
            controlnet_cond = controlnet_cond.permute(0, 2, 4, 1, 3, 5)
            controlnet_cond = controlnet_cond.reshape(batch_size, height * width, -1)

        hidden_states = hidden_states + self.controlnet_x_embedder(controlnet_cond)

        timesteps = timesteps.to(img.dtype) * 1000
        vec_emb = self.timestep_embedding(timesteps)
        if guidance is not None:
            vec_emb = vec_emb + self.guidance_embedding(self.timestep_embedding.time_proj(guidance * 1000))
        vec_emb = vec_emb + self.vector_embedding(y)

        ids = torch.cat((txt_ids, img_ids), dim=1)
        rotary_pos_emb = self.pos_embed(ids)

        double_block_samples = ()
        for id_block, block in enumerate(self.double_blocks):
            with self.get_fp8_context():
                hidden_states, encoder_hidden_states = block(
                    hidden_states=hidden_states,
                    encoder_hidden_states=encoder_hidden_states,
                    rotary_pos_emb=rotary_pos_emb,
                    emb=vec_emb,
                )
                double_block_samples = double_block_samples + (hidden_states,)

        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=0)

        single_block_samples = ()
        for id_block, block in enumerate(self.single_blocks):
            with self.get_fp8_context():
                hidden_states, _ = block(
                    hidden_states=hidden_states,
                    rotary_pos_emb=rotary_pos_emb,
                    emb=vec_emb,
                )
                single_block_samples = single_block_samples + (hidden_states[encoder_hidden_states.shape[0] :, ...],)

        controlnet_double_block_samples = ()
        for double_block_sample, control_block in zip(double_block_samples, self.controlnet_double_blocks):
            double_block_sample, bias = control_block(double_block_sample)
            double_block_sample = double_block_sample + bias if bias else double_block_sample
            controlnet_double_block_samples += (double_block_sample,)

        controlnet_single_block_samples = ()
        for single_block_sample, control_block in zip(single_block_samples, self.controlnet_single_blocks):
            single_block_sample, bias = control_block(single_block_sample)
            single_block_sample = single_block_sample + bias if bias else single_block_sample
            controlnet_single_block_samples += (single_block_sample,)

        controlnet_double_block_samples = [sample * conditioning_scale for sample in controlnet_double_block_samples]
        controlnet_single_block_samples = [sample * conditioning_scale for sample in controlnet_single_block_samples]

        controlnet_double_block_samples = (
            None if len(controlnet_double_block_samples) == 0 else controlnet_double_block_samples
        )
        controlnet_single_block_samples = (
            None if len(controlnet_single_block_samples) == 0 else controlnet_single_block_samples
        )

        return controlnet_double_block_samples, controlnet_single_block_samples


class FluxControlnetForwardWrapper(VisionModule):
    '''
    A wrapper combines flux and flux controlnet forward pass for easier initialization.
    '''

    def __init__(self, flux_config: FluxConfig, flux_controlnet_config: FluxControlNetConfig):
        '''
        Create flux and flux controlnet instances by their config.
        '''
        super().__init__(flux_config)

        self.flux = self.config.configure_model()
        for param in self.flux.parameters():
            param.requires_grad = False

        self.flux_controlnet = FluxControlNet(flux_controlnet_config)
        if flux_controlnet_config.load_from_flux_transformer:
            self.flux_controlnet.load_from_flux_transformer(self.flux)

    def forward(
        self,
        packed_noisy_model_input,
        control_image,
        prompt_embeds,
        pooled_prompt_embeds,
        timesteps,
        latent_image_ids,
        text_ids,
        guidance_vec,
    ):
        '''
        Forward pass for the FluxControlnetForwardWrapper model.
        '''
        # NOTE: This module is wrapped with Fully Sharded Data Parallel (FSDP).
        # To ensure that FSDP can accurately capture the forward and backward
        # processes, it is crucial that the output tensor from this forward
        # function can be used to construct a complete autograd graph.
        controlnet_double_block_samples, controlnet_single_block_samples = self.flux_controlnet(
            img=packed_noisy_model_input,
            controlnet_cond=control_image,
            txt=prompt_embeds,
            y=pooled_prompt_embeds,
            timesteps=timesteps / 1000,
            img_ids=latent_image_ids,
            txt_ids=text_ids,
            guidance=guidance_vec,
        )
        noise_pred = self.flux(
            img=packed_noisy_model_input,
            txt=prompt_embeds,
            y=pooled_prompt_embeds,
            timesteps=timesteps / 1000,
            img_ids=latent_image_ids,
            txt_ids=text_ids,
            guidance=guidance_vec,
            controlnet_double_block_samples=controlnet_double_block_samples,
            controlnet_single_block_samples=controlnet_single_block_samples,
        )

        return noise_pred


class MegatronFluxControlNetModel(MegatronFluxModel):
    """
    Megatron wrapper for flux controlnet model.

    Args:
        flux_params (FluxModelParams): Parameters to configure the Flux model.
        flux_controlnet_config (FluxControlNetConfig): Configuration specific to the FluxControlNet.

    Methods:
        configure_model:
            Configures the model by wrapping the FluxControlNet with the appropriate layers and settings,
            configuring the VAE, scheduler, and text encoders.
        data_step:
            A wrapper around the data-step function specific to FluxControlNet, controlling how data is processed.
        forward:
            Executes a forward pass through FluxControlNet.
        training_step:
            A wrapper step method that calls forward_step with a data batch from data loader.
        forward_step:
            Handles the forward pass specific to training, computing the model's output.
        validation_step:
            Calls inference pipeline with current model weights and save inference result together with the control
            image.
    """

    def __init__(self, flux_params: FluxModelParams, flux_controlnet_config: FluxControlNetConfig, seed: int = 42):
        super().__init__(flux_params, seed=seed)
        self.flux_controlnet_config = flux_controlnet_config
        self.optim.connect(self)

    def configure_model(self):
        '''
        Initialize flux and controlnet modules, vae, scheduler, and text encoders with given configs.
        '''
        if not hasattr(self, "module"):
            self.module = FluxControlnetForwardWrapper(self.config, self.flux_controlnet_config)

            self.configure_vae(self.vae_config)
            self.configure_scheduler()
            self.configure_text_encoders(self.clip_params, self.t5_params)
            # Have to disable requiring grads for those params not getting one, otherwise custom fsdp fails at assert
            # when there is no single layer, encoder_hidden_states related params are not included in computation graph
            for name, param in self.module.named_parameters():
                if self.flux_controlnet_config.num_single_layers == 0:
                    if 'context' in name or 'added' in name:
                        param.requires_grad = False
                # When getting rid of concat, the projection bias in attention and mlp bias are identical
                # So this bias is skipped and not included in the computation graph
                if 'single_blocks' in name and 'self_attention.linear_proj.bias' in name:
                    param.requires_grad = False

    def data_step(self, dataloader_iter):
        '''
        Retrive data batch from dataloader iterator and do necessary processing before feeding into train steps.
        '''
        return self.flux_controlnet_config.data_step_fn(dataloader_iter)

    def forward(self, *args, **kwargs):
        '''
        Calling the controlnet forward pass.
        '''
        # FSDP module -> Bfloat16 module -> ForwardWrapper
        forward_wrapper = self.module.module.module
        return forward_wrapper(*args, **kwargs)

    def training_step(self, batch, batch_idx=None) -> torch.Tensor:
        '''
        A wrapper method takes data batch and returns the results of forward_step.
        '''
        # In mcore the loss-function is part of the forward-pass (when labels are provided)
        return self.forward_step(batch)

    def forward_step(self, batch) -> torch.Tensor:
        '''
        The main forward step function.
        '''
        if self.optim.config.bf16:
            self.autocast_dtype = torch.bfloat16
        elif self.optim.config.fp16:
            self.autocast_dtype = torch.float
        else:
            self.autocast_dtype = torch.float32

        if self.image_precached:
            latents = batch['latents'].cuda(non_blocking=True)
            control_latents = batch['control_latents'].cuda(non_blocking=True)
        else:
            img = batch['images'].cuda(non_blocking=True)
            latents = self.vae.encode(img).to(dtype=self.autocast_dtype)
            hint = batch['hint'].cuda(non_blocking=True)
            control_latents = self.vae.encode(hint).to(dtype=self.autocast_dtype)

        latent_image_ids = self._prepare_latent_image_ids(
            batch_size=latents.shape[0],
            height=latents.shape[2],
            width=latents.shape[3],
            device=latents.device,
            dtype=self.autocast_dtype,
        )

        latents = self._pack_latents(
            latents,
            batch_size=latents.shape[0],
            num_channels_latents=latents.shape[1],
            height=latents.shape[2],
            width=latents.shape[3],
        )

        control_image = self._pack_latents(
            control_latents,
            batch_size=control_latents.shape[0],
            num_channels_latents=control_latents.shape[1],
            height=control_latents.shape[2],
            width=control_latents.shape[3],
        ).transpose(0, 1)

        batch_size = latents.shape[0]
        noise = torch.randn_like(latents, device=latents.device, dtype=latents.dtype)
        u = self.compute_density_for_timestep_sampling(
            "logit_normal",
            batch_size,
        )
        indices = (u * self.scheduler.num_train_timesteps).long()
        timesteps = self.scheduler.timesteps[indices].to(device=latents.device)

        sigmas = self.scheduler.sigmas.to(device=latents.device, dtype=latents.dtype)
        schduler_timesteps = self.scheduler.timesteps.to(device=latents.device)
        step_indices = [(schduler_timesteps == t).nonzero().item() for t in timesteps]
        timesteps = timesteps.to(dtype=latents.dtype)
        sigma = sigmas[step_indices].flatten()
        while len(sigma.shape) < latents.ndim:
            sigma = sigma.unsqueeze(-1)
        packed_noisy_model_input = (1.0 - sigma) * latents + sigma * noise
        packed_noisy_model_input = packed_noisy_model_input.transpose(0, 1)

        if self.config.guidance_embed:
            guidance_vec = torch.full(
                (packed_noisy_model_input.shape[1],),
                self.config.guidance_scale,
                device=latents.device,
                dtype=latents.dtype,
            )
        else:
            guidance_vec = None
        if self.text_precached:
            prompt_embeds = batch['prompt_embeds'].cuda(non_blocking=True).transpose(0, 1)
            pooled_prompt_embeds = batch['pooled_prompt_embeds'].cuda(non_blocking=True)
            text_ids = batch['text_ids'].cuda(non_blocking=True)
        else:
            txt = batch['txt']
            prompt_embeds, pooled_prompt_embeds, text_ids = self.encode_prompt(
                txt, device=latents.device, dtype=latents.dtype
            )

        with torch.cuda.amp.autocast(
            self.autocast_dtype in (torch.half, torch.bfloat16),
            dtype=self.autocast_dtype,
        ):
            noise_pred = self.forward(
                packed_noisy_model_input=packed_noisy_model_input,
                control_image=control_image,
                prompt_embeds=prompt_embeds,
                pooled_prompt_embeds=pooled_prompt_embeds,
                timesteps=timesteps,
                latent_image_ids=latent_image_ids,
                text_ids=text_ids,
                guidance_vec=guidance_vec,
            )

            target = (noise - latents).transpose(0, 1)
            loss = F.mse_loss(noise_pred.float(), target.float(), reduction="mean")
            return loss

    def validation_step(self, batch, batch_idx=None):
        '''
        Initialize flux controlnet pipeline with current model components.

        Saves the inference results together with the hint image to log folder.
        '''
        logging.info("Start validation step")
        from nemo.collections.diffusion.models.flux.pipeline import FluxControlNetInferencePipeline

        pipe = FluxControlNetInferencePipeline(
            params=self.params,
            contorlnet_config=self.flux_controlnet_config,
            flux=self.module.module.module.flux,
            vae=self.vae,
            t5=self.t5,
            clip=self.clip,
            scheduler_steps=self.params.scheduler_steps,
            flux_controlnet=self.module.module.module.flux_controlnet,
        )

        if self.image_precached and self.text_precached:
            latents = batch['latents'].cuda(non_blocking=True)
            control_latents = batch['control_latents'].cuda(non_blocking=True)
            prompt_embeds = batch['prompt_embeds'].cuda(non_blocking=True).transpose(0, 1)
            pooled_prompt_embeds = batch['pooled_prompt_embeds'].cuda(non_blocking=True)
            log_images = pipe(
                latents=latents,
                control_image=control_latents,
                prompt_embeds=prompt_embeds,
                pooled_prompt_embeds=pooled_prompt_embeds,
                height=latents.shape[2] * self.vae_scale_factor,
                width=latents.shape[3] * self.vae_scale_factor,
                num_inference_steps=30,
                num_images_per_prompt=1,
                guidance_scale=7.0,
                dtype=self.autocast_dtype,
                save_to_disk=False,
            )
            log_images[0].save(f"{self.logger.log_dir}/step={self.global_step}_rank{self.local_rank}.png")
        else:
            img = batch['images'].cuda(non_blocking=True)
            hint = batch['hint'].cuda(non_blocking=True)
            text = batch['txt']
            log_images = pipe(
                text,
                control_image=hint,
                num_inference_steps=30,
                num_images_per_prompt=1,
                height=img.shape[2],
                width=img.shape[3],
                guidance_scale=7.0,
                dtype=self.autocast_dtype,
                save_to_disk=False,
            )
            log_images[0].save(f"{self.logger.log_dir}/step={self.global_step}_rank{self.local_rank}_{text}.png")
            hint = pipe.torch_to_numpy(hint)
            hint = pipe.numpy_to_pil(hint)
            hint[0].save(f"{self.logger.log_dir}/step={self.global_step}_rank{self.local_rank}_control.png")
        return torch.tensor([0.0], device=torch.cuda.current_device())
