"""
Wedding Burst Detection — Gradio Web App
Run: python burst_detection_app.py
Requires: pip install gradio timm faiss-cpu imagehash umap-learn hdbscan tqdm pillow exifread pandas scikit-learn torch torchvision
"""

import os, re, json, shutil, logging, zipfile, tempfile
from pathlib import Path
from datetime import datetime
from typing import List, Optional
from math import ceil

import numpy as np
import pandas as pd
import torch
import torch.nn.functional as F
import timm
from torchvision import transforms
from torch.utils.data import Dataset, DataLoader
from PIL import Image, UnidentifiedImageError, ImageDraw, ImageFont
import exifread
import imagehash
import faiss
from sklearn.mixture import GaussianMixture
from sklearn.cluster import DBSCAN
from tqdm import tqdm
import gradio as gr
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches

# ─── Logging ───────────────────────────────────────────────────────────────────
logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
logger = logging.getLogger("BurstApp")

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
SUPPORTED_EXTS = {".jpg", ".jpeg", ".png", ".tif", ".tiff", ".webp", ".heic"}

# ─── Global model cache ────────────────────────────────────────────────────────
_model = None
_transform = None


def get_model():
    global _model, _transform
    if _model is None:
        logger.info("Loading DINOv2...")
        _model = timm.create_model(
            "vit_base_patch14_dinov2.lvd142m", pretrained=True, num_classes=0
        )
        _model.eval().to(DEVICE)
        _transform = transforms.Compose([
            transforms.Resize(518, interpolation=transforms.InterpolationMode.BICUBIC, antialias=True),
            transforms.CenterCrop(518),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        ])
        logger.info("DINOv2 ready.")
    return _model, _transform


# ─── Image Discovery ──────────────────────────────────────────────────────────

def discover_images(root: str) -> List[Path]:
    paths = sorted(
        p for p in Path(root).rglob("*") if p.suffix.lower() in SUPPORTED_EXTS
    )
    return paths


def _parse_exif_dt(tags):
    for field in ("EXIF DateTimeOriginal", "EXIF DateTimeDigitized", "Image DateTime"):
        if field in tags:
            try:
                return datetime.strptime(str(tags[field]), "%Y:%m:%d %H:%M:%S")
            except ValueError:
                pass
    return None


def build_metadata(paths: List[Path]) -> pd.DataFrame:
    records = []
    for idx, p in enumerate(paths):
        dt = None
        try:
            with open(p, "rb") as fh:
                tags = exifread.process_file(fh, stop_tag="DateTimeOriginal", details=False)
            dt = _parse_exif_dt(tags)
        except Exception:
            pass
        nums = re.findall(r"\d+", p.stem)
        fn_order = int(nums[-1]) if nums else idx
        records.append({"file_index": idx, "path": str(p), "filename": p.name,
                        "folder": str(p.parent), "exif_datetime": dt, "fn_order": fn_order})

    df = pd.DataFrame(records)
    has_exif = df["exif_datetime"].notna().mean()
    if has_exif > 0.5:
        df = df.sort_values("exif_datetime", na_position="last").reset_index(drop=True)
        ts = pd.to_datetime(df["exif_datetime"])
        gaps = ts.diff().dt.total_seconds().fillna(0).abs()
    else:
        df = df.sort_values("fn_order").reset_index(drop=True)
        gaps = pd.Series([0.0] * len(df))

    df["seq_index"] = range(len(df))
    df["gap_to_prev_secs"] = gaps.values
    return df


# ─── Embeddings ───────────────────────────────────────────────────────────────

class _ImgDataset(Dataset):
    def __init__(self, paths, transform):
        self.paths = paths
        self.transform = transform
        self._blank = Image.new("RGB", (518, 518))

    def __len__(self): return len(self.paths)

    def __getitem__(self, idx):
        try:
            img = Image.open(self.paths[idx]).convert("RGB")
        except Exception:
            img = self._blank
        return self.transform(img), idx


