# ==============================================================================
# 1. DEPENDANCES & COMPATIBILITE COLAB / ZERO-GPU
# ==============================================================================
#!pip install -q diffusers transformers accelerate datasets gradio torch torchvision huggingface_hub
try:
    import spaces
except ImportError:
    class MockSpaces:
        def GPU(self, fn=None, duration=None): return fn if fn is not None else lambda f: f
    spaces = MockSpaces()

import os, json, torch, torch.nn as nn, numpy as np, gradio as ui
from PIL import Image
from datasets import load_dataset
from torchvision import transforms
from transformers import CLIPTextModel, CLIPTokenizer
from diffusers import UNet2DConditionModel, AutoencoderKL, DDIMScheduler
from huggingface_hub import HfApi


torch.backends.cuda.matmul.allow_tf32 = True
LOU_UNET, LOU_VAE, LOU_TOKENIZER, LOU_TEXT_ENCODER, LOU_SCHEDULER, LOU_IMAGE_SIZE = None, None, None, None, None, 64

# ==============================================================================
# 2. INITIALISATION ARCHITECTURE LOU DE ZERO
# ==============================================================================
def init_lou_architecture(image_size):
    tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
    text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")
    vae = AutoencoderKL(
        in_channels=3, out_channels=3, latent_channels=4,
        down_block_types=("DownEncoderBlock2D", "DownEncoderBlock2D"),
        up_block_types=("UpDecoderBlock2D", "UpDecoderBlock2D"), block_out_channels=(64, 128)
    )
    unet = UNet2DConditionModel(
        sample_size=image_size // 8, in_channels=4, out_channels=4, layers_per_block=1,
        block_out_channels=(64, 128), down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"),
        up_block_types=("CrossAttnUpBlock2D", "UpBlock2D"), cross_attention_dim=768
    )
    noise_scheduler = DDIMScheduler(num_train_timesteps=1000, beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear")
    return tokenizer, text_encoder, unet, vae, noise_scheduler

# ==============================================================================
# 3. TRAIN ET EXPORT REPO STRUCTURÉ (VAE + UNET COMPLET)
# ==============================================================================
@spaces.GPU
def train_lou_model(repo, token, dataset_path, lr, epochs, size, progress=ui.Progress()):
    global LOU_UNET, LOU_VAE, LOU_TOKENIZER, LOU_TEXT_ENCODER, LOU_SCHEDULER, LOU_IMAGE_SIZE
    LOU_IMAGE_SIZE = int(size)
    device = "cuda" if torch.cuda.is_available() else "cpu"

    LOU_TOKENIZER, LOU_TEXT_ENCODER, LOU_UNET, LOU_VAE, LOU_SCHEDULER = init_lou_architecture(LOU_IMAGE_SIZE)
    LOU_UNET.to(device); LOU_TEXT_ENCODER.to(device); LOU_VAE.to(device)

    dataset = load_dataset(dataset_path, split="train")
    tf = transforms.Compose([transforms.Resize((LOU_IMAGE_SIZE, LOU_IMAGE_SIZE)), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])])
    optimizer = torch.optim.AdamW(list(LOU_UNET.parameters()) + list(LOU_VAE.parameters()), lr=float(lr))

    LOU_UNET.train(); LOU_VAE.train()
    for epoch in range(int(epochs)):
        for batch in progress.tqdm(dataset, desc=f"Epoch {epoch + 1}/{epochs}"):
            optimizer.zero_grad()
            clean = tf(batch["image"].convert("RGB")).unsqueeze(0).to(device)
            latents = LOU_VAE.encode(clean).latent_dist.sample() * 0.18215
            noise = torch.randn_like(latents)
            t = torch.randint(0, 1000, (1,), device=device).long()
            noisy = LOU_SCHEDULER.add_noise(latents, noise, t)
            txt = LOU_TOKENIZER(batch["text"], padding="max_length", truncation=True, return_tensors="pt").input_ids.to(device)
            loss = nn.functional.mse_loss(LOU_UNET(noisy, t, LOU_TEXT_ENCODER(txt).last_hidden_state).sample, noise)
            loss.backward(); optimizer.step()

    base_dir = "./lou_model_repo"
    os.makedirs(base_dir, exist_ok=True)
    LOU_UNET.save_pretrained(os.path.join(base_dir, "unet"))
    LOU_VAE.save_pretrained(os.path.join(base_dir, "vae"))
    LOU_TEXT_ENCODER.save_pretrained(os.path.join(base_dir, "text_encoder"))
    LOU_TOKENIZER.save_pretrained(os.path.join(base_dir, "tokenizer"))
    LOU_SCHEDULER.save_pretrained(os.path.join(base_dir, "scheduler"))

    with open(os.path.join(base_dir, "model_index.json"), "w") as f:
        json.dump({"_class_name": "LouDiffusionArchitecture", "_diffusers_version": "0.21.4", "scheduler": ["diffusers", "DDIMScheduler"], "text_encoder": ["transformers", "CLIPTextModel"], "tokenizer": ["transformers", "CLIPTokenizer"], "unet": ["diffusers", "UNet2DConditionModel"], "vae": ["diffusers", "AutoencoderKL"]}, f, indent=2)

    if token and repo:
        try:
            api = HfApi(); api.create_repo(repo_id=repo, token=token, exist_ok=True)
            api.upload_folder(folder_path=base_dir, repo_id=repo, repo_type="model", token=token)
            return f"✅ Entraîné et synchronisé sur HF via hfapi."
        except Exception as e: return f"⚠️ Structure locale ok, erreur hfapi : {str(e)}"
    return f"✅ Entraîné ! Structure locale créée dans : {base_dir}"

