import torch
from PIL import Image, ImageEnhance, ImageFilter
import numpy as np
import gradio as gr
import os

from model_256 import Generator
from feature_extractor import CodeFeatureExtractor

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = Generator().to(device)
checkpoint = torch.load("models/model_256.pth", map_location=device)
model.load_state_dict(checkpoint["state_dict"])
model.eval()

extractor = CodeFeatureExtractor()

def enhance_image(image: Image.Image, upscale_size: int):
    image = image.resize((upscale_size, upscale_size), Image.Resampling.BICUBIC)
    
    image = image.filter(ImageFilter.GaussianBlur(radius=0.8))
    image = ImageEnhance.Color(image).enhance(1.2)
    image = ImageEnhance.Sharpness(image).enhance(1.1)
    
    return image

def generate_from_code(code_text, upscale_size):
    temp_file = "temp.py"
    with open(temp_file, "w", encoding="utf-8") as f:
        f.write(code_text)
    
    features = extractor.extract_from_file(temp_file)
    if features is None:
        return None

    features_tensor = torch.tensor([features], dtype=torch.float32).to(device)
    with torch.no_grad():
        img_tensor = model(features_tensor).cpu()
    
    img_np = (img_tensor[0].permute(1, 2, 0).numpy() * 255).astype(np.uint8)
    img = Image.fromarray(img_np)

    enhanced = enhance_image(img, upscale_size)
    return enhanced

demo = gr.Interface(
    fn=generate_from_code,
    inputs=[
        gr.Textbox(lines=15, placeholder="Paste your Python or JS code...", label="Code Input"),
        gr.Slider(64, 1024, step=64, value=256, label="Output Size (px)")
    ],
    outputs=gr.Image(type="pil"),
    title="Code-to-Art Generator",
    description="Generates artistic images from code features."
)

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