"""miru-codev3-dual: Tokenlearn ⊕ distill_fuse score-fuse via 512-d concat.

Cosine on unit concat equals linear score mix:
  cos(e_q, e_d) = α · cos(A_q, A_d) + (1-α) · cos(B_q, B_d)

Default α=0.5 → e = normalize(concat(√α e_A, √(1-α) e_B)).

Usage (local folder or Hub repo with tokenlearn/ + distill_fuse/):
  from dual_encode import DualConcatEncoder
  enc = DualConcatEncoder.from_pretrained("Takara-DS1/miru-codev3-dual")
  vecs = enc.encode(["def add(a, b): return a + b"])
"""

from __future__ import annotations

import json
from pathlib import Path
from typing import Sequence

import numpy as np
from huggingface_hub import hf_hub_download, snapshot_download
from model2vec import StaticModel


class DualConcatEncoder:
    """Two Model2Vec StaticModels fused by √α-scaled concat + L2 norm → 512-d."""

    def __init__(
        self,
        model_a: StaticModel,
        model_b: StaticModel,
        *,
        alpha: float = 0.5,
        max_length: int = 512,
    ) -> None:
        if model_a.embedding.shape[1] != model_b.embedding.shape[1]:
            raise ValueError(
                f"dim mismatch: {model_a.embedding.shape[1]} vs {model_b.embedding.shape[1]}"
            )
        if not (0.0 <= float(alpha) <= 1.0):
            raise ValueError("alpha must be in [0, 1]")
        self.a = model_a
        self.b = model_b
        self.alpha = float(alpha)
        self.max_length = int(max_length)
        self.dimensions = int(model_a.embedding.shape[1] * 2)

    @classmethod
    def from_pretrained(
        cls,
        repo_or_path: str,
        *,
        alpha: float | None = None,
        tokenlearn_subdir: str = "tokenlearn",
        distill_fuse_subdir: str = "distill_fuse",
        max_length: int = 512,
        revision: str | None = None,
    ) -> DualConcatEncoder:
        root = Path(repo_or_path)
        if root.is_dir() and (root / tokenlearn_subdir).is_dir():
            base = root
        else:
            base = Path(
                snapshot_download(
                    repo_id=repo_or_path,
                    revision=revision,
                    allow_patterns=[
                        f"{tokenlearn_subdir}/*",
                        f"{distill_fuse_subdir}/*",
                        "dual_config.json",
                        "README.md",
                    ],
                )
            )

        cfg_path = base / "dual_config.json"
        cfg: dict = {}
        if cfg_path.exists():
            cfg = json.loads(cfg_path.read_text(encoding="utf-8"))
        use_alpha = float(alpha if alpha is not None else cfg.get("alpha", 0.5))

        a = StaticModel.from_pretrained(str(base / tokenlearn_subdir))
        b = StaticModel.from_pretrained(str(base / distill_fuse_subdir))
        return cls(a, b, alpha=use_alpha, max_length=max_length)

    @classmethod
    def from_branches(
        cls,
        tokenlearn_id: str,
        distill_fuse_id: str,
        *,
        alpha: float = 0.5,
        max_length: int = 512,
    ) -> DualConcatEncoder:
        """Load the two bags from separate Hub repos / local folders."""
        a = StaticModel.from_pretrained(tokenlearn_id)
        b = StaticModel.from_pretrained(distill_fuse_id)
        return cls(a, b, alpha=alpha, max_length=max_length)

    def fuse(self, ea: np.ndarray, eb: np.ndarray) -> np.ndarray:
        sa = np.sqrt(self.alpha)
        sb = np.sqrt(1.0 - self.alpha)
        out = np.concatenate([sa * ea, sb * eb], axis=-1)
        out /= np.linalg.norm(out, axis=-1, keepdims=True) + 1e-12
        return out.astype(np.float32)

    def encode(
        self,
        texts: str | Sequence[str],
        *,
        max_length: int | None = None,
        batch_size: int | None = None,
    ) -> np.ndarray:
        if isinstance(texts, str):
            texts = [texts]
        texts = list(texts)
        ml = max_length or self.max_length
        # model2vec encode accepts list[str]; batching is internal
        kwargs = {"max_length": ml}
        if batch_size is not None:
            kwargs["batch_size"] = batch_size
        ea = np.asarray(self.a.encode(texts, **kwargs), dtype=np.float32)
        eb = np.asarray(self.b.encode(texts, **kwargs), dtype=np.float32)
        if ea.ndim == 1:
            ea = ea[None, :]
            eb = eb[None, :]
        return self.fuse(ea, eb)


def _smoke() -> None:
    here = Path(__file__).resolve().parent
    enc = DualConcatEncoder.from_pretrained(str(here))
    v = enc.encode(["def add(a, b):\n    return a + b"])
    assert v.shape == (1, 512), v.shape
    assert abs(float(np.linalg.norm(v[0])) - 1.0) < 1e-3
    print(f"ok dual encode shape={v.shape} alpha={enc.alpha}")


if __name__ == "__main__":
    _smoke()