import subprocess
subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)

import os
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["CUDA_LAUNCH_BLOCKING"] = "1"

import os.path as osp
import time
import argparse
import shutil
import random
from pathlib import Path
from typing import List
import json

import cv2
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
import PIL.Image as PImage
from torchvision.transforms.functional import to_tensor
from transformers import AutoTokenizer, T5EncoderModel
from huggingface_hub import hf_hub_download
import gradio as gr
import spaces

from models.infinity import Infinity
from models.basic import *
from utils.dynamic_resolution import dynamic_resolution_h_w, h_div_w_templates
from gradio_client import Client

torch._dynamo.config.cache_size_limit = 64
client = Client("Qwen/Qwen2.5-72B-Instruct")

# Define a function to download weights if not present
def download_infinity_weights(weights_path):
    try:
        model_file = weights_path / 'infinity_2b_reg.pth'
        if not model_file.exists():
            hf_hub_download(repo_id="FoundationVision/Infinity", filename="infinity_2b_reg.pth", local_dir=str(weights_path))
        
        vae_file = weights_path / 'infinity_vae_d32reg.pth'
        if not vae_file.exists():
            hf_hub_download(repo_id="FoundationVision/Infinity", filename="infinity_vae_d32reg.pth", local_dir=str(weights_path))
        
    except Exception as e:
        print(f"Error downloading weights: {e}")

def encode_prompt(text_tokenizer, text_encoder, prompt):
    print(f'prompt={prompt}')
    captions = [prompt]
    tokens = text_tokenizer(text=captions, max_length=512, padding='max_length', truncation=True, return_tensors='pt')  # todo: put this into dataset
    input_ids = tokens.input_ids.cuda(non_blocking=True) if torch.cuda.is_available() else tokens.input_ids
    mask = tokens.attention_mask.cuda(non_blocking=True) if torch.cuda.is_available() else tokens.attention_mask
    text_features = text_encoder(input_ids=input_ids, attention_mask=mask)['last_hidden_state'].float()
    lens: List[int] = mask.sum(dim=-1).tolist()
    cu_seqlens_k = F.pad(mask.sum(dim=-1).to(dtype=torch.int32).cumsum_(0), (1, 0))
    Ltext = max(lens)
    kv_compact = []
    for len_i, feat_i in zip(lens, text_features.unbind(0)):
        kv_compact.append(feat_i[:len_i])
    kv_compact = torch.cat(kv_compact, dim=0)
    text_cond_tuple = (kv_compact, lens, cu_seqlens_k, Ltext)
    return text_cond_tuple

def save_slim_model(infinity_model_path, save_file=None, device='cpu', key='gpt_fsdp'):
    print('[Save slim model]')
    full_ckpt = torch.load(infinity_model_path, map_location=device)
    infinity_slim = full_ckpt['trainer'][key]
    # ema_state_dict = cpu_d['trainer'].get('gpt_ema_fsdp', state_dict)
    if not save_file:
        save_file = osp.splitext(infinity_model_path)[0] + '-slim.pth'
    print(f'Save to {save_file}')
    torch.save(infinity_slim, save_file)
    print('[Save slim model] done')
    return save_file

def load_tokenizer(t5_path='google/flan-t5-xl'):
    """
    Load and configure the T5 tokenizer and encoder with optimizations.
    """
    try:
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        bf16_supported = device.type == 'cuda' and torch.cuda.is_bf16_supported()
        dtype = torch.bfloat16 if bf16_supported else torch.float32

        tokenizer = AutoTokenizer.from_pretrained(
            t5_path,
            legacy=True,
            model_max_length=512,
            use_fast=True,
        )

        if device.type == 'cuda':
            torch.cuda.empty_cache()

        encoder = T5EncoderModel.from_pretrained(
            t5_path,
            torch_dtype=dtype,
        )

        encoder.eval().requires_grad_(False).to(device)

        if device.type == 'cuda' and not bf16_supported:
            encoder.half()

        return tokenizer, encoder

    except Exception as e:
        print(f"Error loading tokenizer/encoder: {str(e)}")
        raise RuntimeError("Failed to initialize text models") from e

