"""
dir009 multimodel ensemble (public LB 0.25048).

Pipeline: flip-TTA embeddings from MegaDescriptor and MiewID,
query expansion in each space, similarity averaging, then HAC clustering.
"""

import math
import os

import numpy as np
import pandas as pd
from PIL import Image
from sklearn.cluster import AgglomerativeClustering
from tqdm import tqdm

import torch
import torch.nn.functional as F
import torchvision.transforms as T
import timm
from transformers import AutoModel


def find_comp_root() -> str:
    # Kaggle can mount competition data under different input roots.
    candidates = [
        "/kaggle/input/animal-clef-2026",
        "/kaggle/input/animal-clef-2026/unzipped",
        "/kaggle/input/competitions/animal-clef-2026",
        "/kaggle/input/competitions/animal-clef-2026/animal-clef-2026",
    ]
    for c in candidates:
        if os.path.isfile(os.path.join(c, "metadata.csv")) and os.path.isfile(os.path.join(c, "sample_submission.csv")):
            return c
    for root, _, files in os.walk("/kaggle/input"):
        if "metadata.csv" in files and "sample_submission.csv" in files:
            return root
    raise FileNotFoundError("Could not locate metadata.csv under /kaggle/input")


# Per-dataset params reused from the best validated setup.
SPECIES_CONFIG = {
    "LynxID2025": {"qe_k": 5, "qe_alpha": 0.5, "thr": 0.35},
    "SalamanderID2025": {"qe_k": 3, "qe_alpha": 0.5, "thr": 0.25},
    "SeaTurtleID2022": {"qe_k": 8, "qe_alpha": 0.5, "thr": 0.45},
    "TexasHornedLizards": {"qe_k": 5, "qe_alpha": 0.5, "thr": 0.30},
}

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
BATCH_SIZE = 64 if DEVICE == "cuda" else 8
NUM_WORKERS = 4


class ImageDataset(torch.utils.data.Dataset):
    def __init__(self, df: pd.DataFrame, root: str, transform):
        self._root = root
        self._paths = df["path"].tolist()
        self._image_ids = df["image_id"].astype(int).tolist()
        self._transform = transform

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

    def __getitem__(self, idx: int):
        p = os.path.join(self._root, self._paths[idx])
        img = Image.open(p).convert("RGB")
        x = self._transform(img)
        return x, self._image_ids[idx]


def _to_feature_tensor(out):
    if hasattr(out, "last_hidden_state"):
        out = out.last_hidden_state[:, 0]
    if isinstance(out, (tuple, list)):
        out = out[0]
    if not torch.is_tensor(out):
        raise TypeError(f"Model output is not tensor: {type(out)}")
    return out


def extract_embeddings_flip_tta(df: pd.DataFrame, model, size: int, root: str) -> np.ndarray:
    transform = T.Compose(
        [
            T.Resize((size, size)),
            T.ToTensor(),
            T.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
        ]
    )

    ds = ImageDataset(df, root, transform)
    loader = torch.utils.data.DataLoader(
        ds,
        batch_size=BATCH_SIZE,
        shuffle=False,
        num_workers=NUM_WORKERS,
        pin_memory=(DEVICE == "cuda"),
    )

    feats = []
    image_ids = []
    amp_ctx = torch.cuda.amp.autocast(enabled=(DEVICE == "cuda"))

    with torch.inference_mode():
        for x, ids in tqdm(loader, total=math.ceil(len(ds) / BATCH_SIZE), desc=f"embed{size}", leave=False):
            x = x.to(DEVICE, non_blocking=True)
            # Horizontal flip TTA (same model, two views).
            x_flip = torch.flip(x, dims=[3])

            with amp_ctx:
                out0 = _to_feature_tensor(model(x))
                out1 = _to_feature_tensor(model(x_flip))
                out = out0 + out1

            out = out.float()
            out = F.normalize(out, p=2, dim=1)
            feats.append(out.cpu())

            if torch.is_tensor(ids):
                image_ids.extend(ids.tolist())
            else:
                image_ids.extend(list(ids))

    feats = torch.cat(feats, dim=0).numpy().astype(np.float32)

    df_ids = df["image_id"].astype(int).tolist()
    if image_ids != df_ids:
        raise RuntimeError("DataLoader returned image_ids in unexpected order")

    return feats