@torch.no_grad()
def extract_embeddings(paths, batch_size=16, progress_cb=None):
    model, transform = get_model()
    ds = _ImgDataset(paths, transform)
    dl = DataLoader(ds, batch_size=batch_size, num_workers=0, pin_memory=(DEVICE == "cuda"))
    embeds = np.zeros((len(paths), 768), dtype=np.float32)
    for i, (imgs, idxs) in enumerate(dl):
        out = F.normalize(model(imgs.to(DEVICE)), dim=1).cpu().numpy()
        for li, gi in enumerate(idxs.tolist()):
            embeds[gi] = out[li]
        if progress_cb:
            progress_cb((i + 1) / len(dl))
    return embeds


# ─── Similarity + Burst Logic ─────────────────────────────────────────────────

def compute_phashes(paths):
    out = []
    for p in paths:
        try:
            out.append(imagehash.phash(Image.open(p).convert("RGB")))
        except Exception:
            out.append(None)
    return out


def phash_sim(h1, h2):
    if h1 is None or h2 is None: return 0.0
    return 1.0 - (h1 - h2) / 64.0


def temp_score(gap, max_gap=5.0):
    if max_gap <= 0: return 0.5
    return float(np.exp(-3.0 * gap / max_gap))


def fuse_scores(embeddings, metadata, phashes, cfg):
    N = len(embeddings)
    fused = np.zeros(N, dtype=np.float32)
    cosines = np.zeros(N, dtype=np.float32)
    gaps = metadata["gap_to_prev_secs"].values
    vw, tw, pw = cfg["visual_weight"], cfg["temporal_weight"], 0.10 if cfg["use_phash"] else 0.0
    total_w = vw + tw + pw
    for i in range(N - 1):
        cos_s = float(np.dot(embeddings[i], embeddings[i + 1]))
        ph_s  = phash_sim(phashes[i], phashes[i + 1]) if cfg["use_phash"] else 0.0
        ts    = temp_score(gaps[i + 1], cfg["temporal_gap_secs"])
        cosines[i] = cos_s
        fused[i]   = (vw * cos_s + tw * ts + pw * ph_s) / total_w
    return fused, cosines


def adaptive_threshold(scores, base=0.88):
    try:
        gm = GaussianMixture(n_components=2, random_state=42).fit(scores.reshape(-1, 1))
        means = sorted(gm.means_.flatten())
        return float(np.clip(np.mean(means), 0.70, 0.97))
    except Exception:
        return base


def segment_bursts(fused, metadata, threshold, max_gap):
    N = len(metadata)
    gaps = metadata["gap_to_prev_secs"].values
    ids, cur = [0] * N, 0
    for i in range(N - 1):
        ids[i] = cur
        if not (fused[i] >= threshold and (gaps[i + 1] <= max_gap or gaps[i + 1] == 0)):
            cur += 1
    ids[N - 1] = cur
    return ids


def refine_dbscan(embeddings, burst_ids, metadata, eps=0.10):
    refined = list(burst_ids)
    max_id = max(burst_ids)
    for bid, group in metadata.groupby("burst_id"):
        idxs = list(group.index)
        if len(idxs) <= 8: continue
        sub = embeddings[idxs]
        dist = np.clip(1.0 - sub @ sub.T, 0, 2)
        labels = DBSCAN(eps=eps, min_samples=2, metric="precomputed").fit(dist).labels_
        for li, gi in enumerate(idxs):
            if labels[li] > 0:
                max_id += 1
                refined[gi] = max_id
            elif labels[li] == -1:
                max_id += 1
                refined[gi] = max_id
    return refined


def select_representative(burst_df, embeddings):
    idxs = burst_df.index.tolist()
    sub = embeddings[idxs]
    c = sub.mean(0, keepdims=True)
    c /= np.linalg.norm(c)
    return idxs[int(np.argmax((sub * c).sum(1)))]


