"""Shared pieces: config, encoders, feature caching, projection heads, loss, retrieval metrics."""
from __future__ import annotations
import os, re, json, random, hashlib
from dataclasses import dataclass
import numpy as np, pandas as pd, torch, torch.nn as nn, torch.nn.functional as F, yaml

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))


# ----------------------------------------------------------------------------- config
def _set(d, dotted, val):
    ks = dotted.split(".")
    for k in ks[:-1]:
        d = d.setdefault(k, {})
    d[ks[-1]] = yaml.safe_load(val) if isinstance(val, str) else val


def _get(d, dotted):
    for k in dotted.split("."):
        d = d[k]
    return d


def load_config(path, overrides=()):
    cfg = yaml.safe_load(open(path))
    for o in overrides:  # key.sub=value
        k, v = o.split("=", 1)
        _set(cfg, k, v)
    s = yaml.safe_dump(cfg)
    s = re.sub(r"\$\{([^}]+)\}", lambda m: str(_get(cfg, m.group(1))), s)
    return yaml.safe_load(s)


def rel(p):
    return p if os.path.isabs(p) else os.path.join(ROOT, p)


def seed_all(s):
    random.seed(s); np.random.seed(s); torch.manual_seed(s); torch.cuda.manual_seed_all(s)


def pick_device(name="auto"):
    if name != "auto":
        return torch.device(name)
    if torch.cuda.is_available():
        return torch.device("cuda")
    if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


def safe(name):
    return re.sub(r"[^A-Za-z0-9._-]+", "_", name)


# ----------------------------------------------------------------------------- data
def load_split(cfg, split):
    df = pd.read_parquet(rel(os.path.join(cfg["data"]["dir"], f"{split}.parquet")))
    n = cfg["data"].get("subset")
    if n:
        df = df.iloc[: int(n)].reset_index(drop=True)
    return df


def long_pairs(df, langs):
    """Wide (one row per image, one column per language, NaN = filtered out) -> long (image_idx, lang, text)."""
    rows = []
    for i, r in df.iterrows():
        for l in langs:
            t = r.get(l)
            if isinstance(t, str) and t.strip():
                rows.append((i, l, t))
    return pd.DataFrame(rows, columns=["img_idx", "lang", "text"])


# ----------------------------------------------------------------------------- image encoder (frozen)
class ImageEncoder:
    def __init__(self, cfg, device):
        ic = cfg["image_encoder"]
        self.backend, self.name, self.device = ic["backend"], ic["name"], device
        if self.backend == "hf":
            from transformers import AutoModel, AutoImageProcessor
            self.model = AutoModel.from_pretrained(self.name).to(device).eval()
            self.proc = AutoImageProcessor.from_pretrained(self.name)
        elif self.backend == "open_clip":
            import open_clip
            arch, pre = self.name.split("/", 1)
            self.model, _, self.proc = open_clip.create_model_and_transforms(arch, pretrained=pre)
            self.model = self.model.to(device).eval()
        else:
            raise ValueError(self.backend)

    @torch.no_grad()
    def encode(self, pil_images):
        if self.backend == "hf":
            px = self.proc(images=pil_images, return_tensors="pt")["pixel_values"].to(self.device)
            out = self.model.get_image_features(pixel_values=px)
            if not torch.is_tensor(out):  # some versions return a ModelOutput
                out = getattr(out, "pooler_output", None) if getattr(out, "pooler_output", None) is not None else out[0]
        else:
            px = torch.stack([self.proc(im) for im in pil_images]).to(self.device)
            out = self.model.encode_image(px)
        return out.float().cpu()


def image_features(cfg, df, split, device, batch_size=32):
    """Cached frozen image features for the rows of df, keyed by image_id."""
    from PIL import Image
    ic = cfg["image_encoder"]
    d = rel(os.path.join(ic["cache_dir"], "img_" + safe(ic["name"])))
    os.makedirs(d, exist_ok=True)
    path = os.path.join(d, "all.pt")
    cache = torch.load(path) if os.path.exists(path) else {}
    need = [i for i in df.image_id if i not in cache]
    if need:
        enc = ImageEncoder(cfg, device)
        idir = rel(cfg["data"]["image_dir"])
        import time
        last = time.time()
        for s in range(0, len(need), batch_size):
            ids = need[s: s + batch_size]
            ims = [Image.open(os.path.join(idir, i)).convert("RGB") for i in ids]
            for i, f in zip(ids, enc.encode(ims)):
                cache[i] = f
            if time.time() - last > 20:  # checkpoint the cache so an interrupted run resumes
                print(f"  image feats {s + len(ids)}/{len(need)}", flush=True)
                torch.save(cache, path); last = time.time()
        torch.save(cache, path)
        del enc
    return torch.stack([cache[i] for i in df.image_id])


# ----------------------------------------------------------------------------- text encoder
def text_arm(cfg):
    tc = cfg["text_encoder"]
    a = dict(tc["arms"][tc["arm"]])
    if not a.get("name"):
        raise ValueError(f"text_encoder.arms.{tc['arm']}.name is empty")
    return a


def load_sentence_model(arm, device):
    from sentence_transformers import SentenceTransformer
    kw = {}
    if arm.get("text_only"):  # EmbeddingGemma 2: skip vision/audio towers (270M text-only)
        kw["config_kwargs"] = {"vision_config": None, "audio_config": None}
    if "embeddinggemma" in arm["name"]:
        # EmbeddingGemma: bf16 or fp32 only (fp16 overflows)
        dt = torch.bfloat16 if (device.type == "cuda" and torch.cuda.is_bf16_supported()) else torch.float32
        kw["model_kwargs"] = {"torch_dtype": dt}
    return SentenceTransformer(arm["name"], device=str(device), **kw)


