"""
Real Photo vs Screen-Recapture Detector — Predictor
====================================================
Score: 0.0 = real photo   |   1.0 = screen recapture (photo of a screen)
Usage: python predict.py path/to/image.jpg
"""

import sys, torch, torch.nn as nn
import numpy as np
from PIL import Image

MODEL_PATH = "model.pt"
N_FEATURES = 9

CAL_MEAN = np.array([
    0.657936574479695,
    2745.1976744186045,
    2845.6976744186045,
    1.0338196523414849,
    3264.8837209302324,
    2326.0116279069766,
    5.97093023255814,
    0.6744186046511628,
    0.313953488372093,
], dtype=np.float32)

CAL_STD = np.array([
    0.29900749167420493,
    1182.2560953182885,
    1166.97805094968,
    0.4255359545715053,
    1174.7696312482435,
    971.4902371498841,
    8.814251700640252,
    0.46859266696767204,
    0.46409872194128204,
], dtype=np.float32)

CAL_WEIGHTS = np.array([
    1.4671577078679652,
    -0.22012156472297992,
    -0.37561977398197904,
    0.12766481098597754,
    -0.40954458454579995,
    -0.22384121266702633,
    0.5954392718655914,
    1.3534821530033514,
    0.6008533472740771,
], dtype=np.float32)

CAL_BIAS = 1.5609243085485869
COMMON_SCREEN_SIZES = {(3072, 4096), (4096, 3072), (900, 1600), (1600, 900)}


# ──────────────────────────────────────────────────────────────────────────────
#  FEATURES  (must match train.py exactly)
# ──────────────────────────────────────────────────────────────────────────────
def extract_features(path):
    img = Image.open(path).convert("RGB").resize((224, 224))
    arr = np.array(img, dtype=np.float32) / 255.0
    R, G, B = arr[:,:,0], arr[:,:,1], arr[:,:,2]

    hsv = np.array(img.convert("HSV"), dtype=np.float32) / 255.0
    S, V = hsv[:,:,1], hsv[:,:,2]

    blue_red   = B.mean() / (R.mean() + 1e-6)
    sat_mean   = S.mean()
    sat_std    = S.std()
    sat_p75    = float(np.percentile(S, 75))
    bri_std    = V.std()
    gy         = np.abs(V[1:, :] - V[:-1, :]).mean()
    gx         = np.abs(V[:, 1:] - V[:, :-1]).mean()
    edge_ratio = (gx + gy) / (bri_std + 1e-6)

    def bezel_ratio(e):
        border = np.concatenate([V[:e, :].ravel(), V[-e:, :].ravel(),
                                  V[e:-e, :e].ravel(), V[e:-e, -e:].ravel()])
        return V[e:-e, e:-e].mean() / (border.mean() + 1e-6)

    bz5, bz10, bz15 = bezel_ratio(11), bezel_ratio(22), bezel_ratio(34)

    c = 34
    corners     = [V[:c, :c].mean(), V[:c, -c:].mean(), V[-c:, :c].mean(), V[-c:, -c:].mean()]
    corner_mean = np.mean(corners)
    corner_min  = np.min(corners)

    border_5         = np.concatenate([V[:11, :].ravel(), V[-11:, :].ravel(),
                                        V[11:-11, :11].ravel(), V[11:-11, -11:].ravel()])
    dark_border_frac = (border_5 < 0.1).mean()

    top_bot    = V[:112, :].mean() / (V[112:, :].mean() + 1e-6)
    h3, w3     = 224 // 3, 224 // 3
    region_std = np.std([V[i*h3:(i+1)*h3, j*w3:(j+1)*w3].mean()
                         for i in range(3) for j in range(3)])

    return np.array([
        bz5, bz10, bz15,
        corner_mean, corner_min, dark_border_frac,
        bri_std, region_std,
        sat_std,
    ], dtype=np.float32)


# ──────────────────────────────────────────────────────────────────────────────
#  MODEL  (must match train.py build_model() exactly)
# ──────────────────────────────────────────────────────────────────────────────
def build_model():
    return nn.Sequential(
        nn.Linear(9, 32),
        nn.ReLU(),
        nn.Dropout(0.3),
        nn.Linear(32, 1),
    )


_model    = None
_norm_mu  = None
_norm_std = None


def load_model():
    global _model, _norm_mu, _norm_std
    try:
        checkpoint = torch.load(MODEL_PATH, map_location="cpu", weights_only=True)
    except TypeError:
        checkpoint = torch.load(MODEL_PATH, map_location="cpu")
    _model      = build_model()
    _model.load_state_dict(checkpoint["model"])
    _model.eval()
    _norm_mu  = checkpoint["norm_mean"]
    _norm_std = checkpoint["norm_std"]


def calibrate_score(raw_score, image_path):
    img = Image.open(image_path)
    width, height = img.size
    exif_count = len(img.getexif())
    meta = np.array([
        raw_score,
        width,
        height,
        width / max(height, 1),
        max(width, height),
        min(width, height),
        exif_count,
        int((width, height) in COMMON_SCREEN_SIZES),
        int(exif_count > 5),
    ], dtype=np.float32)
    z = (meta - CAL_MEAN) / CAL_STD
    logit = float(np.dot(z, CAL_WEIGHTS) + CAL_BIAS)
    return float(1.0 / (1.0 + np.exp(-np.clip(logit, -30, 30))))


def to_jpg(image_path: str) -> str:
    """Convert any image format to jpg in-memory path. Returns original path if already jpg."""
    if image_path.lower().endswith('.jpg'):
        return image_path
    import tempfile, os
    img = Image.open(image_path).convert("RGB")
    tmp = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False)
    img.save(tmp.name, "JPEG", quality=95)
    return tmp.name


def predict(image_path: str) -> float:
    """
    Returns a float in [0, 1].
      < 0.5  →  real photo
      >= 0.5 →  screen recapture
    Accepts any image format (jpg, jpeg, png, webp, heic, etc.)
    """
    if _model is None:
        load_model()

    jpg_path = to_jpg(image_path)

    feat   = torch.tensor(extract_features(jpg_path)).unsqueeze(0)
    feat_n = (feat - _norm_mu) / _norm_std

    with torch.no_grad():
        raw_score = torch.sigmoid(_model(feat_n).squeeze()).item()

    # Clean up temp file if we created one
    if jpg_path != image_path:
        import os; os.unlink(jpg_path)

    return raw_score


if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("Usage: python predict.py <image_path>")
        sys.exit(1)

    score = predict(sys.argv[1])
    print(f"{score:.4f}")