# ─── Cluster Image Copying ────────────────────────────────────────────────────

def export_clusters(metadata, output_root):
    output_root = Path(output_root)
    if output_root.exists():
        shutil.rmtree(output_root)
    output_root.mkdir(parents=True)

    burst_sizes = metadata.groupby("burst_id").size()
    for bid, group in metadata.groupby("burst_id"):
        size = burst_sizes[bid]
        tag = "singleton" if size == 1 else f"burst_{int(bid):04d}_n{size}"
        folder = output_root / tag
        folder.mkdir(exist_ok=True)
        for _, row in group.iterrows():
            src = Path(row["path"])
            dst = folder / src.name
            try:
                shutil.copy2(src, dst)
            except Exception as e:
                logger.warning(f"Copy failed {src}: {e}")

    return str(output_root)


def zip_output(output_root):
    zip_path = str(output_root).rstrip("/") + ".zip"
    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
        for f in Path(output_root).rglob("*"):
            if f.is_file():
                zf.write(f, f.relative_to(Path(output_root).parent))
    return zip_path


# ─── Visualization ────────────────────────────────────────────────────────────

def make_cluster_previews(metadata, n_bursts=6, n_per_burst=5):
    """Return a list of PIL grid images, one per burst."""
    burst_sizes = metadata.groupby("burst_id").size()
    top_bursts  = burst_sizes[burst_sizes > 1].sort_values(ascending=False).head(n_bursts).index.tolist()

    grids = []
    for bid in top_bursts:
        group = metadata[metadata["burst_id"] == bid].head(n_per_burst)
        n     = len(group)
        W, H  = 160, 120
        PAD   = 6
        HEADER = 32
        canvas_w = n * W + (n + 1) * PAD
        canvas_h = H + 2 * PAD + HEADER
        canvas   = Image.new("RGB", (canvas_w, canvas_h), (18, 18, 24))
        draw     = ImageDraw.Draw(canvas)

        # Header text
        draw.rectangle([0, 0, canvas_w, HEADER], fill=(35, 35, 50))
        draw.text((PAD, 7), f"Burst {int(bid)}  ·  {burst_sizes[bid]} images", fill=(200, 200, 220))

        for i, (_, row) in enumerate(group.iterrows()):
            x = PAD + i * (W + PAD)
            y = HEADER + PAD
            try:
                img = Image.open(row["path"]).convert("RGB")
                img.thumbnail((W, H), Image.LANCZOS)
                # center-paste on slot
                slot = Image.new("RGB", (W, H), (30, 30, 40))
                ox   = (W - img.width) // 2
                oy   = (H - img.height) // 2
                slot.paste(img, (ox, oy))
            except Exception:
                slot = Image.new("RGB", (W, H), (60, 30, 30))
            canvas.paste(slot, (x, y))
            # Rep badge
            if row.get("is_representative", False):
                draw.rectangle([x, y, x + 28, y + 14], fill=(255, 180, 0))
                draw.text((x + 3, y + 2), "REP", fill=(0, 0, 0))

        grids.append(canvas)

    return grids


def make_histogram(fused_scores, threshold):
    fig, ax = plt.subplots(figsize=(7, 3), facecolor="#12121a")
    ax.set_facecolor("#1a1a28")
    ax.hist(fused_scores[fused_scores > 0], bins=60, color="#6C63FF", edgecolor="none", alpha=0.85)
    ax.axvline(threshold, color="#FF6584", lw=2, linestyle="--", label=f"Threshold {threshold:.3f}")
    ax.set_xlabel("Fused Similarity Score", color="#aaa", fontsize=9)
    ax.set_ylabel("Pairs", color="#aaa", fontsize=9)
    ax.tick_params(colors="#888")
    for spine in ax.spines.values(): spine.set_visible(False)
    ax.legend(fontsize=8, facecolor="#1a1a28", labelcolor="#ccc")
    plt.tight_layout(pad=0.5)
    fig.canvas.draw()
    buf = np.frombuffer(fig.canvas.buffer_rgba(), dtype=np.uint8)
    buf = buf.reshape(fig.canvas.get_width_height()[::-1] + (4,))
    plt.close(fig)
    return Image.fromarray(buf, mode="RGBA").convert("RGB")


