import torch
import numpy as np
from PIL import Image
from typing import List, Union, Optional, Tuple
from dataclasses import dataclass

from diffusers import DiffusionPipeline
from diffusers.utils import BaseOutput
from tqdm import tqdm

@dataclass
class SdxsPipelineOutput(BaseOutput):
    images: Union[List[Image.Image], np.ndarray]

class SdxsPipeline(DiffusionPipeline):
    def __init__(self, vae, text_encoder, tokenizer, unet, scheduler):
        super().__init__()
        self.register_modules(
            vae=vae, 
            text_encoder=text_encoder, 
            tokenizer=tokenizer, 
            unet=unet, 
            scheduler=scheduler
        )
        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)

    def encode_prompt(self, prompt, negative_prompt, device, dtype):
        """
        Полное соответствие функции encode_texts и get_negative_embedding из трейна.
        """
        def get_single_encode(texts, is_negative=False):
            if texts is None or texts == "":
                # Логика get_negative_embedding из трейна
                hidden_dim = self.text_encoder.config.hidden_size
                
                shape = (1, self.text_encoder.config.max_position_embeddings, hidden_dim)
                # В трейне для негатива: zeros для эмбеддингов и ones для маски
                emb = torch.zeros(shape, dtype=dtype, device=device)
                mask = torch.ones((1, self.text_encoder.config.max_position_embeddings), dtype=torch.int64, device=device)
                return emb, mask

            if isinstance(texts, str):
                texts = [texts]

            with torch.no_grad():
                toks = self.tokenizer(
                    texts,
                    padding="max_length",
                    max_length=self.text_encoder.config.max_position_embeddings,
                    truncation=True,
                    return_tensors="pt"
                ).to(device)

                outputs = self.text_encoder(
                    input_ids=toks.input_ids, 
                    attention_mask=toks.attention_mask,
                    output_hidden_states=True
                )

                # 1. Выбираем нужный слой. 
                # -1 — это последний блок трансформера
                # -2 — это предпоследний (стандарт для большинства современных моделей)
                layer_index = -2
                prompt_embeds = outputs.hidden_states[layer_index]

                # 2. ДОБАВЛЯЕМ ФИНАЛЬНУЮ НОРМАЛИЗАЦИЮ
                # В CLIP после всех блоков стоит слой LayerNorm. 
                # Если мы берем скрытые состояния напрямую, мы "проскакиваем" его.
                # Нужно применить его вручную:
                final_layer_norm = self.text_encoder.text_model.final_layer_norm
                prompt_embeds = final_layer_norm(prompt_embeds)

                return prompt_embeds, toks.attention_mask

        # Получаем эмбеддинги
        pos_embeds, pos_mask = get_single_encode(prompt)
        neg_embeds, neg_mask = get_single_encode(negative_prompt, is_negative=True)

        # Выравнивание батча
        batch_size = pos_embeds.shape[0]
        if neg_embeds.shape[0] != batch_size:
            neg_embeds = neg_embeds.repeat(batch_size, 1, 1)
            neg_mask = neg_mask.repeat(batch_size, 1)

        # Конкатенация для CFG: [Negative, Positive]
        text_embeddings = torch.cat([neg_embeds, pos_embeds], dim=0)
        final_mask = torch.cat([neg_mask, pos_mask], dim=0)

        return text_embeddings.to(dtype=dtype), final_mask.to(dtype=torch.int64)

    @torch.no_grad()
    def __call__(
        self,
        prompt: Union[str, List[str]],
        negative_prompt: Optional[Union[str, List[str]]] = None,
        height: int = 1024,
        width: int = 1024,
        num_inference_steps: int = 40, # Как в трейне n_diffusion_steps
        guidance_scale: float = 4.0,   # Как в трейне
        generator: Optional[torch.Generator] = None,
        output_type: str = "pil",
        return_dict: bool = True,
        **kwargs,
    ):
        device = self.device
        self.vae.to(device)
        vae_scaling_factor = getattr(self.vae.config, "scaling_factor", 1.0)
        vae_shift_factor = getattr(self.vae.config, "shift_factor", 0.0)
        
        # 1. Encode Prompt
        dtype = self.text_encoder.dtype
        text_embeddings, attention_mask = self.encode_prompt(
            prompt, negative_prompt, device, dtype
        )

        # 2. Prepare Latents
        batch_size = 1 if isinstance(prompt, str) else len(prompt)
        latent_channels = self.unet.config.in_channels 
        
        latents = torch.randn(
            (batch_size, latent_channels, height // self.vae_scale_factor, width // self.vae_scale_factor),
            generator=generator,
            device=device,
            dtype=dtype
        )

        # 3. Настройка Flow Matching шедулера
        self.scheduler.set_timesteps(num_inference_steps, device=device)
        timesteps = self.scheduler.timesteps

        # 4. Denoising Loop
        for t in tqdm(timesteps, desc="Sampling"):
            # CFG input
            latent_model_input = torch.cat([latents] * 2) if guidance_scale > 1 else latents
            
            # Flow Matching обычно не требует scale_model_input, 
            # но оставим для совместимости с интерфейсом шедулера
            if hasattr(self.scheduler, "scale_model_input"):
                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)

            # Predict
            model_out = self.unet(
                latent_model_input,
                t,
                encoder_hidden_states=text_embeddings,
                encoder_attention_mask=attention_mask,
                return_dict=False,
            )[0]

            # CFG Logic
            if guidance_scale > 1:
                flow_uncond, flow_cond = model_out.chunk(2)
                model_out = flow_uncond + guidance_scale * (flow_cond - flow_uncond)

            # Step (Flow Matching Euler)
            latents = self.scheduler.step(model_out, t, latents, return_dict=False)[0]

        # 5. Decode
        if output_type == "latent":
            return SdxsPipelineOutput(images=latents)

        latents = latents * vae_scaling_factor + vae_shift_factor
        image = self.vae.decode(latents.to(self.vae.dtype), return_dict=False)[0]
        
        # Пост-процессинг
        image = (image / 2 + 0.5).clamp(0, 1)
        image = image.cpu().permute(0, 2, 3, 1).float().numpy()

        if output_type == "pil":
            image = (image * 255).round().astype("uint8")
            image = [Image.fromarray(img) for img in image]

        if not return_dict:
            return image

        return SdxsPipelineOutput(images=image)