def load_infinity(
    rope2d_each_sa_layer, 
    rope2d_normalized_by_hw, 
    use_scale_schedule_embedding, 
    pn, 
    use_bit_label, 
    add_lvl_embeding_only_first_block, 
    model_path='', 
    scale_schedule=None, 
    vae=None, 
    device=None,  # Make device optional
    model_kwargs=None,
    text_channels=2048,
    apply_spatial_patchify=0,
    use_flex_attn=False,
    bf16=True,
):
    print('[Loading Infinity]')
    
    # Set device if not provided
    if device is None:
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f'Using device: {device}')

    # Set autocast dtype based on bf16 and device support
    if bf16 and device == 'cuda' and torch.cuda.is_bf16_supported():
        autocast_dtype = torch.bfloat16
    else:
        autocast_dtype = torch.float32
        bf16 = False  # Disable bf16 if not supported

    text_maxlen = 512
    torch.cuda.empty_cache()

    with torch.amp.autocast(device_type=device, dtype=autocast_dtype), torch.no_grad():
        infinity_test: Infinity = Infinity(
            vae_local=vae, text_channels=text_channels, text_maxlen=text_maxlen,
            shared_aln=True, raw_scale_schedule=scale_schedule,
            checkpointing='full-block',
            customized_flash_attn=False,
            fused_norm=True,
            pad_to_multiplier=128,
            use_flex_attn=use_flex_attn,
            add_lvl_embeding_only_first_block=add_lvl_embeding_only_first_block,
            use_bit_label=use_bit_label,
            rope2d_each_sa_layer=rope2d_each_sa_layer,
            rope2d_normalized_by_hw=rope2d_normalized_by_hw,
            pn=pn,
            apply_spatial_patchify=apply_spatial_patchify,
            inference_mode=True,
            train_h_div_w_list=[1.0],
            **model_kwargs,
        ).to(device)

        print(f'[you selected Infinity with {model_kwargs=}] model size: {sum(p.numel() for p in infinity_test.parameters())/1e9:.2f}B, bf16={bf16}')

        if bf16:
            for block in infinity_test.unregistered_blocks:
                block.bfloat16()

        infinity_test.eval()
        infinity_test.requires_grad_(False)

        print('[Load Infinity weights]')
        state_dict = torch.load(model_path, map_location=device)
        print(infinity_test.load_state_dict(state_dict))
        
        # # Initialize random number generator on the correct device
        # infinity_test.rng = torch.Generator(device=device)
    
    return infinity_test