# ─── ZIP Extraction ───────────────────────────────────────────────────────────

# Persistent temp dir so extracted files survive the whole pipeline call
_TMP_EXTRACT_DIR: Optional[str] = None


def extract_zip(zip_path: str) -> str:
    """
    Extract uploaded ZIP into a fresh temp directory.
    Returns the path to the extraction root.
    """
    global _TMP_EXTRACT_DIR
    # Clean up previous extraction
    if _TMP_EXTRACT_DIR and os.path.isdir(_TMP_EXTRACT_DIR):
        shutil.rmtree(_TMP_EXTRACT_DIR, ignore_errors=True)

    _TMP_EXTRACT_DIR = tempfile.mkdtemp(prefix="burst_upload_")
    with zipfile.ZipFile(zip_path, "r") as zf:
        zf.extractall(_TMP_EXTRACT_DIR)
    logger.info(f"ZIP extracted to {_TMP_EXTRACT_DIR}")
    return _TMP_EXTRACT_DIR


# ─── Main Pipeline ────────────────────────────────────────────────────────────

def run_pipeline(
    zip_upload,
    cosine_threshold,
    temporal_gap_secs,
    visual_weight,
    temporal_weight,
    use_phash,
    batch_size,
    progress=gr.Progress(track_tqdm=True),
):
    if zip_upload is None:
        return None, [], "❌ Please upload a ZIP file containing your images.", None

    cfg = dict(
        cosine_threshold  = cosine_threshold,
        temporal_gap_secs = temporal_gap_secs,
        visual_weight     = visual_weight,
        temporal_weight   = temporal_weight,
        use_phash         = use_phash,
    )

    # 1. Extract ZIP
    # Gradio 4.x passes a file-like object with a .name attribute (temp path)
    progress(0.02, desc="Extracting ZIP…")
    try:
        zip_path = zip_upload.name if hasattr(zip_upload, "name") else str(zip_upload)
        input_folder = extract_zip(zip_path)
    except Exception as e:
        return None, [], f"❌ Failed to extract ZIP: {e}", None

    # Output lives next to the extracted temp dir
    output_folder = os.path.join(tempfile.gettempdir(), "burst_output")
    os.makedirs(output_folder, exist_ok=True)

    # 2. Discover
    progress(0.05, desc="Scanning images…")
    paths = discover_images(input_folder)
    if not paths:
        return None, [], "❌ No supported images found in the ZIP.", None

    # 3. Metadata
    progress(0.08, desc="Reading EXIF…")
    meta = build_metadata(paths)

    # 4. Embeddings
    progress(0.12, desc="Extracting DINOv2 embeddings…")
    path_strs = meta["path"].tolist()
    embeddings = extract_embeddings(path_strs, batch_size=int(batch_size))

    # 4. pHash
    phashes = []
    if use_phash:
        progress(0.55, desc="Computing perceptual hashes…")
        phashes = compute_phashes(path_strs)
    else:
        phashes = [None] * len(path_strs)

    # 5. Fuse
    progress(0.62, desc="Fusing similarity scores…")
    fused, cosines = fuse_scores(embeddings, meta, phashes, cfg)
    meta["fused_score_to_next"] = fused
    meta["cosine_to_next"]      = cosines

    # 6. Threshold
    progress(0.68, desc="Computing adaptive threshold…")
    threshold = adaptive_threshold(fused, base=cosine_threshold)

    # 7. Segment
    progress(0.72, desc="Segmenting bursts…")
    burst_ids = segment_bursts(fused, meta, threshold, temporal_gap_secs)
    meta["burst_id"] = burst_ids

    # 8. DBSCAN refine
    progress(0.78, desc="DBSCAN refinement…")
    burst_ids = refine_dbscan(embeddings, burst_ids, meta)
    uid_map = {o: n for n, o in enumerate(sorted(set(burst_ids)))}
    meta["burst_id"] = [uid_map[b] for b in burst_ids]

    # 9. Representatives
    progress(0.83, desc="Selecting representatives…")
    rep_flags = [False] * len(meta)
    for bid, group in meta.groupby("burst_id"):
        if len(group) == 1:
            rep_flags[group.index[0]] = True
        else:
            rep_flags[select_representative(group, embeddings)] = True
    meta["is_representative"] = rep_flags

    # 10. Export clusters
    progress(0.88, desc="Copying images to cluster folders…")
    cluster_root = os.path.join(output_folder, "burst_clusters")
    export_clusters(meta, cluster_root)

    # 11. CSV / JSON
    progress(0.92, desc="Saving CSV & JSON…")
    os.makedirs(output_folder, exist_ok=True)
    csv_path  = os.path.join(output_folder, "burst_results.csv")
    json_path = os.path.join(output_folder, "burst_results.json")
    meta.drop(columns=["phash"], errors="ignore").to_csv(csv_path, index=False)
    bursts_json = {}
    bsizes = meta.groupby("burst_id").size()
    for bid, group in meta.groupby("burst_id"):
        rep = group[group["is_representative"]]
        bursts_json[int(bid)] = {
            "burst_id": int(bid), "size": len(group),
            "is_singleton": len(group) == 1,
            "representative": rep["filename"].iloc[0] if len(rep) else group["filename"].iloc[0],
            "images": group["filename"].tolist(),
        }
    with open(json_path, "w") as f:
        json.dump(bursts_json, f, indent=2, default=str)

    # 12. Zip
    progress(0.95, desc="Zipping cluster folders…")
    zip_path = zip_output(cluster_root)

    # 13. Visualize
    progress(0.97, desc="Generating previews…")
    grids    = make_cluster_previews(meta, n_bursts=8, n_per_burst=6)
    hist_img = make_histogram(fused, threshold)

    # Summary
    bsizes   = meta.groupby("burst_id").size()
    n_multi  = int((bsizes > 1).sum())
    n_single = int((bsizes == 1).sum())
    reduction = (1 - meta["is_representative"].sum() / len(meta)) * 100
    summary = (
        f"✅ Done!\n\n"
        f"**Images processed:** {len(meta):,}\n"
        f"**Bursts detected:** {len(bsizes):,}\n"
        f"**Multi-image bursts:** {n_multi:,}\n"
        f"**Singletons:** {n_single:,}\n"
        f"**Dataset reduction:** {reduction:.1f}%\n"
        f"**Threshold used:** {threshold:.4f}\n"
        f"**Download the ZIP below ↓**"
    )

    progress(1.0, desc="Complete!")
    return hist_img, grids[:6], summary, zip_path


