"""Сборка пачки «Моей волны» (v1: эвристика). Чистые функции над массивами — без БД и сети."""

from __future__ import annotations

from dataclasses import dataclass, field

import numpy as np
from numpy.typing import NDArray

from pipu_rec.wave.contexts import FEATURES, Rule, fit

MODEL_VER = "heuristic-v1"

Vec = NDArray[np.float32]


@dataclass
class CatalogIndex:
    """Каталог в памяти: строка i — трек ids[i]."""

    ids: list[str]
    artist: list[str]  # основной исполнитель
    emb: Vec  # (n, d), стандартизирован и нормирован; нулевой — нет анализа
    has_emb: NDArray[np.bool_]
    features: dict[str, NDArray[np.float64]]  # NaN — нет признака
    popularity: NDArray[np.float64]  # 0..1
    pos: dict[str, int] = field(default_factory=dict)

    def __post_init__(self) -> None:
        self.pos = {t: i for i, t in enumerate(self.ids)}

    def __len__(self) -> int:
        return len(self.ids)

    @staticmethod
    def build(
        ids: list[str],
        artist: list[str],
        raw_emb: list[Vec | None],
        features: dict[str, list[float | None]],
        plays: list[int],
        dim: int = 256,
    ) -> CatalogIndex:
        n = len(ids)
        emb = np.zeros((n, dim), dtype=np.float32)
        has = np.zeros(n, dtype=bool)
        for i, e in enumerate(raw_emb):
            if e is not None:
                emb[i], has[i] = e, True
        if has.sum() >= 2:
            # Стандартизация по библиотеке: косинус сравнивает отличия треков, а не общий «звук библиотеки».
            mu = emb[has].mean(axis=0)
            sd = emb[has].std(axis=0) + 1e-6
            emb[has] = (emb[has] - mu) / sd
            emb[has] /= np.linalg.norm(emb[has], axis=1, keepdims=True) + 1e-9
        feats = {
            k: np.array([np.nan if v is None else v for v in features.get(k, [None] * n)], dtype=np.float64)
            for k in FEATURES
        }
        p = np.log1p(np.array(plays, dtype=np.float64))
        pop = p / p.max() if n and p.max() > 0 else np.zeros(n)
        return CatalogIndex(ids=ids, artist=artist, emb=emb, has_emb=has, features=feats, popularity=pop)


@dataclass
class Profile:
    """Вкус пользователя: оценки треков и производные."""

    scores: dict[str, float]  # implicit_score > 0 — нравится
    disliked: set[str]
    disliked_artists: set[str]
    recent: set[str]  # сыграны за последние 3 часа
    played: set[str]  # слышал когда-либо
    liked_recent: set[str]  # лайкнутые, игравшие за 24 ч — штраф
    rediscover: set[str]  # любимые, не игравшие > 60 дней
    exploration: float = 0.15


@dataclass
class Session:
    """Состояние сессии волны (Redis)."""

    vec: list[float] | None = None
    served: list[str] = field(default_factory=list)
    recent_artists: list[str] = field(default_factory=list)
    skips_in_row: int = 0


@dataclass(frozen=True)
class Pick:
    id: str
    source: str
    score: float


# Вклад реакции в сессионный вектор: s ← normalize(s + η·r·e).
REACTION_WEIGHT = {"complete": 1.0, "like": 2.0, "skip_early": -1.0, "dislike": -2.0, "skip": -0.5}
ETA = 0.3


