import io
import cv2
import numpy as np
import torch
from PIL import Image
import spaces
from transformers import AutoImageProcessor, AutoModelForImageClassification
import gradio as gr

# 1. Inisialisasi Model & Processor
RESNET_MODEL_ID = "prithivMLmods/Deep-Fake-Detector-Model"
print(f"[*] Inisialisasi arsitektur forensik: {RESNET_MODEL_ID}")
processor = AutoImageProcessor.from_pretrained(RESNET_MODEL_ID)
resnet_model = AutoModelForImageClassification.from_pretrained(RESNET_MODEL_ID).eval()

# 2. Analisis Spektral 2D FFT (CPU)
def analyze_phase_spectrum(img_pil: Image.Image) -> float:
    gray_img = img_pil.convert("L")
    w, h = gray_img.size
    crop_size = 256

    if w >= crop_size and h >= crop_size:
        left = (w - crop_size) // 2
        top = (h - crop_size) // 2
        gray_crop = gray_img.crop((left, top, left + crop_size, top + crop_size))
        gray = np.array(gray_crop, dtype=np.float32) / 255.0
    else:
        gray = np.array(gray_img.resize((256, 256)), dtype=np.float32) / 255.0

    f = np.fft.fft2(gray)
    fshift = np.fft.fftshift(f)
    mag = np.abs(fshift)

    cy, cx = 128, 128
    y, x = np.ogrid[:256, :256]
    low_freq_mask = ((x - cx)**2 + (y - cy)**2) <= (32**2)

    low_energy = float(np.mean(mag[low_freq_mask]) + 1e-5)
    high_energy = float(np.mean(mag[~low_freq_mask]) + 1e-5)

    spectral_ratio = high_energy / low_energy
    score = 1.0 / (1.0 + np.exp(-40.0 * (spectral_ratio - 0.045)))
    return float(np.clip(score, 0.05, 0.95))

# 3. Pipeline Forensik Lengkap (Terdaftar Langsung di ZeroGPU)
@spaces.GPU
def predict_forensic(input_image):
    if input_image is None:
        return {
            "verdict": "INCONCLUSIVE",
            "calibrated_confidence": 0.0,
            "reason": "Payload gambar kosong."
        }

    pil_img = Image.fromarray(input_image).convert("RGB")
    
    # Profiler Degradasi
    gray_np = np.array(pil_img.convert("L"))
    laplacian_var = float(cv2.Laplacian(gray_np, cv2.CV_64F).var())
    w, h = pil_img.size
    is_downscaled = bool((w * h) < (720 * 1280))
    quality_score = float(np.clip(laplacian_var / 100.0, 0.05, 1.0))

    if laplacian_var < 5.0 or (laplacian_var < 35.0 and is_downscaled):
        return {
            "verdict": "INCONCLUSIVE",
            "calibrated_confidence": 0.50,
            "reason": "Severe physical degradation outside safe forensic envelope.",
            "metrics": {"sharpness": round(laplacian_var, 2), "dimensions": f"{w}x{h}", "quality": quality_score}
        }

    # Inferensi Akselerasi ZeroGPU (A100)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    resnet_model.to(device)

    inputs = processor(images=pil_img, return_tensors="pt").to(device)
    with torch.inference_mode():
        outputs = resnet_model(**inputs)
        probs = torch.softmax(outputs.logits, dim=-1)[0]

    prob_real = float(probs[0].item())
    score_visual = float(np.clip(1.0 - prob_real, 0.01, 0.99))
    score_phase = analyze_phase_spectrum(pil_img)

    # Kalibrasi Fusi Bukti
    w_visual = max(0.4, 0.8 * quality_score)
    w_phase = 0.2
    combined_score = (score_visual * w_visual + score_phase * w_phase) / (w_visual + w_phase)
    disagreement = abs(score_visual - score_phase)

    if combined_score <= 0.35 and score_visual < 0.15:
        verdict = "LIKELY NATURAL CAPTURE"
        confidence = round(1.0 - combined_score, 3)
        reason = "Optical sensor texture and visual residuals indicate natural capture."
    elif disagreement > 0.55 and quality_score < 0.25 and combined_score > 0.35:
        verdict = "INCONCLUSIVE"
        confidence = 0.50
        reason = "Visual channel and spectral phase exhibit acute divergence under heavy compression."
    elif combined_score >= 0.55:
        verdict = "LIKELY AI-GENERATED"
        confidence = round(min(combined_score, 0.70 + (0.28 * quality_score)), 3)
        reason = "Forensic visual patterns and high-frequency spectral phase reflect synthetic artifacts."
    elif combined_score <= 0.45:
        verdict = "LIKELY NATURAL CAPTURE"
        confidence = round(1.0 - combined_score, 3)
        reason = "Visual features and phase distributions remain consistent with standard optical capture."
    else:
        verdict = "INCONCLUSIVE"
        confidence = 0.50
        reason = "Technical ambiguity margin (0.45 - 0.55)."

    return {
        "verdict": verdict,
        "calibrated_confidence": round(confidence, 3),
        "reason": reason,
        "metrics": {
            "dimensions": f"{w}x{h}",
            "sharpness": round(laplacian_var, 2),
            "quality_index": round(quality_score, 2),
            "score_visual": round(score_visual, 3),
            "score_phase": round(score_phase, 3),
            "combined_score": round(combined_score, 3)
        }
    }

# 4. Antarmuka Gradio Resmi (Memenuhi Spesifikasi ZeroGPU HF)
with gr.Blocks(title="SecurePixel Engine") as demo:
    gr.Markdown("# 🛡️ SecurePixel Forensics Engine (ZeroGPU)")
    
    with gr.Row():
        image_input = gr.Image(type="numpy", label="Input Frame")
        result_output = gr.JSON(label="Forensic Result")
    
    btn_predict = gr.Button("Analyze Frame", variant="primary")
    btn_predict.click(
        fn=predict_forensic,
        inputs=image_input,
        outputs=result_output,
        api_name="predict"
    )

demo.launch()