def transform(pil_img: PImage.Image, tgt_h: int, tgt_w: int) -> torch.Tensor:
    """
    Transform a PIL image to a tensor with target dimensions while preserving aspect ratio.
    
    Args:
        pil_img: PIL Image to transform
        tgt_h: Target height
        tgt_w: Target width
        
    Returns:
        torch.Tensor: Normalized tensor image in range [-1, 1]
    """
    if not isinstance(pil_img, PImage.Image):
        raise TypeError("Input must be a PIL Image")
    
    if tgt_h <= 0 or tgt_w <= 0:
        raise ValueError("Target dimensions must be positive")
        
    # Calculate resize dimensions preserving aspect ratio
    width, height = pil_img.size
    scale = min(tgt_w / width, tgt_h / height)
    new_width = int(width * scale)
    new_height = int(height * scale)
    
    # Resize using LANCZOS for best quality
    pil_img = pil_img.resize((new_width, new_height), resample=PImage.LANCZOS)
    
    # Create center crop
    arr = np.array(pil_img, dtype=np.uint8)
    
    # Calculate crop coordinates
    y1 = max(0, (new_height - tgt_h) // 2)
    x1 = max(0, (new_width - tgt_w) // 2)
    y2 = y1 + tgt_h
    x2 = x1 + tgt_w
    
    # Crop and convert to tensor
    arr = arr[y1:y2, x1:x2]
    
    # Convert to normalized tensor in one step
    return torch.from_numpy(arr.transpose(2, 0, 1)).float().div_(127.5).sub_(1)

def joint_vi_vae_encode_decode(
    vae: 'VAEModel',  # Type hint would be more specific with actual VAE class
    image_path: str | Path,
    scale_schedule: List[tuple],
    device: torch.device | str,
    tgt_h: int,
    tgt_w: int
) -> tuple[np.ndarray, np.ndarray, torch.Tensor]:
    """
    Encode and decode an image using a VAE model with joint visual-infinity processing.

    Args:
        vae: The VAE model instance
        image_path: Path to input image
        scale_schedule: List of scale tuples for processing
        device: Target device for computation
        tgt_h: Target height for the image
        tgt_w: Target width for the image

    Returns:
        tuple containing:
        - Original image as numpy array (uint8)
        - Reconstructed image as numpy array (uint8)
        - Bit indices tensor
        
    Raises:
        FileNotFoundError: If image file doesn't exist
        RuntimeError: If VAE processing fails
    """
    try:
        # Validate input path
        if not Path(image_path).exists():
            raise FileNotFoundError(f"Image not found at {image_path}")
            
        # Load and preprocess image
        pil_image = Image.open(image_path).convert('RGB')
        inp = transform(pil_image, tgt_h, tgt_w)
        inp = inp.unsqueeze(0).to(device)

        # Normalize scale schedule
        scale_schedule = [(s[0], s[1], s[2]) for s in scale_schedule]

        # Decide whether to use CPU or GPU
        device = 'cuda' if torch.cuda.is_available() else 'cpu'

        # Time the encoding/decoding operations
        with torch.amp.autocast(device, dtype=torch.bfloat16):
            encode_start = time.perf_counter()
            h, z, _, all_bit_indices, _, _ = vae.encode(
                inp, 
                scale_schedule=scale_schedule
            )
            encode_time = time.perf_counter() - encode_start

            decode_start = time.perf_counter()
            recons_img = vae.decode(z)[0]
            decode_time = time.perf_counter() - decode_start

        # Process reconstruction
        if recons_img.dim() == 4:
            recons_img = recons_img.squeeze(1)

        # Log performance metrics
        print(f'VAE encode: {encode_time:.2f}s, decode: {decode_time:.2f}s')
        print(f'Reconstruction shape: {recons_img.shape}, z shape: {z.shape}')

        # Convert to numpy arrays efficiently
        recons_img = (recons_img.add(1).div(2)
                     .permute(1, 2, 0)
                     .mul(255)
                     .cpu()
                     .numpy()
                     .astype(np.uint8))
        
        gt_img = (inp[0].add(1).div(2)
                 .permute(1, 2, 0)
                 .mul(255)
                 .cpu()
                 .numpy()
                 .astype(np.uint8))

        return gt_img, recons_img, all_bit_indices

    except Exception as e:
        print(f"Error in VAE processing: {str(e)}")
        raise RuntimeError("VAE processing failed") from e

def load_visual_tokenizer(args):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    # load vae
    if args.vae_type in [16,18,20,24,32,64]:
        from models.bsq_vae.vae import vae_model
        schedule_mode = "dynamic"
        codebook_dim = args.vae_type
        codebook_size = 2**codebook_dim
        if args.apply_spatial_patchify:
            patch_size = 8
            encoder_ch_mult=[1, 2, 4, 4]
            decoder_ch_mult=[1, 2, 4, 4]
        else:
            patch_size = 16
            encoder_ch_mult=[1, 2, 4, 4, 4]
            decoder_ch_mult=[1, 2, 4, 4, 4]
        vae = vae_model(args.vae_path, schedule_mode, codebook_dim, codebook_size, patch_size=patch_size, 
                        encoder_ch_mult=encoder_ch_mult, decoder_ch_mult=decoder_ch_mult, test_mode=True).to(device)
    else:
        raise ValueError(f'vae_type={args.vae_type} not supported')
    return vae

def load_transformer(vae, args):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model_path = args.model_path

    if args.checkpoint_type == 'torch':
        slim_model_path = model_path
        print(f'Loading checkpoint from {slim_model_path}')
    else:
        raise ValueError(f"Unsupported checkpoint_type: {args.checkpoint_type}")

    model_configs = {
        'infinity_2b': dict(depth=32, embed_dim=2048, num_heads=16, drop_path_rate=0.1, mlp_ratio=4, block_chunks=8),
        'infinity_layer12': dict(depth=12, embed_dim=768, num_heads=8, drop_path_rate=0.1, mlp_ratio=4, block_chunks=4),
        'infinity_layer16': dict(depth=16, embed_dim=1152, num_heads=12, drop_path_rate=0.1, mlp_ratio=4, block_chunks=4),
        'infinity_layer24': dict(depth=24, embed_dim=1536, num_heads=16, drop_path_rate=0.1, mlp_ratio=4, block_chunks=4),
        'infinity_layer32': dict(depth=32, embed_dim=2080, num_heads=20, drop_path_rate=0.1, mlp_ratio=4, block_chunks=4),
        'infinity_layer40': dict(depth=40, embed_dim=2688, num_heads=24, drop_path_rate=0.1, mlp_ratio=4, block_chunks=4),
        'infinity_layer48': dict(depth=48, embed_dim=3360, num_heads=28, drop_path_rate=0.1, mlp_ratio=4, block_chunks=4),
    }

    kwargs_model = model_configs.get(args.model_type)
    if kwargs_model is None:
        raise ValueError(f"Unsupported model_type: {args.model_type}")

    infinity = load_infinity(
        rope2d_each_sa_layer=args.rope2d_each_sa_layer,
        rope2d_normalized_by_hw=args.rope2d_normalized_by_hw,
        use_scale_schedule_embedding=args.use_scale_schedule_embedding,
        pn=args.pn,
        use_bit_label=args.use_bit_label,
        add_lvl_embeding_only_first_block=args.add_lvl_embeding_only_first_block,
        model_path=slim_model_path,
        scale_schedule=None,
        vae=vae,
        device=device,
        model_kwargs=kwargs_model,
        text_channels=args.text_channels,
        apply_spatial_patchify=args.apply_spatial_patchify,
        use_flex_attn=args.use_flex_attn,
        bf16=args.bf16,
    )
    return infinity

def enhance_prompt(prompt):
    SYSTEM = """You are part of a team of bots that creates images. You work with an assistant bot that will draw anything you say.  

When given a user prompt, your role is to transform it into a creative, detailed, and vivid image description that focuses on visual and sensory features. Avoid directly referencing specific real-world people, places, or cultural knowledge unless explicitly requested by the user.  

### Guidelines for Generating the Output:  

1. **Output Format:**  
   Your response must be in the following dictionary format:  
   ```json
   {
     "prompt": "<enhanced image description>",
     "cfg": <cfg value>
   }
   ```  

2. **Enhancing the "prompt" field:**  
   - Use your creativity to expand short or vague prompts into highly detailed, visually rich descriptions.  
   - Focus on describing visual and sensory elements, such as colors, textures, shapes, lighting, and emotions.  
   - Avoid including known real-world information unless the user explicitly requests it. Instead, describe features that evoke the essence or appearance of the scene or subject.  
   - For particularly long user prompts (over 50 words), output them directly without refinement.  
   - Image descriptions must remain between 8-512 words. Any excess text will be ignored.  
   - If the user's request involves rendering specific text in the image, enclose that text in single quotation marks and prefix it with "the text".  

3. **Determining the "cfg" field:**  
   - If the image to be generated is likely to feature a clear face, set `"cfg": 1`.  
   - If the image does not prominently feature a face, set `"cfg": 3`.  

4. **Examples of Enhanced Prompts:**  
   - **User prompt:** "a tree"  
     **Enhanced prompt:** "A towering tree with a textured bark of intricate ridges and grooves stands under a pale blue sky. Its sprawling branches create an umbrella of rich, deep green foliage, with a few golden leaves scattered, catching the sunlight like tiny stars."  
     **Cfg:** `3`  

   - **User prompt:** "a person reading"  
     **Enhanced prompt:** "A figure sits on a cozy armchair, illuminated by the soft, warm glow of a nearby lamp. Their posture is relaxed, and their hands gently hold an open book. Shadows dance across their thoughtful expression, while the fabric of their clothing appears textured and soft, with subtle folds."  
     **Cfg:** `1`  

5. **Your Output:**  
   Always return a single dictionary containing both `"prompt"` and `"cfg"` fields. Avoid any additional commentary or explanations.  

Don't write anything except the dictionary in the output. (Don't start with ```)
"""
    result = client.predict(
            query=prompt,
            history=[],
            system=SYSTEM,
            api_name="/model_chat"
    )

    dict_of_inputs = json.loads(result[1][-1][-1])
    print(dict_of_inputs)

    return gr.update(value=dict_of_inputs["prompt"]), gr.update(value=float(dict_of_inputs['cfg']))

# Set up paths
weights_path = Path(__file__).parent / 'weights'
weights_path.mkdir(exist_ok=True)
download_infinity_weights(weights_path)

# Device setup
dtype = torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float32
print(f"Using dtype: {dtype}")

# Define args
args = argparse.Namespace(
    pn='1M',
    model_path=str(weights_path / 'infinity_2b_reg.pth'),
    cfg_insertion_layer=0,
    vae_type=32,
    vae_path=str(weights_path / 'infinity_vae_d32reg.pth'),
    add_lvl_embeding_only_first_block=1,
    use_bit_label=1,
    model_type='infinity_2b',
    rope2d_each_sa_layer=1,
    rope2d_normalized_by_hw=2,
    use_scale_schedule_embedding=0,
    sampling_per_bits=1,
    text_channels=2048,
    apply_spatial_patchify=0,
    h_div_w_template=1.000,
    use_flex_attn=0,
    cache_dir='/dev/shm',
    checkpoint_type='torch',
    seed=0,
    bf16=1 if dtype == torch.bfloat16 else 0,
    save_file='tmp.jpg',
    enable_model_cache=False,
)

# Load models
print(f"VRAM before forward pass: {torch.cuda.memory_allocated() / 1024 ** 2:.2f} MB")
text_tokenizer, text_encoder = load_tokenizer(t5_path="google/flan-t5-xl")
print(f"VRAM before forward pass: {torch.cuda.memory_allocated() / 1024 ** 2:.2f} MB")
vae = load_visual_tokenizer(args)
print(f"VRAM before forward pass: {torch.cuda.memory_allocated() / 1024 ** 2:.2f} MB")
infinity = load_transformer(vae, args)
print(f"VRAM before forward pass: {torch.cuda.memory_allocated() / 1024 ** 2:.2f} MB")

# Define the image generation function
@spaces.GPU
def generate_image(prompt, cfg, tau, h_div_w, seed):
    args.prompt = prompt
    args.cfg = cfg
    args.tau = tau
    args.h_div_w = h_div_w
    args.seed = seed
    
    # Find the closest h_div_w_template
    h_div_w_template_ = h_div_w_templates[np.argmin(np.abs(h_div_w_templates - h_div_w))]
    
    # Get scale_schedule based on h_div_w_template_
    scale_schedule = dynamic_resolution_h_w[h_div_w_template_][args.pn]['scales']
    scale_schedule = [(1, h, w) for (_, h, w) in scale_schedule]
    
    # Encode the prompt
    text_cond_tuple = encode_prompt(text_tokenizer, text_encoder, prompt)
    
    # Set device if not provided
    device = 'cuda' if torch.cuda.is_available() else 'cpu'

    # Set autocast dtype based on bf16 and device support
    if device == 'cuda' and torch.cuda.is_bf16_supported():
        autocast_dtype = torch.bfloat16
    else:
        autocast_dtype = torch.float32

    torch.cuda.empty_cache()

    with torch.amp.autocast(device_type=device, dtype=autocast_dtype), torch.no_grad():
        infinity.rng = torch.Generator(device=device)
        _, _, img_list = infinity.autoregressive_infer_cfg(
            vae=vae,
            scale_schedule=scale_schedule,
            label_B_or_BLT=text_cond_tuple, g_seed=seed,
            B=1, negative_label_B_or_BLT=None, force_gt_Bhw=None,
            cfg_sc=3, cfg_list=[cfg] * len(scale_schedule), tau_list=[tau] * len(scale_schedule), top_k=900, top_p=0.97,
            returns_vemb=1, ratio_Bl1=None, gumbel=0, norm_cfg=False,
            cfg_exp_k=0.0, cfg_insertion_layer=[args.cfg_insertion_layer],
            vae_type=args.vae_type, softmax_merge_topk=-1,
            ret_img=True, trunk_scale=1000,
            gt_leak=0, gt_ls_Bl=None, inference_mode=True,
            sampling_per_bits=args.sampling_per_bits,
        )
        infinity.rng = torch.Generator(device="cpu")
    
    img = img_list[0]
    image = img.cpu().numpy()
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    image = np.uint8(image)
    
    return image


markdown_description = """### Instructions:

1. Enter a detailed prompt with rich visual features or use the "Enhance Prompt" button to generate a more detailed description.
2. Adjust the "CFG" and "Tau" sliders to control the strength and randomness of the output.
3. Use the "Aspect Ratio" slider to set the aspect ratio of the generated image.
4. Click the "Generate Image" button to create the image based on your prompt.

Arxiv Paper:
[Infinity: Scaling Bitwise AutoRegressive Modeling for High-Resolution Image Synthesis](https://arxiv.org/abs/2412.04431).
"""
html_header = """<div style="text-align: center; margin-bottom: 20px;">
    <h1>Infinity Image Generator by <a href="https://github.com/FoundationVision/Infinity" target="_blank" rel="noopener noreferrer">FoundationVision</a></h1>
    <p style="font-size: 14px; color: #888;">This is not the official implementation from the main developers!</p>
</div>"""
with gr.Blocks() as demo:
    gr.HTML(html_header)
    gr.Markdown(markdown_description)
    
    with gr.Row():
        with gr.Column():
            # Prompt Settings
            gr.Markdown("### Prompt Settings")
            prompt = gr.Textbox(label="Prompt", value="alien spaceship enterprise", placeholder="Enter your prompt here...")
            enhance_prompt_button = gr.Button("Enhance Prompt", variant="secondary")

            # Image Settings
            gr.Markdown("### Image Settings")
            with gr.Row():
                cfg = gr.Slider(label="CFG (Classifier-Free Guidance)", minimum=1, maximum=10, step=0.5, value=3, info="Controls the strength of the prompt.")
                tau = gr.Slider(label="Tau (Temperature)", minimum=0.1, maximum=1.0, step=0.1, value=0.5, info="Controls the randomness of the output.")
            with gr.Row():
                h_div_w = gr.Slider(label="Aspect Ratio (Height/Width)", minimum=0.5, maximum=2.0, step=0.1, value=1.0, info="Set the aspect ratio of the generated image.")
                seed = gr.Number(label="Seed", value=random.randint(0, 10000), info="Set a seed for reproducibility.")
            
            # Generate Button
            generate_button = gr.Button("Generate Image", variant="primary")
        
        with gr.Column():
            # Output Section
            gr.Markdown("### Generated Image")
            output_image = gr.Image(label="Generated Image", type="pil")
    
    # Error Handling
    error_message = gr.Textbox(label="Error Message", visible=False)
    
    # Link the enhance prompt button to the prompt enhancement function
    enhance_prompt_button.click(
        enhance_prompt,
        inputs=prompt,
        outputs=[prompt, cfg],
    )

    # Link the generate button to the image generation function
    generate_button.click(
        generate_image,
        inputs=[prompt, cfg, tau, h_div_w, seed],
        outputs=output_image
    )

# Launch the Gradio app
demo.launch(ssr_mode=False)