def text_features(cfg, texts, device, batch_size=64):
    """Cached frozen sentence embeddings (L2-normalised by the model's own head if it has one)."""
    arm = text_arm(cfg)
    d = rel(os.path.join(cfg["text_encoder"]["cache_dir"], "txt_" + safe(arm["name"])))
    os.makedirs(d, exist_ok=True)
    path = os.path.join(d, "all.pt")
    cache = torch.load(path) if os.path.exists(path) else {}
    key = lambda t: hashlib.md5((arm.get("prompt", "") + t).encode()).hexdigest()
    need = sorted({t for t in texts if key(t) not in cache})
    if need:
        m = load_sentence_model(arm, device)
        for s in range(0, len(need), 512):
            chunk = need[s: s + 512]
            e = m.encode([arm.get("prompt", "") + t for t in chunk], batch_size=batch_size,
                         convert_to_tensor=True, normalize_embeddings=True).float().cpu()
            for t, v in zip(chunk, e):
                cache[key(t)] = v
            print(f"  text feats {s + len(chunk)}/{len(need)}", flush=True)
            torch.save(cache, path)
        del m
    return torch.stack([cache[key(t)] for t in texts])


class TrainableText(nn.Module):
    """Sentence-transformers model with only its top-N transformer layers trainable."""

    def __init__(self, cfg, device):
        super().__init__()
        self.arm = text_arm(cfg)
        self.st = load_sentence_model(self.arm, device)
        n = cfg["text_encoder"]["unfreeze_last_n_layers"]
        for p in self.st.parameters():
            p.requires_grad = False
        layers = None
        for mod in self.st.modules():  # find the deepest ModuleList of transformer blocks
            for name, child in mod.named_children():
                if isinstance(child, nn.ModuleList) and name in ("layer", "layers") and len(child) > 1:
                    layers = child
        if layers is None:
            raise RuntimeError("could not locate transformer layers to unfreeze")
        for blk in list(layers)[-n:]:
            for p in blk.parameters():
                p.requires_grad = True
        self.max_len = cfg["data"]["max_text_len"]
        self.device = device

    def forward(self, texts):
        texts = [self.arm.get("prompt", "") + t for t in texts]
        feats = (self.st.preprocess if hasattr(self.st, "preprocess") else self.st.tokenize)(texts)
        feats = {k: v[:, : self.max_len].to(self.device) if torch.is_tensor(v) and v.dim() == 2 else v for k, v in feats.items()}
        return self.st(feats)["sentence_embedding"].float()


# ----------------------------------------------------------------------------- heads + loss
class Head(nn.Module):
    def __init__(self, din, dout, hidden=0, dropout=0.0):
        super().__init__()
        self.net = nn.Linear(din, dout) if not hidden else nn.Sequential(
            nn.Linear(din, hidden), nn.GELU(), nn.Dropout(dropout), nn.Linear(hidden, dout))

    def forward(self, x):
        return F.normalize(self.net(x), dim=-1)


class AlignModel(nn.Module):
    def __init__(self, d_img, d_txt, cfg):
        super().__init__()
        m = cfg["model"]
        self.img = Head(d_img, m["proj_dim"], m["proj_hidden"], m["dropout"])
        self.txt = Head(d_txt, m["proj_dim"], m["proj_hidden"], m["dropout"])
        self.logit_scale = nn.Parameter(torch.tensor(float(m["init_logit_scale"])))


def clip_loss(zi, zt, img_ids, logit_scale, mask_same_image=True):
    """Symmetric InfoNCE. With multilingual captions several texts can share an image in one batch;
    those off-diagonal same-image pairs are masked out of the negatives."""
    s = logit_scale.clamp(max=4.6052).exp()
    logits = s * zi @ zt.t()
    if mask_same_image:
        same = img_ids[:, None] == img_ids[None, :]
        eye = torch.eye(len(img_ids), dtype=torch.bool, device=logits.device)
        logits = logits.masked_fill(same & ~eye, float("-inf"))
    y = torch.arange(len(img_ids), device=logits.device)
    return 0.5 * (F.cross_entropy(logits, y) + F.cross_entropy(logits.t(), y))


# ----------------------------------------------------------------------------- metrics
@torch.no_grad()
def retrieval_recall(zi, zt, ks=(1, 5, 10)):
    """zi[i] and zt[i] are a matched image/caption pair (one caption per image within a language)."""
    sim = zi @ zt.t()
    n = sim.shape[0]
    gt = torch.arange(n)
    r_i2t = (sim.argsort(dim=1, descending=True) == gt[:, None]).float().argmax(dim=1)
    r_t2i = (sim.t().argsort(dim=1, descending=True) == gt[:, None]).float().argmax(dim=1)
    out = {"n": n}
    for k in ks:
        out[f"i2t_R@{k}"] = round(100 * (r_i2t < k).float().mean().item(), 2)
        out[f"t2i_R@{k}"] = round(100 * (r_t2i < k).float().mean().item(), 2)
    out["mean_R"] = round(float(np.mean([out[f"{d}_R@{k}"] for d in ("i2t", "t2i") for k in ks])), 2)
    return out


def dump(obj, path):
    os.makedirs(os.path.dirname(path), exist_ok=True)
    json.dump(obj, open(path, "w"), indent=2)