import os
import torch
import torch.nn as nn
from typing import Optional
from dataclasses import dataclass
from transformers.utils import ModelOutput
from transformers.modeling_utils import PreTrainedModel
from transformers.configuration_utils import PretrainedConfig
from .build import load_sd_model, load_Florence2_model
from .vlv_utils import initiate_time_steps, normalize


class SDConfig(PretrainedConfig):
    """Configuration class for SDModel."""
    model_type = "sd"
    
    def __init__(self, **kwargs):
        super().__init__(**kwargs)


class MLP(nn.Module):
    def __init__(self, input_dim, output_dim):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(input_dim, output_dim),
            nn.GELU(),
            nn.Linear(output_dim, output_dim),
        )
        
    def forward(self, x):
        return self.layers(x)

@dataclass
class SDOutput(ModelOutput):
    loss: Optional[torch.FloatTensor] = None

class SDModel(PreTrainedModel):
    config_class = SDConfig
    
    def __init__(
        self,
        config=None,
        training_args = None,
    ):
        if config is None:
            config = SDConfig()
        super().__init__(config)
        self.training_args = training_args
        if self.training_args.fp32:
            self._dtype = torch.float32
        else:
            self._dtype = torch.bfloat16
        self._device = torch.device(self.training_args.device if hasattr(self.training_args, 'device') else "cuda" if torch.cuda.is_available() else "cpu")
        
        self.vae, self.tokenizer, self.text_encoder, self.unet, self.scheduler = load_sd_model(training_args)
        torch.cuda.empty_cache()
        self.unet.eval()
        self.text_encoder.eval()
        self.model, self.processor = load_Florence2_model(training_args)

        self.unet = self.unet.to(self._dtype).to(device=self._device)
        self.text_encoder = self.text_encoder.to(self._dtype).to_empty(device=self._device)
        self.model = self.model.to(self._dtype).to_empty(device=self._device)
        self.vae = self.vae.to(torch.float32).to_empty(device=self._device)

        self.batch_size = self.training_args.batch_size 

        hidden_dim = 1024 
        self.language_proj = nn.Sequential(
            nn.Linear(1024, hidden_dim, dtype=self._dtype),
            nn.GELU(),
            nn.Linear(hidden_dim, 1024, dtype=self._dtype)
        ).to_empty(device=self._device)
        for param in self.language_proj.parameters():
            param.requires_grad = True

        self.num_queries = self.training_args.learnable_token_length
        self.query_embed = nn.Parameter(torch.randn(1, self.num_queries, 1024, dtype=self._dtype))
        self.query_embed.requires_grad = True
        
        self.unet.enable_gradient_checkpointing()

    def _unet_pred_noise(self, x_start, t, noise, context):
        t = t.to(dtype=torch.long)
        
        dtype = self.unet.dtype
        x_start = x_start.to(dtype)
        noise = noise.to(dtype)
        context = context.to(dtype)
        
        nt = t.shape[0]
        noised_latent = self.scheduler.add_noise(x_start, noise, t)
        
        pred_noise = self.unet(
            noised_latent, 
            t, 
            encoder_hidden_states=context.expand(nt, -1, -1)
        ).sample

        return pred_noise
    
    def generate_images(self, images):
        batch_size = self.training_args.eval_batch_size
        prompt = ["<MORE_DETAILED_CAPTION>"] * batch_size
        inputs = self.processor(text=prompt, images=images, return_tensors="pt").to(self._device).to(self._dtype)
        
        if inputs["input_ids"] is not None:
            inputs_embeds = self.model.language_model.get_input_embeddings()(inputs["input_ids"]).to(self._dtype)
        if inputs["pixel_values"] is not None:
            image_features = self.model._encode_image(inputs["pixel_values"]).to(self._dtype)
            inputs_embeds, attention_mask = self.model._merge_input_ids_with_image_features(image_features, inputs_embeds)
        if inputs_embeds is not None:
            attention_mask = attention_mask.to(inputs_embeds.dtype)
        encoder_outputs = self.model.language_model.model.encoder(
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
            output_hidden_states=True,
            return_dict=True
        )
        
        decoder_input_embeds = self.query_embed.expand(batch_size, -1, -1)
        decoder_attention_mask = torch.ones(
            (batch_size, self.num_queries), 
            dtype=self._dtype, 
            device=self._device
        )
        
        encoder_hidden_states = encoder_outputs.last_hidden_state.to(self._dtype)
        decoder_input_embeds = decoder_input_embeds.to(self._dtype)
        attention_mask = attention_mask.to(self._dtype)
        
        decoder_outputs = self.model.language_model.model.decoder(
            inputs_embeds=decoder_input_embeds,
            attention_mask=decoder_attention_mask,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=attention_mask,
            output_hidden_states=True,
            return_dict=True
        )
        
        last_decoder_hidden_state = decoder_outputs.last_hidden_state
        conditional_context = self.language_proj(last_decoder_hidden_state)

        un_token = self.tokenizer("", padding="max_length", truncation=True,max_length=77, return_tensors="pt").input_ids.to(self._device)
        un_context_embeddings = self.text_encoder(un_token).last_hidden_state
        un_context_embeddings = un_context_embeddings.expand(batch_size, -1, -1)
        if self.training_args.use_text_encoder:
            context_embeddings = self.text_encoder(
                inputs_embeds=conditional_context.to(self._dtype)
            ).last_hidden_state

        latent_shape = (batch_size, 4, self.training_args.image_size // 8, self.training_args.image_size // 8)
        latents = torch.randn(latent_shape, device=self._device, dtype=self._dtype)

        scheduler = self.scheduler
        scheduler.set_timesteps(self.training_args.num_inference_steps)
        with torch.no_grad():
            for t in scheduler.timesteps:
                latent_model_input = torch.cat([latents, latents], dim=0)
                latent_model_input = scheduler.scale_model_input(latent_model_input, t)

                combined_embeddings = torch.cat([un_context_embeddings, context_embeddings], dim=0).to(self._dtype)
                noise_pred = self.unet(
                    latent_model_input, t, encoder_hidden_states=combined_embeddings
                )[0]

                noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2, dim=0)
                noise_pred = noise_pred_uncond + self.training_args.guidance_scale * (noise_pred_cond - noise_pred_uncond)

                latents = scheduler.step(noise_pred, t, latents)[0]

        scaled_latents = latents / 0.18215
        with torch.no_grad():
            decoded_latents = self.vae.decode(scaled_latents.to(torch.float32))[0]

        return decoded_latents
    
    def get_conditional_context(self, images, batch_size=None):
        if batch_size is None:
            batch_size = self.batch_size
        prompt = ["<MORE_DETAILED_CAPTION>"] * batch_size
        inputs = self.processor(text=prompt, images=images, return_tensors="pt").to(self._device).to(self._dtype)
        
        if inputs["input_ids"] is not None:
            inputs_embeds = self.model.language_model.get_input_embeddings()(inputs["input_ids"]).to(self._dtype)
        if inputs["pixel_values"] is not None:
            image_features = self.model._encode_image(inputs["pixel_values"]).to(self._dtype)
            inputs_embeds, attention_mask = self.model._merge_input_ids_with_image_features(image_features, inputs_embeds)
        if inputs_embeds is not None:
            attention_mask = attention_mask.to(inputs_embeds.dtype)
        encoder_outputs = self.model.language_model.model.encoder(
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
            output_hidden_states=True,
            return_dict=True
        )
        
        decoder_input_embeds = self.query_embed.expand(batch_size, -1, -1)
        decoder_attention_mask = torch.ones(
            (batch_size, self.num_queries), 
            dtype=self._dtype, 
            device=self._device
        )
        
        encoder_hidden_states = encoder_outputs.last_hidden_state.to(self._dtype)
        decoder_input_embeds = decoder_input_embeds.to(self._dtype)
        attention_mask = attention_mask.to(self._dtype)
        
        decoder_outputs = self.model.language_model.model.decoder(
            inputs_embeds=decoder_input_embeds,
            attention_mask=decoder_attention_mask,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=attention_mask,
            output_hidden_states=True,
            return_dict=True
        )
        
        last_decoder_hidden_state = decoder_outputs.last_hidden_state
        return last_decoder_hidden_state
    
    def forward(
        self,
        image=None,
        filename=None,
        **kwargs,
    ) -> SDOutput:
        images_for_language_model = image
        normalize_images = normalize(image, rescale=True)
        x0=self.vae.encode(normalize_images.to(torch.float32)).latent_dist.sample()
        latent = x0 * 0.18215
        
        total_timestep = self.scheduler.num_train_timesteps

        timesteps = initiate_time_steps(0, total_timestep, self.batch_size, self.training_args).long()
        timesteps = timesteps.to(self._device)
        c, h, w = latent.shape[1:]
        if not self.training_args.use_same_noise_among_timesteps:
            noise = torch.randn((self.batch_size, c, h, w), device=self._device, dtype=self._dtype)
        else:
            noise = torch.randn((1, c, h, w), device=self._device, dtype=self._dtype)
            noise = noise.repeat(self.batch_size, 1, 1, 1)

        conditional_context = self.get_conditional_context(images_for_language_model)
        conditional_context = self.language_proj(conditional_context)

        if self.training_args.use_text_encoder:
            text_encoder_output = self.text_encoder(input_ids=None, inputs_embeds=conditional_context.to(self._dtype))
            pred_noise = self._unet_pred_noise(x_start=latent, t=timesteps, noise=noise, context=text_encoder_output.last_hidden_state.to(self._dtype)).to(self._dtype)
        else:
            pred_noise = self._unet_pred_noise(x_start=latent, t=timesteps, noise=noise, context=conditional_context.to(self._dtype)).to(self._dtype)
        
        if self.training_args.loss == "l1":
            loss = torch.nn.functional.l1_loss(pred_noise, noise)
        else:
            loss = torch.nn.functional.mse_loss(pred_noise, noise)
            
        return SDOutput(loss=loss)