# ==============================================================================
# 4. INFERENCE EN ESPACE LATENT (VITESSE MAXIMALE FP16)
# ==============================================================================
@spaces.GPU
def generate_lou_images(prompt, steps, num_images, seed, cfg, auto):
    global LOU_UNET, LOU_VAE, LOU_TOKENIZER, LOU_TEXT_ENCODER, LOU_SCHEDULER, LOU_IMAGE_SIZE
    if LOU_UNET is None: return [Image.new("RGB", (64, 64), color="gray")], "❌ Lancez d'abord l'entraînement."

    device = "cuda" if torch.cuda.is_available() else "cpu"
    LOU_UNET.to(device=device, dtype=torch.float16); LOU_TEXT_ENCODER.to(device=device, dtype=torch.float16); LOU_VAE.to(device=device, dtype=torch.float16)
    LOU_UNET.eval(); LOU_TEXT_ENCODER.eval(); LOU_VAE.eval()

    if seed >= 0: torch.manual_seed(int(seed)); np.random.seed(int(seed))
    t_in = LOU_TOKENIZER([prompt], padding="max_length", truncation=True, return_tensors="pt").input_ids.to(device)
    u_in = LOU_TOKENIZER([""], padding="max_length", truncation=True, return_tensors="pt").input_ids.to(device)

    with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.float16):
        context = torch.cat([LOU_TEXT_ENCODER(u_in).last_hidden_state, LOU_TEXT_ENCODER(t_in).last_hidden_state])

    generated = []
    LOU_SCHEDULER.set_timesteps(int(steps))
    with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.float16):
        for _ in range(int(num_images)):
            latents = torch.randn((1, 4, LOU_IMAGE_SIZE // 8, LOU_IMAGE_SIZE // 8), device=device, dtype=torch.float16)
            for t in LOU_SCHEDULER.timesteps:
                p_uncond, p_text = LOU_UNET(torch.cat([latents] * 2), t, encoder_hidden_states=context).sample.chunk(2)
                latents = LOU_SCHEDULER.step(p_uncond + cfg * (p_text - p_uncond), t, latents).prev_sample
            img_tensor = (LOU_VAE.decode(latents / 0.18215).sample / 2 + 0.5).clamp(0, 1).squeeze(0)
            generated.append(Image.fromarray((img_tensor.cpu().float().permute(1, 2, 0).numpy() * 255).astype(np.uint8)))

    return generated, f"🎉 Images {LOU_IMAGE_SIZE}x{LOU_IMAGE_SIZE} générées par Lou !"

# ==============================================================================
# 5. CODE GRADIO INTERFACE MINIMALISTE COMPACTE
# ==============================================================================
with ui.Blocks(theme=ui.themes.Default()) as app:
    with ui.Tab("🏋️ Train"):
        with ui.Row():
            repo_i = ui.Textbox(label="Repo HF", placeholder="user/repo"); token_i = ui.Textbox(label="Token HF", type="password")
            ds_i = ui.Textbox(label="Dataset", value="diffusers/pokemon-gpt4-captions"); lr_i = ui.Number(label="LR", value=1e-4)
            ep_i = ui.Slider(1, 50, value=3, step=1, label="Époques"); sz_i = ui.Slider(32, 1024, value=64, step=32, label="Taille Image")
        t_btn = ui.Button("🚀 Lancer l'entraînement de Lou", variant="primary"); t_out = ui.Textbox(label="Statut", lines=1)

    with ui.Tab("🎨 Playground"):
        with ui.Row():
            pr_i = ui.Textbox(label="Prompt", value="a cute pokemon", scale=2); st_i = ui.Slider(5, 100, value=20, step=1, label="Steps")
            num_i = ui.Slider(1, 4, value=1, step=1, label="Images"); sd_i = ui.Number(label="Seed", value=-1)
            gd_i = ui.Slider(1.0, 20.0, value=7.5, step=0.5, label="Guidance"); au_i = ui.Checkbox(label="Autotraining", value=True)
        g_btn = ui.Button("🔮 Générer", variant="secondary"); out_gal = ui.Gallery(label="Outputs", height=200); g_out = ui.Textbox(label="Statut Gen", lines=1)

    t_btn.click(train_lou_model, inputs=[repo_i, token_i, ds_i, lr_i, ep_i, sz_i], outputs=t_out)
    g_btn.click(generate_lou_images, inputs=[pr_i, st_i, num_i, sd_i, gd_i, au_i], outputs=[out_gal, g_out])
app.queue(default_concurrency_limit=4)
app.launch(inline=True, share=True,debug=True)