import gradio as gr
import cv2
import numpy as np
from pathlib import Path
from ultralytics import YOLO

# ── Constants ─────────────────────────────────────────────────────────────────
MODEL_PATH = "yolov8_cattle_keypoints.pt"

KP_NAMES = [
    "left_ear_tip", "right_ear_tip", "left_ear_base", "right_ear_base",
    "left_eye", "right_eye", "nose_left", "nose_right", "nose_tip",
    "mouth_left", "mouth_right", "chin_left", "chin_right",
]

KP_COLORS = [
    (255,  69,   0),
    ( 30, 144, 255),
    (255, 165,   0),
    (  0, 191, 255),
    (154, 205,  50),
    (238, 130, 238),
    (  0, 255, 127),
    (255, 215,   0),
    (255, 255,   0),
    (255,  20, 147),
    (  0, 255, 255),
    (255, 140,   0),
    (147, 112, 219),
]

SKELETON = [
    (0, 2), (1, 3),
    (2, 4), (3, 5),
    (4, 5),
    (6, 8), (7, 8),
    (6, 9), (7, 10),
    (9, 10),
    (9, 11), (10, 12),
    (11, 12),
    (4, 6), (5, 7),
]

# ── Model loading ─────────────────────────────────────────────────────────────
_model = None

def load_model():
    global _model
    if _model is None:
        if not Path(MODEL_PATH).exists():
            raise FileNotFoundError(
                f"Model file '{MODEL_PATH}' not found. "
                "Upload yolov8_cattle_keypoints.pt to the Space root."
            )
        _model = YOLO(MODEL_PATH)
    return _model


# ── Inference ─────────────────────────────────────────────────────────────────
def run_inference(image: np.ndarray, conf_threshold: float, show_labels: bool):
    if image is None:
        return None, "Upload an image first."

    m = load_model()
    results = m.predict(source=image, conf=float(conf_threshold), verbose=False)

    annotated  = image.copy()
    table_rows = []

    for result in results:
        boxes          = result.boxes
        keypoints_data = result.keypoints
        if boxes is None or keypoints_data is None or len(boxes) == 0:
            continue

        for det_idx in range(len(boxes)):
            conf = float(boxes.conf[det_idx])
            kps  = keypoints_data.data[det_idx].cpu().numpy()  # (13, 3)

            # Bounding box
            x1, y1, x2, y2 = boxes.xyxy[det_idx].cpu().numpy().astype(int)
            cv2.rectangle(annotated, (x1, y1), (x2, y2), (255, 255, 255), 2)
            cv2.putText(
                annotated, f"cattle {conf:.2f}",
                (x1, max(y1 - 8, 0)),
                cv2.FONT_HERSHEY_SIMPLEX, 0.55, (255, 255, 255), 1, cv2.LINE_AA,
            )

            # Skeleton
            for (i, j) in SKELETON:
                if i >= len(kps) or j >= len(kps):
                    continue
                xi, yi, vi = kps[i]
                xj, yj, vj = kps[j]
                if vi < 0.5 or vj < 0.5:
                    continue
                cv2.line(
                    annotated,
                    (int(xi), int(yi)), (int(xj), int(yj)),
                    (180, 180, 180), 1, cv2.LINE_AA,
                )

            # Keypoints
            row = {"detection": det_idx + 1, "confidence": f"{conf:.3f}"}
            for kp_idx, (kx, ky, kv) in enumerate(kps):
                name  = KP_NAMES[kp_idx]
                color = KP_COLORS[kp_idx]
                if kv > 0.5:
                    cv2.circle(annotated, (int(kx), int(ky)), 6, color, -1)
                    cv2.circle(annotated, (int(kx), int(ky)), 7, (0, 0, 0), 1)
                    if show_labels:
                        cv2.putText(
                            annotated, name,
                            (int(kx) + 8, int(ky) - 4),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.38,
                            color, 1, cv2.LINE_AA,
                        )
                    row[name] = f"({int(kx)}, {int(ky)})  vis={kv:.2f}"
                else:
                    row[name] = "not visible"

            table_rows.append(row)

    if table_rows:
        lines = [f"### {len(table_rows)} detection(s) found\n"]
        for row in table_rows:
            lines.append(f"**Detection {row['detection']}** — conf {row['confidence']}")
            for name in KP_NAMES:
                lines.append(f"  - `{name}`: {row.get(name, 'n/a')}")
            lines.append("")
        results_md = "\n".join(lines)
    else:
        results_md = "### No cattle detected\nTry lowering the confidence threshold."

    return annotated, results_md


