# SPDX-FileCopyrightText: Copyright © 2025 Idiap Research Institute <contact@idiap.ch>
# SPDX-FileContributor: Francois Poh <francois.poh22@imperial.ac.uk>
# SPDX-FileContributor: Anjith George <anjith.george@idiap.ch>
# SPDX-FileContributor: Samuel Michel <samuel.michel@idiap.ch>
#
# SPDX-License-Identifier: GPL-3.0-or-later

from __future__ import annotations

import os
import time
from functools import lru_cache
from collections.abc import Callable

import numpy as np
import pandas as pd
import torch
import gradio as gr
import matplotlib.pyplot as plt

import onnxruntime as ort

from lib.models import get_model
from lib.align import get_preprocessor

from calibrate_score import (
    fit_calibrator_from_csv,
    fit_calibrator_from_scores,
    apply_calibrator,
)

ort.set_default_logger_severity(3)

# =====================================================
# Configuration
# =====================================================

MODEL_VARIANTS = ["clip", "lora", "ires100", "ires100-tune"]
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

PROJECT_URL = "https://www.idiap.ch/paper/artface/"
ARXIV_URL = "https://arxiv.org/abs/2508.20626"

DATASET_DIRS = {
    "Historical Faces": "historicalface",
    "Museum": "museum",
}

def get_calibration_files(folder):
    return {
        model: os.path.join(folder, f"{model}.csv")
        for model in MODEL_VARIANTS
    }

# =====================================================
# Original Palette & Friendly Professional Styling
# =====================================================

PRIMARY = "#F97316"  # Original Orange
BG = "#FBFBFC"
CARD = "#FFFFFF"
BORDER = "#E7E7EA"
TEXT = "#0B1220"
MUTED = "#556070"

CSS = f"""
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;600;700;800;900&display=swap');

body, .gradio-container {{
  background: {BG} !important;
  font-family: 'Inter', sans-serif !important;
  color: {TEXT};
}}

#wrap {{
  max-width: 940px;
  margin: 0 auto;
}}

.gr-block, .gr-box, .gr-panel {{
  background: {CARD} !important;
  border: 1px solid {BORDER} !important;
  border-radius: 16px !important;
  box-shadow: none !important;
}}

.gr-button-primary {{
  background: {TEXT} !important;
  border: none !important;
  color: white !important;
  font-weight: 800 !important;
  border-radius: 12px !important;
}}

.kicker {{
  display:flex;
  align-items:center;
  justify-content:space-between;
  gap: 1rem;
  padding: 1.5rem 0 .5rem 0;
}}

.title {{
  font-size: 1.8rem;
  font-weight: 900;
  letter-spacing: -0.02em;
  color: {TEXT};
}}

.subtitle {{
  color: {MUTED};
  font-weight: 600;
  font-size: .95rem;
}}

.chip {{
  border: 1px solid {BORDER};
  border-radius: 999px;
  padding: .25rem .7rem;
  font-weight: 800;
  font-size: .75rem;
  color: {MUTED};
  background: #fff;
}}

.section-h {{
  font-size: 1rem;
  font-weight: 900;
  margin-bottom: .2rem;
}}

.hint {{
  color: {MUTED};
  font-weight: 600;
  font-size: .9rem;
  margin-bottom: 1rem;
}}

.fused-box {{
  display:flex;
  align-items:center;
  justify-content:space-between;
  padding: 1.25rem;
  border-radius: 16px;
  border: 1px solid {BORDER};
  background: #fff;
  margin: 1rem 0;
}}

.fused-title {{
  font-weight: 900;
  font-size: 1.1rem;
  margin-bottom: .2rem;
}}

.fused-meta {{
  color: {MUTED};
  font-weight: 700;
  font-size: .85rem;
  line-height: 1.4;
}}

.llr-val {{
  font-size: 2.2rem;
  font-weight: 950;
  color: {PRIMARY};
  letter-spacing: -0.03em;
}}

.pill {{
  display: inline-block;
  padding: .2rem .6rem;
  border-radius: 6px;
  font-weight: 800;
  font-size: .75rem;
  margin-top: .5rem;
}}

.footer {{
  text-align:center;
  color: {MUTED};
  font-weight: 600;
  padding: 2rem 0;
  font-size: .85rem;
}}
"""

