import os
os.environ["CUDA_VISIBLE_DEVICES"] = ""
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"

import gradio as gr
import numpy as np
import cv2
import tensorflow as tf
from PIL import Image
from huggingface_hub import hf_hub_download
from skimage.filters import gabor
from skimage.color import rgb2gray as sk_rgb2gray
from scipy.spatial import cKDTree

# ── Config ───────────────────────────────────────────────────
MODEL_REPO = "Unearthed/artifact-classifier"
MODEL_FILE = "lithic_classifier_hf_fixed.keras"
IMG_SIZE   = (224, 224)

CONF_THRESHOLD        = 0.55
ENTROPY_THRESHOLD     = 1.35
TEXTURE_THRESHOLD     = 15.0
REGULARITY_THRESHOLD  = 0.30
SYMMETRY_THRESHOLD    = 0.38
ARTIFACT_OVERRIDE_THR = 0.45

class_names = ["arrowheads", "axehammerstones", "notartifact", "scrapers"]

DISPLAY_CATEGORY = {
    "arrowheads":      "Projectile Points",
    "axehammerstones": "Other Lithics",
    "scrapers":        "Other Lithics",
    "notartifact":     "Non-Artifacts",
}

def display_category(cls_name):
    for key, label in DISPLAY_CATEGORY.items():
        if key in cls_name.lower():
            return label
    return cls_name

# ── Load model at startup ─────────────────────────────────────
print("Loading model...")
model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)
model = tf.keras.models.load_model(model_path, compile=False)
print("Model ready!")

# ── Feature extraction ───────────────────────────────────────
def build_5channel_input(img_np):
    rgb  = img_np.astype(np.float32) / 255.0
    gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY).astype(np.float32)
    blur = cv2.GaussianBlur(gray, (5, 5), 0)
    lap  = np.abs(cv2.Laplacian(blur, cv2.CV_32F))
    lap  = (lap - lap.min()) / (lap.max() - lap.min() + 1e-8)
    gray_sk = sk_rgb2gray(img_np).astype(np.float64)
    resps = [np.abs(gabor(gray_sk, frequency=0.15, theta=t)[0])
             for t in [0, np.pi/4, np.pi/2]]
    gab = np.mean(resps, axis=0).astype(np.float32)
    gab = (gab - gab.min()) / (gab.max() - gab.min() + 1e-8)
    return np.concatenate([rgb, lap[:,:,np.newaxis], gab[:,:,np.newaxis]], axis=-1)

def texture_complexity_score(img_np):
    gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)
    return float(cv2.Laplacian(gray, cv2.CV_64F).var())

def edge_regularity_score(img_np):
    gray  = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)
    edges = cv2.Canny(gray, 50, 150)
    ys, xs = np.where(edges > 0)
    if len(xs) < 20:
        return 0.0
    pts  = np.stack([xs, ys], axis=1).astype(np.float32)
    tree = cKDTree(pts)
    dists, _ = tree.query(pts, k=2)
    return float(np.clip(1.0 / (1.0 + np.std(dists[:, 1])), 0, 1))

