import torch
import torch.nn as nn
from torchvision import transforms, models
import open_clip
from PIL import Image, ImageFilter
import numpy as np
import os

# --- 1. SETUP & DEVICE ---
DEVICE = torch.device("mps") if torch.backends.mps.is_available() else torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {DEVICE}")

# --- 2. LOAD MODELS ---

# A. Load openclip (ViT-L-14)
print("Loading openclip...")
openclip_model, _, openclip_preprocess = open_clip.create_model_and_transforms(
    'ViT-L-14', pretrained='datacomp_xl_s13b_b90k'
)
openclip_model.to(DEVICE)

# Define your openclip Forensic Head Architecture (matches your training)
class openclipHead(nn.Module):
    def __init__(self, input_dim):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, 512),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(512, 1)
        )
    def forward(self, x): return self.net(x)

# Load openclip Weights
openclip_head = openclipHead(input_dim=768).to(DEVICE)
openclip_head.load_state_dict(torch.load('models/openclip_forensic_head.pth', map_location=DEVICE))
openclip_head.eval()

# B. Load ConvNeXt-Base
print("Loading ConvNeXt...")
cn_backbone = models.convnext_base(weights=None) # Architecture only
cn_backbone.to(DEVICE)
cn_backbone.eval()

class ConvNextHead(nn.Module):
    def __init__(self, input_dim):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, 512),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(512, 1)
        )
    def forward(self, x): return self.net(x)

cn_head = ConvNextHead(input_dim=1024).to(DEVICE)
cn_head.load_state_dict(torch.load('models/convnext_forensic_head.pth', map_location=DEVICE))
cn_head.eval()

# ConvNext Preprocessing (Standard ImageNet)
cn_preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# --- 3. FEATURE EXTRACTION (Heuristics) ---

def extract_simple_features(image_path):
    img = Image.open(image_path).convert('RGB')
    img_array = np.array(img) / 255.0
    
    edges = np.abs(np.diff(np.mean(img_array, axis=2), axis=0)).mean() + \
            np.abs(np.diff(np.mean(img_array, axis=2), axis=1)).mean()
    
    img_smooth = np.array(img.filter(ImageFilter.GaussianBlur(2))) / 255.0
    noise = np.mean((img_array - img_smooth) ** 2) * 1000
    
    return {
        'noise_level': noise,
        'edge_density': edges,
        'is_too_clean': (noise < 0.05 and edges < 0.12) # Adjusted thresholds
    }

# --- 4. THE ENSEMBLE INFERENCE ---

def run_ensemble(image_path):
    img = Image.open(image_path).convert('RGB')
    
    # openclip Score
    img_openclip = openclip_preprocess(img).unsqueeze(0).to(DEVICE)
    with torch.no_grad():
        sig_feat = openclip_model.encode_image(img_openclip)
        sig_feat /= sig_feat.norm(dim=-1, keepdim=True)
        sig_logit = openclip_head(sig_feat)
        prob_openclip = torch.sigmoid(sig_logit).item()
    
    # ConvNeXt Score
    img_cn = cn_preprocess(img).unsqueeze(0).to(DEVICE)
    with torch.no_grad():
        feat = cn_backbone.features(img_cn)
        feat = cn_backbone.avgpool(feat)
        feat = torch.flatten(feat, 1)
        cn_logit = cn_head(feat)
        prob_cn = torch.sigmoid(cn_logit).item()
    
    # Average the two for the "Raw Ensemble Score"
    raw_ensemble_score = (prob_openclip + prob_cn) / 2
    
    # Calibration
    features = extract_simple_features(image_path)
    if features['is_too_clean']:
        calibrated_score = raw_ensemble_score * 0.55 # 45% discount for product shots
        reason = "Clean product-shot detected. Reducing probability."
    else:
        calibrated_score = raw_ensemble_score
        reason = "Standard analysis applied."
    
    return {
        'openclip_score': prob_openclip,
        'convnext_score': prob_cn,
        'raw_ensemble': raw_ensemble_score,
        'calibrated': min(calibrated_score, 0.95),
        'reason': reason,
        'features': features
    }

# --- 5. TEST IT ---
test_image = "/Users/rishitbaitule/Downloads/b.jpg" # Update this path!

if os.path.exists(test_image):
    results = run_ensemble(test_image)
    
    print("-" * 30)
    print(f"Individual openclip: {results['openclip_score']:.2%}")
    print(f"Individual ConvNeXt: {results['convnext_score']:.2%}")
    print("-" * 30)
    print(f"ENSEMBLE RAW SCORE: {results['raw_ensemble']:.2%}")
    print(f"CALIBRATED SCORE:   {results['calibrated']:.2%}")
    print(f"REASON: {results['reason']}")
    print("-" * 30)
else:
    print("Image not found. Please check test_image path.")