TITLE_HTML = f"""
<div class="kicker">
  <div>
    <div class="title"><span style="color:{PRIMARY};">ArtFace</span> Demo</div>
    <div class="subtitle">Forensic face identification for historical portraits</div>
  </div>
  <div style="display:flex; gap:0.5rem;">
    <div class="chip">GPU: {"ON" if torch.cuda.is_available() else "OFF"}</div>
    <div class="chip">Calibration: ENFSI</div>
  </div>
</div>
<div style="display:flex; gap:1.2rem; margin-bottom: 1.5rem;">
  <a href="{PROJECT_URL}" target="_blank" style="color:{TEXT}; font-weight:800; text-decoration:none; font-size:.9rem;">Project Page</a>
  <a href="{ARXIV_URL}" target="_blank" style="color:{TEXT}; font-weight:800; text-decoration:none; font-size:.9rem;">arXiv Paper</a>
</div>
"""

# =====================================================
# Backend
# =====================================================

aligner = get_preprocessor("align")
MODELS = {}
for name in MODEL_VARIANTS:
    model, prep = get_model(name).torch()
    model.eval().to(DEVICE)
    MODELS[name] = (model, prep)


@lru_cache(maxsize=None)
def get_cached_dynamic_calibrator(selected_models_key, fuse_method, cal_folder):
    selected_models = list(selected_models_key)
    calibration_files = get_calibration_files(cal_folder)
    key_cols = ["probe_subject_id", "bio_ref_subject_id"]
    merged = None
    for name in selected_models:
        df = pd.read_csv(calibration_files[name])[key_cols + ["score"]].rename(columns={"score": f"score_{name}"})
        merged = df if merged is None else merged.merge(df, on=key_cols, how="inner")

    score_cols = [f"score_{name}" for name in selected_models]
    labels = (merged["probe_subject_id"] == merged["bio_ref_subject_id"]).astype(int).values

    scores_mat = merged[score_cols].values
    if fuse_method == "median":
        fused = np.median(scores_mat, axis=1)
    elif fuse_method == "max":
        fused = np.max(scores_mat, axis=1)
    else:
        fused = np.mean(scores_mat, axis=1)

    calibrator = fit_calibrator_from_scores(fused, labels, P_tar=0.5)
    calibrator.update({"cohort_scores": fused, "cohort_labels": labels})
    return calibrator


def make_plot(cal, target_score):
    w, b = float(cal["w"]), float(cal["b"])
    llrs = ((w * cal["cohort_scores"]) + b) / np.log(10)
    target_llr = ((w * target_score) + b) / np.log(10)

    fig, ax = plt.subplots(figsize=(9, 5), facecolor=BG)
    ax.hist(llrs[cal["cohort_labels"] == 1], bins=40, alpha=0.6, label="Genuines", color="tab:blue", orientation='horizontal', density=True)
    ax.hist(llrs[cal["cohort_labels"] == 0], bins=40, alpha=0.6, label="Impostors", color="tab:orange", orientation='horizontal', density=True)

    ax.axhline(target_llr, color=TEXT, lw=3, ls='--', label=f"Result: {target_llr:.2f}")

    yticks = [-7, -5, -3, -1, 0, 1, 3, 5, 7]
    ylabs = ["Extreme $H_I$", "V.Strong $H_I$", "Strong $H_I$", "Weak $H_I$", "Neutral", "Weak $H_G$", "Strong $H_G$", "V.Strong $H_G$", "Extreme $H_G$"]
    ax.set_yticks(yticks)
    ax.set_yticklabels(ylabs, fontsize=9)
    ax.set_ylim([-8, 8])
    ax.set_ylabel("ENFSI Verbal Scale")
    ax.set_xlabel("Frequency (Counts)")
    ax.set_title("Calibrated Likelihood Distribution (Cohort Counts)", fontweight="bold")
    ax.grid(axis='y', alpha=0.2)
    ax.legend(frameon=False, loc="upper right")
    plt.tight_layout()
    return fig


