
from diffusers import FluxTransformer2DModel, ModelMixin
from diffusers.pipelines import FluxFillPipeline
import lightning as L
from peft import LoraConfig, get_peft_model_state_dict
import prodigyopt
from safetensors.torch import load_file
import torch
import torch.nn.functional as F

from .transformer import tranformer_forward
from .condition import Condition
from .pipeline_tools import encode_images, prepare_text_input


class OminiModelLORAFIll(L.LightningModule):
    def __init__(
        self,
        fill_model: FluxTransformer2DModel,
        flux_pipe_id: str,
        lora_path: str = None,
        reuse_lora_path: str = None,
        lora_config: dict = None,
        device: str = "cuda",
        dtype: torch.dtype = torch.bfloat16,
        model_config: dict = {},
        optimizer_config: dict = None,
        gradient_checkpointing: bool = False,
        odm_loss_config: dict = None,
        ocr_loss_config: dict = None,
        byt5_encoder_config: dict = None,
    ):
        # Initialize the LightningModule
        super().__init__()
        self.model_config = model_config
        self.optimizer_config = optimizer_config

        # Load the Flux pipeline
        
        self.transformer = fill_model
        self.transformer.gradient_checkpointing = gradient_checkpointing
        self.transformer.train()

        self.flux_pipe: FluxFillPipeline = (
            FluxFillPipeline.from_pretrained(flux_pipe_id, transformer=self.transformer, text_encoder=None, text_encoder_2=None).to(dtype=dtype).to(device)
        )

        # Freeze the Flux pipeline
        # self.flux_pipe.text_encoder.requires_grad_(False).eval()
        # self.flux_pipe.text_encoder_2.requires_grad_(False).eval()
        # self.flux_pipe.vae.requires_grad_(False).eval()

        self.vae_scale_factor = 8

        # Initialize LoRA layers
        self.lora_layers = self.init_lora(lora_path, lora_config)
        # reuse the weight
        if reuse_lora_path is not None:
            print(f"reuse the lora path: {reuse_lora_path}")
            state_dict = load_file(reuse_lora_path)
            state_dict1 = {x.replace('lora_A', 'lora_A.default').replace('lora_B', 'lora_B.default').replace('transformer.', ''): v for x, v in state_dict.items()}
            self.transformer.load_state_dict(state_dict1, strict=False)

        self.to(device).to(dtype)

    def init_lora(self, lora_path: str, lora_config: dict):
        assert lora_path or lora_config
        if lora_path:
            # TODO: Implement this
            raise NotImplementedError
        else:
            self.transformer.add_adapter(LoraConfig(**lora_config))
            # TODO: Check if this is correct (p.requires_grad)
            lora_layers = filter(
                lambda p: p.requires_grad, self.transformer.parameters()
            )
        return list(lora_layers)

    def save_lora(self, path: str):
        FluxFillPipeline.save_lora_weights(
            save_directory=path,
            transformer_lora_layers=get_peft_model_state_dict(self.transformer),
            safe_serialization=True,
        )

    def configure_optimizers(self):
        # Freeze the transformer
        self.transformer.requires_grad_(False)
        opt_config = self.optimizer_config

        # Set the trainable parameters
        self.trainable_params = self.lora_layers

        # Unfreeze trainable parameters
        for p in self.trainable_params:
            p.requires_grad_(True)

        # Initialize the optimizer
        if opt_config["type"] == "AdamW":
            optimizer = torch.optim.AdamW(self.trainable_params, **opt_config["params"])
        elif opt_config["type"] == "Prodigy":
            optimizer = prodigyopt.Prodigy(
                self.trainable_params,
                **opt_config["params"],
            )
        elif opt_config["type"] == "SGD":
            optimizer = torch.optim.SGD(self.trainable_params, **opt_config["params"])
        else:
            raise NotImplementedError

        return optimizer

    def training_step(self, batch, batch_idx):
        step_loss = self.step(batch)
        self.log_loss = (
            step_loss.item()
            if not hasattr(self, "log_loss")
            else self.log_loss * 0.95 + step_loss.item() * 0.05
        )
        return step_loss
