import os
import requests
import tarfile
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
import json
import math
from tqdm import tqdm
from transformers import BertTokenizer, BertModel
import gradio as gr

# Configuration
class Config:
    device = "cuda" if torch.cuda.is_available() else "cpu"
    image_size = 64
    batch_size = 32
    num_epochs = 50
    learning_rate = 1e-4
    timesteps = 1000
    text_embed_dim = 768
    num_images_options = [1, 4, 6]
    
    # URLs for COCO dataset download
    coco_images_url = "http://images.cocodataset.org/zips/train2017.zip"
    coco_annotations_url = "http://images.cocodataset.org/annotations/annotations_trainval2017.zip"
    data_dir = "./coco_data"
    images_dir = os.path.join(data_dir, "train2017")
    annotations_path = os.path.join(data_dir, "annotations/instances_train2017.json")
    
    def __init__(self):
        os.makedirs(self.data_dir, exist_ok=True)

config = Config()

# Download COCO dataset
def download_and_extract_coco():
    if os.path.exists(config.images_dir) and os.path.exists(config.annotations_path):
        print("COCO dataset already downloaded")
        return
    
    print("Downloading COCO dataset...")
    
    # Download images
    images_zip_path = os.path.join(config.data_dir, "train2017.zip")
    if not os.path.exists(images_zip_path):
        response = requests.get(config.coco_images_url, stream=True)
        with open(images_zip_path, "wb") as f:
            for chunk in tqdm(response.iter_content(chunk_size=1024)):
                if chunk:
                    f.write(chunk)
    
    # Download annotations
    annotations_zip_path = os.path.join(config.data_dir, "annotations_trainval2017.zip")
    if not os.path.exists(annotations_zip_path):
        response = requests.get(config.coco_annotations_url, stream=True)
        with open(annotations_zip_path, "wb") as f:
            for chunk in tqdm(response.iter_content(chunk_size=1024)):
                if chunk:
                    f.write(chunk)
    
    # Extract files
    print("Extracting images...")
    with tarfile.open(images_zip_path, "r:zip") as tar:
        tar.extractall(config.data_dir)
    
    print("Extracting annotations...")
    with tarfile.open(annotations_zip_path, "r:zip") as tar:
        tar.extractall(config.data_dir)
    
    print("COCO dataset ready")

download_and_extract_coco()

# Text model
class TextEncoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
        self.model = BertModel.from_pretrained('bert-base-uncased')
        for param in self.model.parameters():
            param.requires_grad = False
    
    def forward(self, texts):
        inputs = self.tokenizer(texts, return_tensors="pt", padding=True, truncation=True, max_length=64)
        inputs = {k: v.to(config.device) for k, v in inputs.items()}
        outputs = self.model(**inputs)
        return outputs.last_hidden_state[:, 0, :]

text_encoder = TextEncoder().to(config.device)

# Diffusion model
class ConditionalUNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
        self.down1 = DownBlock(64, 128)
        self.down2 = DownBlock(128, 256)
        
        self.text_proj = nn.Linear(config.text_embed_dim, 256)
        self.merge = nn.Linear(256 + 256, 256)
        
        self.up1 = UpBlock(256, 128)
        self.up2 = UpBlock(128, 64)
        self.final = nn.Conv2d(64, 3, kernel_size=3, padding=1)
        
    def forward(self, x, t, text_emb):
        x1 = F.relu(self.conv1(x))
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        
        text_emb = self.text_proj(text_emb)
        text_emb = text_emb.unsqueeze(-1).unsqueeze(-1)
        text_emb = text_emb.expand(-1, -1, x3.size(2), x3.size(3))
        
        x = torch.cat([x3, text_emb], dim=1)
        b, c, h, w = x.shape
        x = x.permute(0, 2, 3, 1).reshape(b*h*w, c)
        x = self.merge(x)
        x = x.reshape(b, h, w, 256).permute(0, 3, 1, 2)
        
        x = self.up1(x)
        x = self.up2(x)
        return self.final(x)

class DownBlock(nn.Module):
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(),
            nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
    
    def forward(self, x):
        return self.conv(x)

class UpBlock(nn.Module):
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        self.conv = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(),
            nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU()
        )
    
    def forward(self, x):
        x = self.up(x)
        return self.conv(x)

# Diffusion process
betas = linear_beta_schedule(config.timesteps).to(config.device)
alphas = 1. - betas
alphas_cumprod = torch.cumprod(alphas, dim=0)
sqrt_alphas_cumprod = torch.sqrt(alphas_cumprod)
sqrt_one_minus_alphas_cumprod = torch.sqrt(1. - alphas_cumprod)

def linear_beta_schedule(timesteps):
    beta_start = 0.0001
    beta_end = 0.02
    return torch.linspace(beta_start, beta_end, timesteps)