# ─── Gradio UI ────────────────────────────────────────────────────────────────

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

body, .gradio-container {
    background: #0d0d14 !important;
    font-family: 'Syne', sans-serif !important;
    color: #d0d0e0 !important;
}

.gr-panel, .gr-box, .gr-form, .panel {
    background: #13131f !important;
    border: 1px solid #2a2a3f !important;
    border-radius: 12px !important;
}

h1, h2, h3 { font-family: 'Syne', sans-serif !important; }

.title-block {
    text-align: center;
    padding: 28px 0 8px;
}
.title-block h1 {
    font-size: 2.4rem;
    font-weight: 800;
    letter-spacing: -1px;
    background: linear-gradient(135deg, #a78bfa, #f472b6, #fb923c);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    margin: 0;
}
.title-block p {
    color: #6b6b8a;
    font-family: 'DM Mono', monospace;
    font-size: 0.82rem;
    margin-top: 6px;
}

.run-btn {
    background: linear-gradient(135deg, #7c3aed, #db2777) !important;
    border: none !important;
    color: white !important;
    font-family: 'Syne', sans-serif !important;
    font-weight: 600 !important;
    font-size: 1.05rem !important;
    border-radius: 10px !important;
    padding: 14px !important;
    transition: opacity 0.2s !important;
}
.run-btn:hover { opacity: 0.85 !important; }

.dl-btn button {
    background: #1e1e30 !important;
    border: 1px solid #4c4c7a !important;
    color: #a78bfa !important;
    font-family: 'DM Mono', monospace !important;
    border-radius: 8px !important;
}

label { color: #8888aa !important; font-size: 0.82rem !important; font-family: 'DM Mono', monospace !important; }
input, textarea, select { background: #1a1a28 !important; color: #d0d0e0 !important; border-color: #2a2a3f !important; }
.gr-slider input { accent-color: #a78bfa; }
"""

def build_ui():
    with gr.Blocks(css=CSS, title="Wedding Burst Detector") as app:

        gr.HTML("""
        <div class="title-block">
            <h1>📷 Burst Detector</h1>
            <p>DINOv2 · Temporal Fusion · DBSCAN · Production-Grade</p>
        </div>
        """)

        with gr.Row():
            # ── Left: controls ────────────────────────────────────────────
            with gr.Column(scale=1):
                gr.Markdown("### ⚙️ Configuration")

                zip_upload = gr.File(
                    label="Upload Images ZIP",
                    file_types=[".zip"],
                    file_count="single",
                )
                gr.Markdown(
                    "<small style='color:#6b6b8a;font-family:DM Mono,monospace'>"
                    "ZIP your wedding photo folder(s) and drop here. "
                    "Subfolders are scanned recursively.</small>"
                )

                gr.Markdown("#### Similarity")
                cosine_thr = gr.Slider(0.70, 0.98, value=0.88, step=0.01,
                                       label="Cosine Threshold (↑ stricter)")
                temporal_gap = gr.Slider(1.0, 30.0, value=5.0, step=0.5,
                                         label="Max Temporal Gap (seconds)")

                gr.Markdown("#### Weights")
                with gr.Row():
                    vis_w  = gr.Slider(0.0, 1.0, value=0.75, step=0.05, label="Visual Weight")
                    temp_w = gr.Slider(0.0, 1.0, value=0.25, step=0.05, label="Temporal Weight")

                gr.Markdown("#### Performance")
                with gr.Row():
                    batch_sz = gr.Slider(4, 64, value=16, step=4, label="Batch Size")
                    use_ph   = gr.Checkbox(value=True, label="Use pHash")

                run_btn = gr.Button("🚀  Run Burst Detection", elem_classes=["run-btn"])

            # ── Right: outputs ────────────────────────────────────────────
            with gr.Column(scale=2):
                gr.Markdown("### 📊 Results")
                summary_md = gr.Markdown("*Upload a ZIP and click Run.*")
                histogram   = gr.Image(label="Similarity Score Distribution", type="pil", height=220)

                gr.Markdown("### 🖼️ Cluster Previews (top bursts by size)")
                gallery = gr.Gallery(
                    label="Burst Grids",
                    columns=2,
                    height=420,
                    object_fit="contain",
                )

                gr.Markdown("### 💾 Download Clustered ZIP")
                zip_out = gr.File(label="burst_clusters.zip", elem_classes=["dl-btn"])

        # ── Wire up ───────────────────────────────────────────────────────
        run_btn.click(
            fn=run_pipeline,
            inputs=[zip_upload, cosine_thr, temporal_gap, vis_w, temp_w, use_ph, batch_sz],
            outputs=[histogram, gallery, summary_md, zip_out],
        )

    return app


if __name__ == "__main__":
    app = build_ui()
    app.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        inbrowser=True,
    )