# SPDX-FileCopyrightText: 2025 Idiap Research Institute
# SPDX-FileContributor: Anjith George
# SPDX-License-Identifier: BSD-3-Clause

"""EdgeFace demo"""
from __future__ import annotations

from pathlib import Path
import cv2
import gradio as gr
import numpy as np
import torch
import torch.nn.functional as F
from torchvision import transforms

from utils import align_crop

# ───────────────────────────────
# Data & models
# ───────────────────────────────
DATA_DIR = Path("data")
EXTS = (".jpg", ".jpeg", ".png", ".bmp", ".webp")
PRELOADED = sorted(p for p in DATA_DIR.iterdir() if p.suffix.lower() in EXTS)

EDGE_MODELS = [
    "edgeface_base",
    "edgeface_s_gamma_05",
    "edgeface_xs_gamma_06",
    "edgeface_xxs",
]

# ───────────────────────────────
# Styling (orange palette)
# ───────────────────────────────
PRIMARY = "#F97316"
PRIMARY_DARK = "#C2410C"
ACCENT_LIGHT = "#FFEAD2"
BG_LIGHT = "#FFFBF7"
TEXT_DARK = "#0F172A"

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

/* ─── palette ───────────────────────────────────────────── */
body {{
  font-family:'Inter',sans-serif;
  background:{BG_LIGHT};
  color:{TEXT_DARK};
}}

a {{
  color:{PRIMARY};
  text-decoration:none;
  font-weight:600;
}}
a:hover {{color:{PRIMARY_DARK}}}

/* ─── headline ──────────────────────────────────────────── */
#titlebar {{
  text-align:center;
  margin-top:2.4rem;
  margin-bottom:.9rem;
}}
#edgeface-title {{
  font-size:2.6rem;
  font-weight:800;
  margin:0;
  line-height:1.25;
  color: #0F172A;
}}
#edgeface-title .brand {{
  background:linear-gradient(90deg,{PRIMARY} 0%,{PRIMARY_DARK} 90%);
  -webkit-background-clip:text;
  color:transparent;
}}

/* ─── card look ─────────────────────────────────────────── */
.gr-block,
.gr-box,
.gr-row,
#cite-wrapper {{
  border:1px solid #F8C89B;
  border-radius:10px;
  background:#fff;
  box-shadow:0 3px 6px rgba(0,0,0,.05);
}}
.gr-gallery-item {{background:#fff}}

/* ─── controls / inputs ─────────────────────────────────── */
.gr-button-primary,
#copy-btn {{
  background:linear-gradient(90deg,{PRIMARY} 0%,{PRIMARY_DARK} 100%);
  border:none;
  color:#fff;
  border-radius:6px;
  font-weight:600;
  transition:transform .12s ease,box-shadow .12s ease;
}}
.gr-button-primary:hover,
#copy-btn:hover {{
  transform:translateY(-2px);
  box-shadow:0 4px 12px rgba(249,115,22,.35);
}}
.gr-dropdown input {{border:1px solid {PRIMARY}99}}

.preview img,
.preview canvas {{object-fit:contain!important}}

/* ─── hero section ─────────────────────────────────────── */
#hero-wrapper {{text-align:center}}

#hero-badge {{
  display:inline-block;
  padding:.85rem 1.2rem;
  border-radius:8px;
  background:{ACCENT_LIGHT};
  border:1px solid {PRIMARY}55;
  font-size:.95rem;
  font-weight:600;
  margin-bottom:.5rem;
}}

#hero-links {{
  font-size:.95rem;
  font-weight:600;
  margin-bottom:1.6rem;
}}
#hero-links img {{
  height:22px;
  vertical-align:middle;
  margin-left:.55rem;
}}

/* ─── score area ───────────────────────────────────────── */
#score-area {{
  text-align:center;      /* ← centres the badge */
}}

.match-badge {{
  display:inline-block;
  padding:.35rem .9rem;
  border-radius:9999px;
  font-weight:600;
  font-size:1.25rem;      /* ← slightly larger */
}}