def forward_diffusion_sample(x_0, t, device=config.device):
    noise = torch.randn_like(x_0)
    sqrt_alphas_cumprod_t = sqrt_alphas_cumprod[t].view(-1, 1, 1, 1)
    sqrt_one_minus_alphas_cumprod_t = sqrt_one_minus_alphas_cumprod[t].view(-1, 1, 1, 1)
    return sqrt_alphas_cumprod_t * x_0 + sqrt_one_minus_alphas_cumprod_t * noise, noise

# COCO Dataset
class CocoDataset(Dataset):
    def __init__(self, root_dir, annotations_file, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        
        with open(annotations_file, 'r') as f:
            data = json.load(f)
        
        self.images = []
        self.captions = []
        
        image_id_to_captions = {}
        for ann in data['annotations']:
            if ann['image_id'] not in image_id_to_captions:
                image_id_to_captions[ann['image_id']] = []
            image_id_to_captions[ann['image_id']].append(ann['caption'])
        
        for img in data['images']:
            if img['id'] in image_id_to_captions:
                self.images.append(img)
                self.captions.append(image_id_to_captions[img['id']][0])
    
    def __len__(self):
        return len(self.images)
    
    def __getitem__(self, idx):
        img_path = os.path.join(self.root_dir, self.images[idx]['file_name'])
        image = Image.open(img_path).convert('RGB')
        caption = self.captions[idx]
        
        if self.transform:
            image = self.transform(image)
        
        return image, caption

# Transformations
transform = transforms.Compose([
    transforms.Resize((config.image_size, config.image_size)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])

# Model initialization
model = ConditionalUNet().to(config.device)
optimizer = torch.optim.Adam(model.parameters(), lr=config.learning_rate)

# Training
def train():
    dataset = CocoDataset(config.images_dir, config.annotations_path, transform)
    dataloader = DataLoader(dataset, batch_size=config.batch_size, shuffle=True)
    
    for epoch in range(config.num_epochs):
        for batch_idx, (images, captions) in enumerate(tqdm(dataloader)):
            images = images.to(config.device)
            
            # Get text embeddings
            text_emb = text_encoder(captions)
            
            # Sample random timesteps
            t = torch.randint(0, config.timesteps, (images.size(0),), device=config.device)
            
            # Forward diffusion
            x_noisy, noise = forward_diffusion_sample(images, t)
            
            # Predict noise
            pred_noise = model(x_noisy, t, text_emb)
            
            # Loss and backpropagation
            loss = F.mse_loss(pred_noise, noise)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
            if batch_idx % 100 == 0:
                print(f"Epoch {epoch}, Batch {batch_idx}, Loss: {loss.item():.4f}")
        
        # Save model
        torch.save(model.state_dict(), f"model_epoch_{epoch}.pth")

# Generation
@torch.no_grad()
def generate(prompt, num_images=1):
    model.eval()
    num_images = int(num_images)
    
    text_emb = text_encoder([prompt]*num_images)
    x = torch.randn((num_images, 3, config.image_size, config.image_size)).to(config.device)
    
    for t in reversed(range(config.timesteps)):
        t_tensor = torch.full((num_images,), t, device=config.device)
        pred_noise = model(x, t_tensor, text_emb)
        alpha_t = alphas[t].view(1, 1, 1, 1)
        alpha_cumprod_t = alphas_cumprod[t].view(1, 1, 1, 1)
        beta_t = betas[t].view(1, 1, 1, 1)
        
        if t > 0:
            noise = torch.randn_like(x)
        else:
            noise = torch.zeros_like(x)
        
        x = (1 / torch.sqrt(alpha_t)) * (
            x - ((1 - alpha_t) / torch.sqrt(1 - alpha_cumprod_t)) * pred_noise
        ) + torch.sqrt(beta_t) * noise
    
    x = torch.clamp(x, -1, 1)
    x = (x + 1) / 2
    
    images = []
    for img in x:
        img = transforms.ToPILImage()(img.cpu())
        images.append(img)
    
    return images

# GUI
def generate_and_display(prompt, num_images):
    images = generate(prompt, num_images)
    
    fig, axes = plt.subplots(1, len(images), figsize=(5*len(images), 5))
    if len(images) == 1:
        axes.imshow(images[0])
        axes.axis('off')
    else:
        for ax, img in zip(axes, images):
            ax.imshow(img)
            ax.axis('off')
    plt.tight_layout()
    return fig

with gr.Blocks() as demo:
    gr.Markdown("## GPUDiff-V1: diffussion powerful image generator!")
    with gr.Row():
        prompt_input = gr.Textbox(label="Prompt", placeholder="Enter image description...")
        num_select = gr.Dropdown(choices=config.num_images_options, value=1, label="Number of images")
    generate_btn = gr.Button("Generate")
    output = gr.Plot()
    
    generate_btn.click(
        fn=generate_and_display,
        inputs=[prompt_input, num_select],
        outputs=output
    )

if __name__ == "__main__":
    
    train()
    
    demo.launch()