def symmetry_score(img_np):
    gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY).astype(np.float32)
    h, w = gray.shape
    left  = gray[:, :w//2]
    right = gray[:, w//2:][:, ::-1]
    min_w = min(left.shape[1], right.shape[1])
    left, right = left[:, :min_w], right[:, :min_w]
    corr = np.corrcoef(left.flatten(), right.flatten())[0, 1]
    return float(np.clip((corr + 1) / 2, 0, 1))

def stem_notch_score(img_np):
    gray   = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)
    _, bw  = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    h, w   = bw.shape
    def row_width(region):
        widths = []
        for row in region:
            cols = np.where(row > 0)[0]
            if len(cols) > 2:
                widths.append(cols[-1] - cols[0])
        return np.mean(widths) if widths else 0
    full_width = row_width(bw)
    base_width = row_width(bw[int(h * 0.75):, :])
    if full_width < 1:
        return 0.0
    return float(np.clip(1.0 - (base_width / full_width), 0, 1))

def is_likely_broken_artifact(tex, reg, sym):
    return tex > 80 and reg > 0.55 and sym < 0.70

# ── Main prediction ──────────────────────────────────────────

def predict(image):
    print("=== PREDICT CALLED ===")
    if image is None:
        print("=== IMAGE IS NONE ===")
        return "Please upload an image."
    
    print("=== CONVERTING IMAGE ===")
    img_np = np.array(image.convert("RGB").resize(IMG_SIZE))
    print(f"=== IMAGE SHAPE: {img_np.shape} ===")
    
    tex_score = texture_complexity_score(img_np)
    print(f"=== TEX SCORE: {tex_score} ===")
    reg_score  = edge_regularity_score(img_np)
    sym_score  = symmetry_score(img_np)
    stem_score = stem_notch_score(img_np)

    lines = []
    lines.append("=" * 55)
    lines.append("  FEATURE ANALYSIS (archaeological gates)")
    lines.append("=" * 55)
    lines.append(f"  Texture complexity  : {tex_score:.1f}  (>15 = real object, >80 = worked stone)")
    lines.append(f"  Edge regularity     : {reg_score:.3f} (>0.30 = systematic flaking)")
    lines.append(f"  Bilateral symmetry  : {sym_score:.3f} (>0.38 = hafted point geometry)")
    lines.append(f"  Stem/notch evidence : {stem_score:.3f} (>0.25 = base constriction for hafting)")
    lines.append("")

    lines.append(f"[Gate 1] Texture complexity: {tex_score:.1f}")
    if tex_score < TEXTURE_THRESHOLD:
        lines.append("  ❌ REJECTED — surface too uniform (not a real stone photo)")
        return "\n".join(lines)

    lines.append(f"[Gate 2] Edge regularity: {reg_score:.3f}")

    x5    = build_5channel_input(img_np)
    preds = model.predict(np.expand_dims(x5, 0), verbose=0)[0]

    top_idx  = int(np.argmax(preds))
    top_cls  = class_names[top_idx]
    top_conf = float(preds[top_idx])
    ent      = float(-np.sum(preds * np.log(preds + 1e-8)))

    notartifact_idx    = next((i for i, c in enumerate(class_names) if "notartifact" in c.lower()), None)
    artifact_score_sum = sum(float(preds[i]) for i, c in enumerate(class_names) if "notartifact" not in c.lower())
    broken = is_likely_broken_artifact(tex_score, reg_score, sym_score)

    lines.append("")
    lines.append("  Model raw scores:")
    merged_scores = {}
    for i, cls in enumerate(class_names):
        lbl = display_category(cls)
        merged_scores[lbl] = merged_scores.get(lbl, 0.0) + float(preds[i])
    top_display = display_category(top_cls)
    for lbl, score in merged_scores.items():
        bar = "█" * int(score * 30)
        tag = " ← predicted" if lbl == top_display else ""
        lines.append(f"    {lbl:22s} {score:.3f}  {bar}{tag}")

    lines.append("")
    lines.append(f"  Top class   : {display_category(top_cls)}")
    lines.append(f"  Confidence  : {top_conf:.3f}")
    lines.append(f"  Entropy     : {ent:.3f}")
    lines.append(f"  Artifact sum: {artifact_score_sum:.3f} (combined non-notartifact score)")

    if broken:
        lines.append("")
        lines.append(f"  🔍 FRAGMENT INDICATOR: High texture ({tex_score:.0f}) + "
                     f"High regularity ({reg_score:.2f}) + Reduced symmetry ({sym_score:.2f})")
        lines.append("     → Characteristic signature of a BROKEN/FRAGMENTARY worked stone tool")

    override_triggered = False
    if top_cls == "notartifact":
        if tex_score > 80 and reg_score > 0.55 and artifact_score_sum > ARTIFACT_OVERRIDE_THR:
            artifact_scores = {c: float(preds[i]) for i, c in enumerate(class_names) if i != notartifact_idx}
            top_cls  = max(artifact_scores, key=artifact_scores.get)
            top_conf = artifact_scores[top_cls]
            override_triggered = True
            lines.append("")
            lines.append("  ⚡ ARCHAEOLOGICAL OVERRIDE TRIGGERED")
            lines.append(f"     Overriding notartifact → {display_category(top_cls)}")
        else:
            lines.append("")
            lines.append("  ✅ PREDICTION: NON-ARTIFACTS")
            lines.append(f"     Feature gates did not trigger override:")
            lines.append(f"       Texture {tex_score:.0f}, Regularity {reg_score:.2f}, "
                         f"Artifact sum {artifact_score_sum:.2f}")
            return "\n".join(lines)

    lines.append(f"[Gate 3] Bilateral symmetry: {sym_score:.3f}")
    if "arrowhead" in top_cls.lower():
        if sym_score < SYMMETRY_THRESHOLD:
            if broken or override_triggered:
                lines.append(f"  ⚠️  Low symmetry ({sym_score:.2f}) but broken fragment detected — gate bypassed")
            else:
                lines.append(f"  ⚠️  Low symmetry for arrowhead prediction — may be debitage")

    lines.append(f"[Gate 4] Confidence: {top_conf:.3f} | Entropy: {ent:.3f}")
    if not override_triggered:
        if top_conf < CONF_THRESHOLD or ent > ENTROPY_THRESHOLD:
            lines.append("")
            lines.append("  ⚠️  NOT CONFIDENT — ambiguous result")
            lines.append("     (Broken or unusual specimen — consider manual review)")
            return "\n".join(lines)

    lines.append("")
    disp = display_category(top_cls)
    if override_triggered:
        lines.append(f"  ✅ FINAL PREDICTION: {disp.upper()} (via archaeological override)")
        lines.append(f"     (Internal model class: {top_cls})")
    else:
        lines.append(f"  ✅ PREDICTION: {disp.upper()}")
        lines.append(f"     (Internal model class: {top_cls})")

    if broken:
        lines.append("  📌 NOTE: Fragmentary/broken specimen. Morphology may be incomplete.")
    if stem_score > 0.25 and ("arrowhead" in top_cls.lower() or override_triggered):
        lines.append(f"  📌 Stem/notch evidence detected (score {stem_score:.2f}) — consistent with hafted projectile point")

    interp = {
        "arrowhead": (
            "Projectile point. Key features to verify:\n"
            "     • Bifacial flaking on both faces\n"
            "     • Thinned cross-section (blade thinner than base)\n"
            "     • Stem or notching for hafting\n"
            "     • Regular, intentional edge serrations\n"
            "     • Conchoidal fracture surface texture (not granular)"
        ),
        "scraper": (
            "Scraper. Key features to verify:\n"
            "     • Unifacial working (one face only)\n"
            "     • Steep retouch angle on working edge (~70°)\n"
            "     • Cortex usually present on dorsal face\n"
            "     • No stem or hafting preparation"
        ),
        "axehammerstone": (
            "Axe/Hammerstone. Key features to verify:\n"
            "     • Pecking marks from repeated impact use\n"
            "     • Ground or polished facets\n"
            "     • Heavier, more rounded form\n"
            "     • May show battered/crushed areas"
        ),
    }
    for key, msg in interp.items():
        if key in top_cls.lower():
            lines.append("")
            lines.append("  📋 Archaeological checklist:")
            lines.append(f"     {msg}")
            break

    lines.append("=" * 55)
    return "\n".join(lines)

# ── Gradio UI ────────────────────────────────────────────────
with gr.Blocks(title="Unearthed — Artifact Classifier") as demo:
    gr.Markdown("""
    # 🪨 Unearthed — Artifact Classifier
    Upload a photo of a stone object to classify it.
    **Categories: Projectile Points · Other Lithics · Non-Artifacts**
    """)
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil", label="Upload image")
            submit_btn  = gr.Button("Classify", variant="primary")
        with gr.Column():
            output_text = gr.Textbox(
                label="Analysis Result",
                lines=30,
                max_lines=40,
            )

    submit_btn.click(fn=predict, inputs=image_input, outputs=output_text)

demo.launch()