from flask import Flask, request, jsonify
from flask_cors import CORS
import torch
import torch.nn as nn
import open_clip
from PIL import Image, ImageFilter
from torchvision import transforms, models
from io import BytesIO
import requests
import base64
import numpy as np

app = Flask(__name__)
CORS(app)

DEVICE = torch.device("mps") if torch.backends.mps.is_available() else torch.device("cpu")


class ForensicHead(nn.Module):
    def __init__(self, input_dim=768):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, 512),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(512, 1),
            nn.Sigmoid()
        )

    def forward(self, x):
        return self.net(x)


print("Loading Models...")

model, _, preprocess = open_clip.create_model_and_transforms(
    "ViT-L-14",
    pretrained="datacomp_xl_s13b_b90k"
)
model = model.to(DEVICE)
model.eval()

tokenizer = open_clip.get_tokenizer("ViT-L-14")

AI_FLAWS = [
    "plastic skin, overly smooth textures, and lack of realistic pores",
    "distorted anatomical shapes like strange hands, limbs, or face",
    "inconsistent lighting, impossible shadows, or unnatural highlights",
    "garbled, blurred, or nonsensical background details and text",
    "asymmetrical facial features or floating elements",
    "blending errors where subjects melt unnaturally into the background"
]

REAL_TRAITS = [
    "natural texture with visible realistic imperfections",
    "physically consistent lighting, shadows, and reflections",
    "natural anatomical proportions and distinct physical boundaries",
    "sharp, coherent background elements and depth of field",
    "authentic noise and realistic color balance"
]

print("Encoding explainability vectors...")
ai_tokens = tokenizer(AI_FLAWS).to(DEVICE)
real_tokens = tokenizer(REAL_TRAITS).to(DEVICE)

with torch.no_grad():
    ai_text_features = model.encode_text(ai_tokens)
    ai_text_features /= ai_text_features.norm(dim=-1, keepdim=True)
    
    real_text_features = model.encode_text(real_tokens)
    real_text_features /= real_text_features.norm(dim=-1, keepdim=True)

head = ForensicHead(input_dim=768)
head.load_state_dict(torch.load("models/openclip_forensic_head.pth", map_location=DEVICE))
head = head.to(DEVICE)
head.eval()

cn_backbone = models.convnext_base(weights=None)
cn_backbone.to(DEVICE)
cn_backbone.eval()

class ConvNextHead(nn.Module):
    def __init__(self, input_dim=1024):
        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()

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])
])

print("Models loaded")


def load_image_from_url(url: str) -> Image.Image:
    headers = {
        "User-Agent": "Mozilla/5.0"
    }
    response = requests.get(url, headers=headers, timeout=8)
    response.raise_for_status()
    return Image.open(BytesIO(response.content)).convert("RGB")


def load_image_from_data_url(data_url: str) -> Image.Image:
    if "," not in data_url:
        raise ValueError("Invalid data URL")

    _, encoded = data_url.split(",", 1)
    raw = base64.b64decode(encoded)
    return Image.open(BytesIO(raw)).convert("RGB")


def load_any_image(payload: dict) -> Image.Image:
    if "image_url" in payload and payload["image_url"]:
        src = payload["image_url"]
        if src.startswith("data:"):
            return load_image_from_data_url(src)
        return load_image_from_url(src)

    if "image" in payload and payload["image"]:
        return load_image_from_data_url(payload["image"])

    raise ValueError("No image data provided")