/* ─── citation card ────────────────────────────────────── */
#cite-wrapper {{
  position:relative;
  padding:.9rem 1rem;
  margin-top:2rem;
}}
#cite-wrapper code {{
  font-family: SFMono-Regular, Consolas, monospace;
  font-size: .84rem;
  white-space: pre-wrap;
  color: #0F172A; 
}}
#copy-btn {{
  position:absolute;
  top:.55rem;
  right:.6rem;
  padding:.18rem .7rem;
  font-size:.72rem;
  line-height:1;
}}
"""


# ───────────────────────────────
# Torch / transforms
# ───────────────────────────────
_tx = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize([0.5, 0.5, 0.5],[0.5, 0.5, 0.5]),
])
def get_edge_model(name:str)->torch.nn.Module:
    if name not in get_edge_model.cache:
        mdl=torch.hub.load("otroshi/edgeface",name,source="github",pretrained=True).eval()
        mdl.to("cuda" if torch.cuda.is_available() else "cpu")
        get_edge_model.cache[name]=mdl
    return get_edge_model.cache[name]
get_edge_model.cache={}

# ───────────────────────────────
# Helpers
# ───────────────────────────────
def _as_rgb(path:Path)->np.ndarray:
    return cv2.cvtColor(cv2.imread(str(path)),cv2.COLOR_BGR2RGB)

def badge(text:str,colour:str)->str:
    return f'<div class="match-badge" style="background:{colour}22;color:{colour}">{text}</div>'

# ───────────────────────────────
# Face comparison
# ───────────────────────────────
def compare(img_left,img_right,variant):
    crop_a,crop_b=align_crop(img_left),align_crop(img_right)
    if crop_a is None and crop_b is None:
        return None,None,badge("No face detected","#DC2626")
    if crop_a is None:
        return None,None,badge("No face in A","#DC2626")
    if crop_b is None:
        return None,None,badge("No face in B","#DC2626")
    mdl=get_edge_model(variant);dev=next(mdl.parameters()).device
    with torch.no_grad():
        ea=mdl(_tx(cv2.cvtColor(crop_a,cv2.COLOR_RGB2BGR))[None].to(dev))[0]
        eb=mdl(_tx(cv2.cvtColor(crop_b,cv2.COLOR_RGB2BGR))[None].to(dev))[0]
    pct=float(F.cosine_similarity(ea[None],eb[None]).item()*100)
    pct=max(0,min(100,pct))
    colour="#15803D" if pct>=80 else "#CA8A04" if pct>=50 else "#DC2626"
    return crop_a,crop_b,badge(f"{pct:.2f}% match",colour)

# ───────────────────────────────
# Static HTML
# ───────────────────────────────
TITLE_HTML = """
<h1 id='edgeface-title'>
  <span class="brand">EdgeFace:</span> Efficient Face Recognition Model for Edge Devices
</h1>
"""

# <div id="hero-badge">
#   🏆 Winner of IJCB 2023 Efficient Face Recognition Competition
# </div><br/>
  
HERO_HTML = f"""
<div id="hero-wrapper">

  <div id="hero-links">
    <a href="https://www.idiap.ch/paper/edgeface/">Project</a>&nbsp;•&nbsp;
    <a href="https://publications.idiap.ch/attachments/papers/2024/George_IEEETBIOM_2024.pdf">Paper</a>&nbsp;•&nbsp;
    <a href="https://arxiv.org/abs/2307.01838">arXiv</a>&nbsp;•&nbsp;
    <a href="https://gitlab.idiap.ch/bob/bob.paper.tbiom2023_edgeface">Code</a>&nbsp;•&nbsp;
    <img src="https://hitscounter.dev/api/hit?url=https%3A%2F%2Fhuggingface.co%2Fspaces%2idiap%2FEdgeFace&label=Visitors&icon=award-fill&color=%23dc3545" alt="Visitors">
  </div>
</div>
"""

CITATION_HTML = """
<div id="cite-wrapper">
  <button id="copy-btn" onclick="
      navigator.clipboard.writeText(document.getElementById('bibtex').innerText)
        .then(()=>{this.textContent='✔︎';setTimeout(()=>this.textContent='Copy',1500);});
  ">Copy</button>
<code id="bibtex">@article{edgeface,
  title  = {{EdgeFace: Efficient Face Recognition Model for Edge Devices}},
  author = {{George, A. and Ecabert, C. and Otroshi, H. and Kotwal, K. and Marcel, S.}},
  journal= {{IEEE Trans. Biometrics, Behavior, & Identity Science}},
  year   = {{2024}}
}</code>
</div>
"""

# ───────────────────────────────
# Gradio UI
# ───────────────────────────────
with gr.Blocks(css=CSS, title="EdgeFace Demo") as demo:
    gr.HTML(TITLE_HTML, elem_id="titlebar")
    gr.HTML(HERO_HTML)

    with gr.Row():
        gal_a = gr.Gallery(PRELOADED, columns=[5], height=120,
                           label="Image A", object_fit="contain")
        gal_b = gr.Gallery(PRELOADED, columns=[5], height=120,
                           label="Image B", object_fit="contain")

    with gr.Row():
        # img_a = gr.Image(type="numpy", height=300, label="Image A",
        #                  elem_classes="preview")
        # img_b = gr.Image(type="numpy", height=300, label="Image B",
        #                  elem_classes="preview")
        
        img_a = gr.Image(type="numpy", height=300, label="Image A (click or drag-drop)",
                 interactive=True, elem_classes="preview")
        img_b = gr.Image(type="numpy", height=300, label="Image B (click or drag-drop)",
                        interactive=True, elem_classes="preview")

    def _fill(evt: gr.SelectData):
        return _as_rgb(PRELOADED[evt.index]) if evt.index is not None else None
    gal_a.select(_fill, outputs=img_a)
    gal_b.select(_fill, outputs=img_b)

    variant_dd = gr.Dropdown(EDGE_MODELS, value="edgeface_base",
                             label="Model variant")
    btn = gr.Button("Compare", variant="primary")

    with gr.Row():
        out_a = gr.Image(label="Aligned A (112×112)")
        out_b = gr.Image(label="Aligned B (112×112)")
    score_html = gr.HTML(elem_id="score-area")

    btn.click(compare, [img_a, img_b, variant_dd],
              [out_a, out_b, score_html])

    gr.HTML(CITATION_HTML)

# ───────────────────────────────
if __name__ == "__main__":
    demo.launch(share=True)