# ── CSS ───────────────────────────────────────────────────────────────────────
CSS = """
@import url('https://fonts.googleapis.com/css2?family=Bebas+Neue&family=DM+Mono:wght@400;500&family=DM+Sans:wght@300;400;500&display=swap');

:root {
    --bg:      #0d0f0e;
    --surface: #161a18;
    --border:  #2a332e;
    --accent:  #4ffe9a;
    --muted:   #7a8c80;
    --text:    #e8ede9;
    --radius:  4px;
}
body, .gradio-container {
    background: var(--bg) !important;
    font-family: 'DM Sans', sans-serif !important;
    color: var(--text) !important;
}
.hdr {
    padding: 2.5rem 0 1.5rem;
    text-align: center;
    border-bottom: 1px solid var(--border);
    margin-bottom: 2rem;
}
.hdr h1 {
    font-family: 'Bebas Neue', sans-serif;
    font-size: clamp(2.8rem, 7vw, 5.5rem);
    letter-spacing: 0.08em;
    color: var(--accent);
    margin: 0;
    line-height: 1;
    text-shadow: 0 0 40px rgba(79,254,154,0.25);
}
.hdr p {
    font-family: 'DM Mono', monospace;
    font-size: 0.78rem;
    color: var(--muted);
    letter-spacing: 0.15em;
    text-transform: uppercase;
    margin: 0.6rem 0 0;
}
.tags {
    display: flex; gap: 0.5rem; flex-wrap: wrap;
    justify-content: center; margin-top: 0.8rem;
}
.tag {
    font-family: 'DM Mono', monospace;
    font-size: 0.68rem; letter-spacing: 0.1em;
    text-transform: uppercase; padding: 0.25rem 0.65rem;
    border: 1px solid var(--border); border-radius: 2px; color: var(--muted);
}
.tag.hot { border-color: var(--accent); color: var(--accent); }
button.primary {
    background: var(--accent) !important;
    color: #0d0f0e !important;
    font-family: 'Bebas Neue', sans-serif !important;
    font-size: 1.1rem !important;
    letter-spacing: 0.12em !important;
    border: none !important;
    border-radius: var(--radius) !important;
    padding: 0.7rem 2rem !important;
    transition: opacity 0.15s, transform 0.1s !important;
}
button.primary:hover { opacity: 0.85 !important; transform: translateY(-1px) !important; }
input[type=range]    { accent-color: var(--accent) !important; }
input[type=checkbox] { accent-color: var(--accent) !important; }
"""

HEADER_HTML = """
<div class="hdr">
  <h1>CattleFace · Pose</h1>
  <p>YOLOv8 · 13-point facial landmark detection for bovines</p>
  <div class="tags">
    <span class="tag hot">13 keypoints</span>
    <span class="tag">ears · eyes · nose · mouth · chin</span>
    <span class="tag hot">real-time inference</span>
    <span class="tag">UARK-AICV benchmark</span>
  </div>
</div>
"""

FOOTER_HTML = """
<div style="text-align:center;padding:1.5rem 0 0.5rem;
            font-family:'DM Mono',monospace;font-size:0.7rem;
            color:#7a8c80;letter-spacing:0.08em;">
  MODEL · YOLOv8s-pose &nbsp;|&nbsp;
  DATASET · UARK-AICV/CattleFace-RGBT-benchmark &nbsp;|&nbsp;
  13 FACIAL LANDMARKS
</div>
"""

# ── Layout ────────────────────────────────────────────────────────────────────
with gr.Blocks(title="CattleFace Pose") as demo:
    gr.HTML(HEADER_HTML)

    with gr.Row():
        with gr.Column(scale=1):
            inp_image = gr.Image(
                label="Input Image",
                type="numpy",
                sources=["upload", "webcam", "clipboard"],
            )
            conf_slider = gr.Slider(
                minimum=0.05, maximum=0.95, value=0.25, step=0.05,
                label="Confidence Threshold",
            )
            show_labels = gr.Checkbox(value=True, label="Show keypoint labels")
            run_btn = gr.Button("Detect Landmarks", variant="primary")

        with gr.Column(scale=1):
            out_image = gr.Image(label="Annotated Output", type="numpy")
            out_text  = gr.Markdown(label="Keypoint Details")

    run_btn.click(
        fn=run_inference,
        inputs=[inp_image, conf_slider, show_labels],
        outputs=[out_image, out_text],
    )

    gr.HTML(FOOTER_HTML)

if __name__ == "__main__":
    demo.launch(css=CSS)