def get_explanation(label: str, img_feat_tensor: torch.Tensor, heuristics: dict, confidence: float) -> str:
    """Combines Zero-Shot CLIP semantic extraction with raw OpenCV Image Processing."""
    noise = heuristics.get('noise_level', 0)
    edges = heuristics.get('edge_density', 0)
    
    if label == "AI":
        # find the closest semantic flaw using dot product similarity
        similarity = (100.0 * img_feat_tensor @ ai_text_features.T).softmax(dim=-1)
        top_idx = similarity.argmax().item()
        semantic_reason = AI_FLAWS[top_idx]
        
        technical_reason = []
        if noise < 0.025:
            technical_reason.append(f"an unnatural lack of sensor noise ({noise:.3f})")
        if edges < 0.1:
            technical_reason.append("abnormally soft structural contours")
            
        tech_str = (" coupled directly with " + " and ".join(technical_reason)) if technical_reason else ""
        return f"The model detected {semantic_reason}{tech_str}.<br/><br/><strong>Assessed as Synthetic ({confidence*100:.1f}% confidence)</strong>"
    else:
        # Feature Extraction: Find closest authentic trait
        similarity = (100.0 * img_feat_tensor @ real_text_features.T).softmax(dim=-1)
        top_idx = similarity.argmax().item()
        semantic_reason = REAL_TRAITS[top_idx]
        
        technical_reason = []
        if noise > 0.04:
            technical_reason.append(f"expected natural grain matrix ({noise:.3f})")
        if edges >= 0.1:
            technical_reason.append("well-defined structural boundaries")
            
        tech_str = (" supported by " + " and ".join(technical_reason)) if technical_reason else ""
        return f"The model identified {semantic_reason}{tech_str}.<br/><br/><strong>Assessed as Authentic</strong>"


def extract_simple_features(img: Image.Image):
    """Extract image characteristics WITHOUT ML"""
    img_rgb = img.convert('RGB')
    img_array = np.array(img_rgb) / 255.0

    # Edge density (real photos have more natural edges, less uniform)
    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()

    # Noise level (real photos have noise, AI images are smoother)
    img_smooth = np.array(img_rgb.filter(ImageFilter.GaussianBlur(2))) / 255.0
    noise = np.mean((img_array - img_smooth) ** 2) * 1000  # scale for visibility

    # Color balance (product photos often have strong color gradients)
    hsv = img.convert('HSV')
    hsv_array = np.array(hsv) / 255.0
    color_variance = np.var(hsv_array[:, :, 0])  # hue variance

    return {
        'edge_density': edges,
        'noise_level': noise,
        'color_variance': color_variance,
        'is_too_clean': (noise < 0.02 and edges < 0.1),  # product photo signature
    }


def calibrate_output(img: Image.Image, raw_confidence: float):
    """Adjust model confidence based on image characteristics"""
    features = extract_simple_features(img)

    if features['is_too_clean']:
        adjusted_confidence = raw_confidence * 0.67
    else:
        adjusted_confidence = raw_confidence
    adjusted_confidence = min(adjusted_confidence, 0.90)

    return adjusted_confidence, features


@app.route("/predict", methods=["POST"])
def predict():
    try:
        data = request.get_json(force=True, silent=False)
        image = load_any_image(data)

        openclip_weight = 0.95
        convnext_weight = 0.05

        # openclip inference
        img_openclip = preprocess(image).unsqueeze(0).to(DEVICE)
        with torch.no_grad():
            features = model.encode_image(img_openclip)
            features = features / features.norm(dim=-1, keepdim=True)
            prob_openclip = float(head(features).item())
            
            global_image_features = features.squeeze(0).clone()

        # convnext inference
        img_cn = cn_preprocess(image).unsqueeze(0).to(DEVICE)
        with torch.no_grad():
            cn_feat = cn_backbone.features(img_cn)
            cn_feat = cn_backbone.avgpool(cn_feat)
            cn_feat = torch.flatten(cn_feat, 1)
            cn_logit = cn_head(cn_feat)
            prob_cn = torch.sigmoid(cn_logit).item()
            
        total_ml_weight = openclip_weight + convnext_weight
        if total_ml_weight > 0:
            raw_ensemble_score = (prob_openclip * openclip_weight + prob_cn * convnext_weight) / total_ml_weight
        else:
            raw_ensemble_score = (prob_openclip + prob_cn) / 2.0

        prob, img_features = calibrate_output(image, raw_ensemble_score)

        label = "AI" if prob >= 0.75 else "Real"
        confidence = prob if label == "AI" else "✓"

        return jsonify({
            "label": label,
            "confidence": confidence,
            "explanation": get_explanation(label, global_image_features, img_features, confidence),
            "scores": {
                "openclip": prob_openclip,
                "convnext": prob_cn,
                "ensemble_raw": raw_ensemble_score,
                "calibrated_final": prob
            }
        })

    except Exception as e:
        return jsonify({
            "error": str(e)
        }), 500


if __name__ == "__main__":
    app.run(debug=True)