"""
Eruku Model - Styled Handwritten Text Recognition

This module implements the Eruku model for autoregressive styled text image generation.
Based on the papers:
- "Zero-Shot Styled Text Image Generation, but Make It Autoregressive" (CVPR 2025)
- "Autoregressive Styled Text Image Generation, but Make it Reliable" (WACV 2026)
"""

import torch
import torch.nn as nn
from typing import Optional, Tuple, List, Union
from transformers import PreTrainedModel, T5ForConditionalGeneration, T5Config, AutoTokenizer
from diffusers import AutoencoderKL
from einops import rearrange, repeat
from torch.nn.utils.rnn import pad_sequence
from torchvision.transforms import Normalize
from PIL import Image
import numpy as np

from .configuration_eruku import ErukuConfig


# Number of special tokens: SOG, EOG, IMG
SPECIAL_TOKEN_COUNT = 3


def pad_images(images: List[torch.Tensor], padding_value: float = 1.0) -> torch.Tensor:
    """Pad a list of images to the same width."""
    images = [rearrange(img, 'c h w -> w c h') for img in images]
    padded = rearrange(pad_sequence(images, padding_value=padding_value), 'w b c h -> b c h w')
    return padded.contiguous()


class ErukuPreTrainedModel(PreTrainedModel):
    """
    Base class for Eruku models.
    """
    config_class = ErukuConfig
    base_model_prefix = "eruku"
    supports_gradient_checkpointing = True
    
    def _init_weights(self, module):
        """Initialize weights - handled by sub-components."""
        pass


