import gradio as gr
import torch
from torch import nn
from transformers import CLIPVisionModel, CLIPImageProcessor
from PIL import Image, ImageEnhance, ImageOps

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

# -----------------------------
# Aesthetic Scoring Model
# -----------------------------
class AestheticScorer(nn.Module):
    def __init__(self):
        super().__init__()
        self.vision_model = CLIPVisionModel.from_pretrained("openai/clip-vit-base-patch16")
        self.mlp = nn.Sequential(
            nn.Linear(self.vision_model.config.hidden_size, 512),
            nn.ReLU(),
            nn.Linear(512, 1)
        )

    def forward(self, pixel_values):
        outputs = self.vision_model(pixel_values=pixel_values)
        pooled = outputs.pooler_output
        return self.mlp(pooled)


processor = CLIPImageProcessor.from_pretrained("openai/clip-vit-base-patch16")
ae_model = AestheticScorer().to(device)
ae_model.eval()

# -----------------------------
# Cinematic Enhancement
# -----------------------------
def apply_cinematic(image: Image.Image) -> Image.Image:
    img = image.convert("RGB")

    # slight letterboxing for cinematic aspect ratio feel
    w, h = img.size
    bar_height = int(h * 0.07)  # 7% black bars
    new_img = Image.new("RGB", (w, h), (0, 0, 0))
    new_img.paste(img, (0, 0))

    # adjust colors for cinematic mood
    enhancer = ImageEnhance.Color(img)
    img = enhancer.enhance(1.2)  # richer colors

    enhancer = ImageEnhance.Contrast(img)
    img = enhancer.enhance(1.15)  # more contrast

    enhancer = ImageEnhance.Brightness(img)
    img = enhancer.enhance(0.95)  # slightly darker

    # add subtle warm tint (teal & orange style)
    r, g, b = img.split()
    r = r.point(lambda i: i * 1.05)
    b = b.point(lambda i: i * 0.95)
    img = Image.merge("RGB", (r, g, b))

    # add black bars
    new_img.paste(img, (0, 0))
    draw = ImageOps.expand(img, border=(0, bar_height, 0, bar_height), fill="black")
    return draw


# -----------------------------
# Scoring
# -----------------------------
def aesthetic_score(image: Image.Image) -> float:
    inputs = processor(images=image, return_tensors="pt").to(device)
    with torch.no_grad():
        score = ae_model(inputs['pixel_values']).cpu().item()
    return round(float(score), 2)


# -----------------------------
# Pipeline
# -----------------------------
def process(image):
    # score original
    original_score = aesthetic_score(image)

    # cinematic version
    cinematic_img = apply_cinematic(image)

    # score cinematic
    cinematic_score = aesthetic_score(cinematic_img)

    # ranking suggestion
    suggestions = sorted(
        [("Original", original_score), ("Cinematic", cinematic_score)],
        key=lambda x: x[1],
        reverse=True,
    )

    result_text = "📊 Aesthetic Ranking:\n"
    for rank, (name, score) in enumerate(suggestions, 1):
        result_text += f"{rank}. {name} — Score: {score}\n"

    return cinematic_img, result_text


# -----------------------------
# Gradio UI
# -----------------------------
with gr.Blocks() as demo:
    gr.Markdown("## 🎬 Cinematic Photo Enhancer")
    with gr.Row():
        inp = gr.Image(type="pil", label="Upload your photo")
        out_img = gr.Image(type="pil", label="Cinematic Output")
    out_txt = gr.Textbox(label="Suggestions & Scores")

    run_btn = gr.Button("Enhance Photo")
    run_btn.click(process, inputs=inp, outputs=[out_img, out_txt])

demo.launch()