def query_expansion(feats: np.ndarray, k: int = 5, alpha: float = 0.5) -> np.ndarray:
    n = int(feats.shape[0])
    if n == 0:
        return feats
    k = int(min(max(k, 1), n))

    sim = (feats @ feats.T).astype(np.float32)
    idx = np.argpartition(sim, -k, axis=1)[:, -k:]

    expanded = feats + float(alpha) * feats[idx].mean(axis=1)
    expanded = expanded / np.clip(np.linalg.norm(expanded, axis=1, keepdims=True), 1e-12, None)
    return expanded.astype(np.float32)


def distance_from_similarity(sim: np.ndarray) -> np.ndarray:
    dist = np.clip(1.0 - sim.astype(np.float32), 0.0, 1.0)
    np.fill_diagonal(dist, 0.0)
    return dist


def cluster_agglomerative(dist: np.ndarray, thr: float, linkage: str = "average") -> np.ndarray:
    kwargs = dict(n_clusters=None, linkage=linkage, distance_threshold=float(thr))
    try:
        model = AgglomerativeClustering(metric="precomputed", **kwargs)
    except TypeError:
        model = AgglomerativeClustering(affinity="precomputed", **kwargs)
    return model.fit_predict(dist).astype(int)


def main():
    torch.backends.cudnn.benchmark = True

    root = find_comp_root()
    print(f"Using competition input root: {root}")

    meta = pd.read_csv(os.path.join(root, "metadata.csv"))
    meta_test = meta[meta["split"] == "test"].copy()
    assert meta_test["image_id"].is_unique

    sample = pd.read_csv(os.path.join(root, "sample_submission.csv"))

    print("Loading models from HuggingFace (internet enabled)")
    mega = timm.create_model("hf-hub:BVRA/MegaDescriptor-L-384", pretrained=True)
    if hasattr(mega, "reset_classifier"):
        mega.reset_classifier(0)
    mega.eval().to(DEVICE)

    miew = AutoModel.from_pretrained("conservationxlabs/miewid-msv3", trust_remote_code=True)
    miew.eval().to(DEVICE)

    results = []
    for dataset_name, df_ds in meta_test.groupby("dataset"):
        df_ds = df_ds.reset_index(drop=True)
        cfg = SPECIES_CONFIG[dataset_name]
        print(f"Dataset {dataset_name}: n={len(df_ds)}")

        feats_mega = extract_embeddings_flip_tta(df_ds, mega, 384, root)
        feats_miew = extract_embeddings_flip_tta(df_ds, miew, 512, root)

        # Apply the same QE params to both views.
        feats_mega = query_expansion(feats_mega, k=int(cfg['qe_k']), alpha=float(cfg["qe_alpha"]))
        feats_miew = query_expansion(feats_miew, k=int(cfg['qe_k']), alpha=float(cfg["qe_alpha"]))

        # Fuse similarities instead of raw features to keep each model in its own space.
        sim = 0.5 * (feats_mega @ feats_mega.T) + 0.5 * (feats_miew @ feats_miew.T)
        dist = distance_from_similarity(sim)

        labels = cluster_agglomerative(dist, thr=float(cfg['thr']), linkage="average")
        n_clusters = int(np.unique(labels).shape[0])
        print(f"  clusters={n_clusters} thr={cfg['thr']} qe_k={cfg['qe_k']}")

        results.append(
            pd.DataFrame(
                {
                    "image_id": df_ds["image_id"].astype(int).values,
                    "cluster": [f"cluster_{dataset_name}_{c}" for c in labels],
                }
            )
        )

    pred = pd.concat(results, ignore_index=True)
    pred_map = pred.set_index("image_id")["cluster"]

    sample["cluster"] = sample["image_id"].astype(int).map(pred_map)
    if sample["cluster"].isna().any():
        missing = sample.loc[sample["cluster"].isna(), "image_id"].head(20).tolist()
        raise RuntimeError(f"Missing predictions for image_ids: {missing}")

    out_path = "submission.csv"
    sample.to_csv(out_path, index=False)
    print(f"Wrote {out_path} with {len(sample)} rows")
    print(sample.head())


if __name__ == "__main__":
    main()