def taste_centroids(index: CatalogIndex, profile: Profile, max_k: int = 6, seed: int = 0) -> Vec:
    """Центроиды вкуса: k-means по эмбеддингам треков с положительной оценкой (веса = оценка).
    k растёт с историей: 1 центроид на ~8 любимых треков, не больше max_k."""
    rows = [(index.pos[t], s) for t, s in profile.scores.items() if s > 0 and t in index.pos]
    rows = [(i, s) for i, s in rows if index.has_emb[i]]
    if not rows:
        return np.zeros((0, index.emb.shape[1]), dtype=np.float32)
    x = index.emb[[i for i, _ in rows]]
    w = np.array([s for _, s in rows])
    k = int(min(max_k, max(1, len(rows) // 8)))
    rng = np.random.default_rng(seed)
    centers = x[rng.choice(len(x), size=k, replace=False)]
    for _ in range(10):
        assign = np.argmax(x @ centers.T, axis=1)
        for c in range(k):
            m = assign == c
            if m.any():
                v = (x[m] * w[m, None]).sum(axis=0)
                centers[c] = v / (np.linalg.norm(v) + 1e-9)
    return centers.astype(np.float32)


def update_session(
    index: CatalogIndex, session: Session, feedback: list[tuple[str, str]], centroids: Vec
) -> None:
    """Обновляет сессионный вектор и счётчик скипов по реакциям на выданные треки.
    Три скипа подряд — «смена курса»: вектор переводится к наименее похожему центроиду вкуса."""
    d = index.emb.shape[1]
    vec = np.array(session.vec, dtype=np.float32) if session.vec else np.zeros(d, dtype=np.float32)
    for track_id, outcome in feedback:
        i = index.pos.get(track_id)
        r = REACTION_WEIGHT.get(outcome)
        if r is None:
            continue
        session.skips_in_row = session.skips_in_row + 1 if outcome in ("skip", "skip_early", "dislike") else 0
        if i is None or not index.has_emb[i]:
            continue
        vec = vec + ETA * r * index.emb[i]
        norm = float(np.linalg.norm(vec))
        vec = vec / norm if norm > 0 else vec
    if session.skips_in_row >= 3 and len(centroids):
        sims = centroids @ vec if np.any(vec) else np.zeros(len(centroids))
        vec = centroids[int(np.argmin(sims))].copy()
        session.skips_in_row = 0
    session.vec = vec.tolist() if np.any(vec) else None


def _top(scores: NDArray[np.float64], mask: NDArray[np.bool_], k: int) -> NDArray[np.intp]:
    idx = np.nonzero(mask)[0]
    if len(idx) <= k:
        return idx
    part = idx[np.argpartition(-scores[idx], k)[:k]]
    return part


def build_batch(
    index: CatalogIndex,
    profile: Profile,
    session: Session,
    rules: list[Rule],
    n: int = 5,
    exclude: set[str] | None = None,
    seed_ids: list[str] | None = None,
    rng: np.random.Generator | None = None,
) -> list[Pick]:
    """Кандидаты → фильтры → оценка → соответствие контексту → MMR с правилами разнообразия."""
    rng = rng or np.random.default_rng()
    size = len(index)
    if size == 0:
        return []
    centroids = taste_centroids(index, profile)
    svec = np.array(session.vec, dtype=np.float32) if session.vec else None

    # 2. Фильтры.
    excluded = profile.disliked | profile.recent | set(session.served) | (exclude or set())
    allowed = np.array(
        [
            t not in excluded and a not in profile.disliked_artists
            for t, a in zip(index.ids, index.artist, strict=True)
        ]
    )
    if not allowed.any():
        allowed = np.array([t not in profile.disliked for t in index.ids])
    if not allowed.any():
        return []

    # Сходства со вкусом и сессией.
    sim_taste = (index.emb @ centroids.T).max(axis=1) if len(centroids) else np.zeros(size)
    sim_session = index.emb @ svec if svec is not None else np.zeros(size)
    sim_taste = np.where(index.has_emb, sim_taste, 0.0)
    sim_session = np.where(index.has_emb, sim_session, 0.0)

    # Привязанность к исполнителю — сглаженное среднее оценок его треков.
    totals: dict[str, list[float]] = {}
    for t, s in profile.scores.items():
        i = index.pos.get(t)
        if i is not None:
            totals.setdefault(index.artist[i], []).append(s)
    affinity = {a: sum(v) / (len(v) + 2) for a, v in totals.items()}
    artist_aff = np.array([np.tanh(affinity.get(a, 0.0) / 3) for a in index.artist])

    # 1. Источники кандидатов.
    sources: dict[int, str] = {}

    def add(idxs: NDArray[np.intp] | list[int], source: str) -> None:
        for i in idxs:
            sources.setdefault(int(i), source)

    if seed_ids:
        seeds = [index.pos[s] for s in seed_ids if s in index.pos and index.has_emb[index.pos[s]]]
        if seeds:
            q = index.emb[seeds].mean(axis=0)
            add(_top((index.emb @ q).astype(np.float64), allowed & index.has_emb, 100), "seed")
    query = np.maximum(sim_session, sim_taste) if svec is not None or len(centroids) else None
    if query is not None:
        # «Новое» помечаем раньше audio_sim: неслышанный трек рядом со вкусом считается исследованием.
        never = np.array([t not in profile.played for t in index.ids])
        add(_top(sim_taste + 0.3 * sim_session, allowed & index.has_emb & never, 50), "fresh")
        add(_top(query, allowed & index.has_emb, 150), "audio_sim")
    liked_artists = {a for a, v in affinity.items() if v > 0.3}
    if liked_artists:
        by_artist = [int(i) for i in np.nonzero(allowed)[0] if index.artist[i] in liked_artists]
        add(rng.permutation(by_artist)[:60].tolist() if by_artist else [], "artist")
    redis_ = [index.pos[t] for t in profile.rediscover if t in index.pos and allowed[index.pos[t]]]
    add(rng.permutation(redis_)[:30].tolist() if redis_ else [], "rediscover")
    add(_top(index.popularity + rng.random(size) * 0.3, allowed, 30), "popular")
    # Холодный старт и исследование: случайные треки из разных частей библиотеки.
    add(rng.permutation(np.nonzero(allowed)[0])[:40].tolist(), "explore")

    cand = np.array(sorted(sources))
    # 3. Оценка v1.
    score = (
        0.35 * sim_session[cand]
        + 0.25 * sim_taste[cand]
        + 0.15 * artist_aff[cand]
        + 0.05 * index.popularity[cand]
        + 0.02 * rng.random(len(cand))  # разбивает ничьи, особенно при холодном старте
    )
    liked_recent = np.array([index.ids[i] in profile.liked_recent for i in cand])
    score = score - 0.2 * liked_recent

    # 4. Соответствие контексту; если отсеялось слишком много — порог ослабляется.
    if rules:
        f = fit({k: v[cand] for k, v in index.features.items()}, rules)
        # Нужен запас для разнообразия, но не больше 50 и не больше четверти пула.
        enough = min(50, max(3 * n, len(cand) // 4))
        for threshold in (0.35, 0.15, 0.0):
            keep = f >= threshold
            if keep.sum() >= enough or threshold == 0.0:
                break
        cand, score, f = cand[keep], score[keep], f[keep]
        score = score + 0.5 * f  # контекст важнее вкуса, когда выбран явно
    if len(cand) == 0:
        return []

    # 5. MMR + правила: 1 трек исполнителя на пачку, не больше 2 на последние 10; доля «нового».
    picks: list[Pick] = []
    chosen: list[int] = []
    recent_artists = session.recent_artists[-10:]
    want_fresh = max(0, round(n * profile.exploration))
    order = list(range(len(cand)))
    lam = 0.75
    while len(picks) < n and order:
        best_j, best_v = -1, -np.inf
        for j in order:
            i = int(cand[j])
            a = index.artist[i]
            if any(index.artist[c] == a for c in chosen) or recent_artists.count(a) >= 2:
                continue
            redundancy = max((float(index.emb[i] @ index.emb[c]) for c in chosen), default=0.0)
            v = lam * float(score[j]) - (1 - lam) * redundancy
            if v > best_v:
                best_j, best_v = j, v
        if best_j < 0:
            # Правила разнообразия не дают набрать пачку — ослабляем их, но не повторяем треки.
            if recent_artists:
                recent_artists = []
                continue
            if chosen:
                chosen_artists = {index.artist[c] for c in chosen}
                rest = [j for j in order if index.artist[int(cand[j])] in chosen_artists]
                if not rest:
                    break
                best_j = max(rest, key=lambda j: score[j])
            else:
                break
        i = int(cand[best_j])
        chosen.append(i)
        order.remove(best_j)
        picks.append(Pick(id=index.ids[i], source=sources[i], score=round(float(score[best_j]), 4)))

    # Исследование: «новых» (неслышанных) меньше доли exploration — заменяем ими слабейшие выборы.
    fresh_left = sorted(
        (j for j in order if sources[int(cand[j])] == "fresh"), key=lambda j: score[j], reverse=True
    )
    while sum(p.source == "fresh" for p in picks) < want_fresh and fresh_left:
        j = fresh_left.pop(0)
        i = int(cand[j])
        others = [k for k, p in enumerate(picks) if p.source != "fresh"]
        if not others:
            break
        weakest = min(others, key=lambda k: picks[k].score)
        if any(index.artist[index.pos[p.id]] == index.artist[i] for k, p in enumerate(picks) if k != weakest):
            continue
        picks[weakest] = Pick(id=index.ids[i], source="fresh", score=round(float(score[j]), 4))
    return picks