class ErukuForConditionalGeneration(ErukuPreTrainedModel):
    """
    Eruku model for conditional styled text image generation.
    
    The model takes a style image (handwritten/typewritten text sample),
    optional style text (transcription of the style image), and generation
    text (text to render), and produces an image of the generation text
    in the style of the reference image.
    
    Example usage:
    ```python
    from transformers import AutoModel
    from PIL import Image
    import torch
    
    # Load model
    model = AutoModel.from_pretrained("blowing-up-groundhogs/eruku", trust_remote_code=True)
    model.eval()
    
    # Generate handwriting
    result = model.generate_handwriting(
        style_image=Image.open("style.png"),
        style_text="Hello",  # optional - text in style image
        gen_text="World",    # text to generate
    )
    result.save("output.png")
    ```
    """
    
    def __init__(self, config: ErukuConfig):
        super().__init__(config)
        self.config = config
        
        # Character-level tokenizer
        self.tokenizer = AutoTokenizer.from_pretrained(config.tokenizer_name_or_path)
        self.tokenizer.add_tokens(["<sog>"])
        
        # T5 backbone
        t5_config = T5Config.from_pretrained(config.t5_name_or_path)
        t5_config.vocab_size = len(self.tokenizer)
        self.T5 = T5ForConditionalGeneration(t5_config)
        self.T5.lm_head = nn.Identity()
        
        # Image normalization
        self.normalize = Normalize(0.5, 0.5)
        
        # Special token embeddings
        self.sos = nn.Embedding(1, t5_config.d_model)
        self.sog = nn.Embedding(1, t5_config.d_model)
        self.eog = nn.Embedding(1, t5_config.d_model)
        
        # VAE for image encoding/decoding
        self.vae = AutoencoderKL.from_pretrained(config.vae_name_or_path)
        self._freeze_module(self.vae)
        
        # Projection layers
        vae_dim = config.vae_latent_dim * config.channels * config.slices_per_query
        self.query_emb = nn.Linear(vae_dim, t5_config.d_model)
        self.t5_to_vae = nn.Linear(t5_config.d_model, vae_dim)
        self.t5_to_special = nn.Linear(t5_config.d_model, SPECIAL_TOKEN_COUNT)
        
        # Unconditional embedding for CFG
        self.uncond_embedding = nn.Embedding(1, t5_config.d_model)
        
        # CFG configuration
        self.drop_text = False
        self.drop_img = False
        
        # Einops rearrangements
        self.z_rearrange = lambda x: rearrange(x, 'b w (q c h) -> b c h (w q)', 
                                                c=config.channels, q=config.slices_per_query)
        
        self.post_init()
    
    def _freeze_module(self, module: nn.Module):
        """Freeze all parameters in a module."""
        module.eval()
        for param in module.parameters():
            param.requires_grad = False
    
    def _img_encode(self, img: torch.Tensor) -> torch.Tensor:
        """Encode image to VAE latent space."""
        img = self.normalize(img)
        img = img.contiguous()
        return self.vae.encode(img.float()).latent_dist.sample()
    
    @torch.no_grad()
    def get_model_inputs(
        self,
        style_img: List[torch.Tensor],
        style_len: Union[int, List[int]],
        max_img_len: int = 1024 * 1024
    ) -> dict:
        """
        Prepare model inputs from style images.
        
        Args:
            style_img: List of style image tensors [C, H, W]
            style_len: Width(s) of style images
            max_img_len: Maximum image length in pixels
            
        Returns:
            Dictionary with decoder_inputs_embeds
        """
        bs = len(style_img)
        decoder_inputs_embeds_list = []
        
        # Pad images to same width
        style_img_padded = pad_images([el.to(self.T5.device) for el in style_img])
        style_img_embeds = self._img_encode(style_img_padded)
        
        for el in range(bs):
            if isinstance(style_len, int):
                sl = style_len
            else:
                sl = int(style_len[el])
            
            # Ensure width is within bounds
            sl = max(64, min(sl, style_img_embeds.shape[-1] * 8))
            
            # Style image embeddings + SOG marker
            sample_embeds = torch.cat([
                style_img_embeds[el, :, :, :sl // 8],
                torch.ones(1, 8, 1).to(self.T5.device),  # SOG placeholder
            ], dim=-1)
            
            sample_embeds = rearrange(sample_embeds, 'c h w -> w (h c)', h=8, c=1)
            decoder_inputs_embeds_list.append(sample_embeds)
        
        decoder_inputs_embeds = pad_sequence(
            decoder_inputs_embeds_list, 
            padding_value=1, 
            batch_first=True
        )[:, :max_img_len // 8]
        
        return {'decoder_inputs_embeds': decoder_inputs_embeds}
    
    @torch.inference_mode()
    def generate(
        self,
        decoder_inputs_embeds_vae: torch.Tensor,
        style_text: List[str],
        gen_text: List[str],
        cfg_scale: float = 1.25,
        max_new_tokens: int = 512
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Generate styled text image autoregressively.
        
        Args:
            decoder_inputs_embeds_vae: VAE embeddings of style image
            style_text: List of style text strings (can be empty)
            gen_text: List of generation text strings
            cfg_scale: Classifier-free guidance scale (1.25 recommended)
            max_new_tokens: Maximum tokens to generate
            
        Returns:
            Tuple of (generated_image, special_sequence)
        """
        # Encode text
        encoded_text = self.tokenizer(
            [f"{style}<sog>{gen}" for style, gen in zip(style_text, gen_text)],
            padding=True,
            return_tensors="pt"
        )
        text_input_ids = encoded_text['input_ids'].to(self.T5.device)
        text_mask = encoded_text['attention_mask'].to(self.T5.device)
        
        # Initialize generation
        sog = repeat(self.sog.weight, '1 d -> b 1 d', b=1)
        sos = repeat(self.sos.weight, '1 d -> b 1 d', b=1)
        
        z_sequence = [decoder_inputs_embeds_vae]
        special_sequence = torch.ones(decoder_inputs_embeds_vae.size(1)) * 3
        
        # Build initial decoder inputs
        decoder_inputs_embeds = self.query_emb(torch.cat(z_sequence, dim=1))
        if len(style_text[0]) != 0:
            decoder_inputs_embeds = torch.cat([sos, decoder_inputs_embeds], dim=1)
        else:
            decoder_inputs_embeds = torch.cat([sos, decoder_inputs_embeds, sog], dim=1)
            vae_latent = self.t5_to_vae(sog)
            special_sequence = torch.cat([special_sequence, torch.zeros(1)])
            z_sequence.append(vae_latent)
        
        # Autoregressive generation
        for i in range(max_new_tokens):
            if cfg_scale != 1.0:
                # Classifier-free guidance
                conditional_text_embeds = self.T5.shared(text_input_ids)
                if self.drop_text:
                    unconditional_text_embeds = self.uncond_embedding.weight.expand_as(conditional_text_embeds)
                else:
                    unconditional_text_embeds = conditional_text_embeds
                
                if self.drop_img:
                    unconditional_decoder_embeds = self.uncond_embedding.weight.expand_as(decoder_inputs_embeds)
                else:
                    unconditional_decoder_embeds = decoder_inputs_embeds
                
                output_uncond = self.T5(
                    inputs_embeds=unconditional_text_embeds,
                    attention_mask=text_mask,
                    decoder_inputs_embeds=unconditional_decoder_embeds
                ).logits[:, -1:]
                
                output_cond = self.T5(
                    input_ids=text_input_ids,
                    attention_mask=text_mask,
                    decoder_inputs_embeds=decoder_inputs_embeds
                ).logits[:, -1:]
                
                output = output_uncond + (output_cond - output_uncond) * cfg_scale
            else:
                output = self.T5(
                    input_ids=text_input_ids,
                    attention_mask=text_mask,
                    decoder_inputs_embeds=decoder_inputs_embeds
                ).logits[:, -1:]
            
            # Predict special token
            special_prediction = self.t5_to_special(output)
            predicted_special = torch.argmax(special_prediction, dim=-1).item()
            
            if predicted_special == 0:  # SOG
                decoder_inputs_embeds = torch.cat([decoder_inputs_embeds, sog], dim=1)
                vae_latent = self.t5_to_vae(output)
                special_sequence = torch.cat([special_sequence, torch.zeros(1)])
            elif predicted_special == 1:  # EOG - stop generation
                special_sequence = torch.cat([special_sequence, torch.ones(1)])
                vae_latent = self.t5_to_vae(output)
                z_sequence.append(vae_latent)
                break
            else:  # IMG token
                vae_latent = self.t5_to_vae(output)
                decoder_inputs_embeds = torch.cat([
                    decoder_inputs_embeds, 
                    self.query_emb(vae_latent)
                ], dim=1)
                special_sequence = torch.cat([special_sequence, torch.ones(1) * 2])
            
            z_sequence.append(vae_latent)
        
        # Decode to image
        z_sequence = [el.to(self.vae.device) for el in z_sequence]
        z_sequence = torch.cat(z_sequence, dim=1)
        z_sequence = self.z_rearrange(z_sequence)

        z_sequence = z_sequence[:,:,:,decoder_inputs_embeds_vae.size(1)+1:]

        img = torch.clamp(self.vae.decode(z_sequence).sample, -1, 1)
        
        return img, special_sequence.to(self.T5.device)
    
    def generate_handwriting(
        self,
        style_image: Image.Image,
        gen_text: str,
        style_text: str = "",
        cfg_scale: float = 1.25,
        max_new_tokens: int = 512,
        device: Optional[str] = None
    ) -> Image.Image:
        """
        High-level API for generating handwriting.
        
        This is the recommended entry point for inference.
        
        Args:
            style_image: PIL Image containing handwriting style reference
            gen_text: Text to generate in the style
            style_text: Optional transcription of text in style_image
            cfg_scale: Classifier-free guidance scale (default: 1.25)
            max_new_tokens: Maximum generation length
            device: Device to use (auto-detected if None)
            
        Returns:
            PIL Image of generated handwriting
        """
        import torchvision.transforms as T
        
        if device is None:
            device = next(self.parameters()).device
        
        # Preprocess style image
        style_img = style_image.convert('RGB')
        
        # Resize to height 64 maintaining aspect ratio
        width, height = style_img.size
        new_width = int(64 * width / height)
        style_img = style_img.resize((new_width, 64), Image.LANCZOS)
        
        # Convert to tensor
        style_tensor = T.ToTensor()(style_img).to(device)
        style_len = style_tensor.shape[-1]
        
        # Get model inputs
        inputs = self.get_model_inputs(
            style_img=[style_tensor],
            style_len=style_len,
            max_img_len=1024 * 1024
        )
        
        # Generate
        output_img, _ = self.generate(
            decoder_inputs_embeds_vae=inputs['decoder_inputs_embeds'],
            style_text=[style_text],
            gen_text=[gen_text],
            cfg_scale=cfg_scale,
            max_new_tokens=max_new_tokens
        )
        
        # Crop out the style image part (keep only generated portion)
        style_width_latent = style_len // 8 + 1  # +1 for SOG token
        # output_img = output_img[:, :, :, style_width_latent * 8:]
        
        # Trim whitespace
        output_img = self._trim_white(output_img)
        
        # Convert to PIL
        output_img = (torch.clamp(output_img, -1, 1) + 1) * 127.5
        output_img = output_img.byte().squeeze().cpu().numpy()
        
        if len(output_img.shape) == 2:
            return Image.fromarray(output_img, mode='L')
        elif output_img.shape[0] == 3:
            output_img = np.transpose(output_img, (1, 2, 0))
            return Image.fromarray(output_img, mode='RGB')
        else:
            return Image.fromarray(output_img[0], mode='L')
    
    @staticmethod
    def _trim_white(img: torch.Tensor, threshold: float = 0.9, padding: int = 8) -> torch.Tensor:
        """Trim white margins from generated image."""
        start_idx, end_idx = 0, img.size(-1)
        vertical_min = img[0, 0].min(-2).values.tolist()
        
        # Skip initial non-white columns
        for v in vertical_min:
            if v >= threshold:
                break
            start_idx += 1
        
        # Skip initial white columns
        for v in vertical_min:
            if v < threshold:
                break
            start_idx += 1
        
        # Skip trailing white columns
        for v in vertical_min[::-1]:
            if v < threshold:
                break
            end_idx -= 1
        
        start_idx = max(start_idx - padding, 0)
        end_idx = min(end_idx + padding, img.size(-1))
        
        if start_idx >= end_idx:
            return img
        
        return img[..., start_idx:end_idx]
    
    def forward(self, **kwargs):
        """Forward pass - mainly for training compatibility."""
        raise NotImplementedError(
            "Direct forward() is not supported. Use generate_handwriting() for inference."
        )


# Register for AutoModel
ErukuConfig.register_for_auto_class()
ErukuForConditionalGeneration.register_for_auto_class("AutoModel")