import torch
from PIL import Image
import torchvision.transforms as transforms
from util import *
from cnn1 import UNet
from transformers import CLIPModel, AutoTokenizer
from torchvision.transforms.functional import adjust_contrast
import torch.nn.functional as F

def load_models(weights_path):
    
    clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(DEVICE)
    tokenizer = AutoTokenizer.from_pretrained("openai/clip-vit-base-patch32")
    vgg19 = models.vgg19(pretrained=True).features.to(DEVICE)
    for param in vgg19.parameters():
        param.requires_grad_(False)
    
    cnn_model = UNet(text_dim=512).to(DEVICE)
    cnn_model.load_state_dict(torch.load(weights_path, map_location=DEVICE))
    
    return clip_model, tokenizer, vgg19, cnn_model

def reset_cnn_model():
    """Reset CNN model weights for each new image"""
    global cnn_model
    if cnn_model is not None:
        # Reinitialize the model
        for layer in cnn_model.modules():
            if isinstance(layer, (torch.nn.Conv2d, torch.nn.Linear)):
                torch.nn.init.xavier_uniform_(layer.weight)
                if layer.bias is not None:
                    torch.nn.init.zeros_(layer.bias)

def style_transfer(img, prompt, weights_path, source="a Photo"):
    """
    Apply style transfer to an uploaded image
    
    Args:
        image: PIL Image uploaded by user
        prompt: Text style prompt
        num_steps: Number of optimization steps (fewer = faster, more = better quality)
    
    Returns:
        Stylized PIL Image
    """
    if img is None:
        return None
    
    img = load_image(img).to(DEVICE)
    
    # Load models
    clip_model, tokenizer, vgg19, cnn_model = load_models(weights_path)

    cnn_model.eval()

    with torch.no_grad():
        edited_text = prompt_ensemble(prompt)
        text_tokens = tokenizer(edited_text, padding=True, return_tensors="pt").to(DEVICE)
        text_features = clip_model.get_text_features(**text_tokens)
        text_features = text_features.mean(axis=0, keepdim=True)
        text_features /= text_features.norm(dim=-1, keepdim=True)
        
    with torch.inference_mode():
        target = cnn_model(img, text_embedding=text_features).to(DEVICE)
    
    # Post-process output
    output_image = target.clone().detach()
    output_image = torch.clamp(output_image, 0, 1)
    output_image = adjust_contrast(output_image, 1.5)
    
    # Convert tensor back to PIL Image
    output_image = output_image.squeeze(0).cpu()
    output_image = transforms.ToPILImage()(output_image)
    
    return output_image