def process(img1, img2, models, method, cal_dataset):
    if not img1 or not img2:
        return [None] * 4 + [pd.DataFrame()]
    a1, a2 = aligner(img1), aligner(img2)
    if not a1 or not a2:
        return [None] * 2 + ["No face detected", None, pd.DataFrame()]

    cal_folder = DATASET_DIRS[cal_dataset]

    start = time.time()
    scores = {}
    for n in models:
        m, prep = MODELS[n]
        x1 = prep(a1).unsqueeze(0).to(DEVICE)
        x2 = prep(a2).unsqueeze(0).to(DEVICE)
        with torch.no_grad():
            e1, e2 = m(x1)[0].cpu().numpy(), m(x2)[0].cpu().numpy()
        scores[n] = float(np.dot(e1, e2) / (np.linalg.norm(e1) * np.linalg.norm(e2) + 1e-12))

    if method == "mean":
        f_score = np.mean(list(scores.values()))
    elif method == "median":
        f_score = np.median(list(scores.values()))
    else:
        f_score = np.max(list(scores.values()))

    dur = time.time() - start

    try:
        cal = get_cached_dynamic_calibrator(tuple(sorted(models)), method, cal_folder)
        res = apply_calibrator(f_score, cal)
        llr_val, interp = res['llr_10'], res['interpretation']
        plot = make_plot(cal, f_score)
    except Exception as e:
        llr_val, interp, plot = 0.0, f"Error: {str(e)}", None

    pill_bg = "#DCFCE7" if llr_val > 1 else ("#FEF2F2" if llr_val < -1 else "#F1F5F9")
    pill_tx = "#166534" if llr_val > 1 else ("#991B1B" if llr_val < -1 else "#475569")

    html = f"""
    <div class="fused-box">
      <div>
        <div class="fused-title">Likelihood Ratio (Log₁₀)</div>
        <div class="fused-meta">
          Method: <b>{method}</b> · Models: {len(models)} · Dataset: <b>{cal_dataset}</b> · ⏱ {dur:.2f}s<br>
          <span class="pill" style="background:{pill_bg}; color:{pill_tx};">Verdict: {interp}</span>
        </div>
      </div>
      <div style="text-align:right;">
        <div class="llr-val">{llr_val:.3f}</div>
        <div class="fused-meta">LLR Score</div>
      </div>
    </div>
    """

    rows = [{"Model": n, "Similarity": round(s, 4)} for n, s in scores.items()]
    return a1, a2, html, plot, pd.DataFrame(rows)


# =====================================================
# UI
# =====================================================

with gr.Blocks(title="ArtFace") as demo:
    with gr.Column(elem_id="wrap"):
        gr.HTML(TITLE_HTML)

        with gr.Group():
            gr.HTML('<div class="section-h">Inputs</div><div class="hint">Upload Reference and Probe images for alignment and identification.</div>')
            with gr.Row():
                i1 = gr.Image(label="Image A", type="pil", height=300)
                i2 = gr.Image(label="Image B", type="pil", height=300)

            with gr.Row():
                run = gr.Button("Compare Identities", variant="primary")
                clear = gr.ClearButton([i1, i2], value="Clear")

        with gr.Accordion("Analysis Settings", open=False):
            sel = gr.CheckboxGroup(MODEL_VARIANTS, value=["lora", "ires100-tune"], label="Active Models")
            met = gr.Radio(["mean", "median", "max"], value="mean", label="Fusion Method")
            cal_dir = gr.Dropdown(
                choices=list(DATASET_DIRS.keys()),
                value="Historical Faces",
                label="Calibration Dataset",
            )

        gr.HTML('<div style="height:1px; background:#E7E7EA; margin: 1.5rem 0;"></div>')

        gr.HTML('<div class="section-h">Results</div><div class="hint">Calibrated LLR based on ENFSI standards.</div>')

        with gr.Row():
            o1 = gr.Image(label="Aligned A", height=150, interactive=False)
            o2 = gr.Image(label="Aligned B", height=150, interactive=False)

        res_html = gr.HTML('<div style="padding:1rem; text-align:center; color:#556070;">Execute analysis to see likelihood score.</div>')

        with gr.Row():
            res_table = gr.Dataframe(label="Individual Scores", interactive=False)
            res_plot = gr.Plot(label="Likelihood Distribution")

        gr.HTML('<div class="footer">Research Demo · Idiap Research Institute · 2026</div>')

    run.click(process, [i1, i2, sel, met, cal_dir], [o1, o2, res_html, res_plot, res_table])

if __name__ == "__main__":
    demo.launch(
        theme=gr.themes.Soft(primary_hue="orange", neutral_hue="slate"),
        css=CSS, share=True,
    )