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

from __future__ import annotations

import math
import time
from dataclasses import dataclass, field
from typing import Protocol

import numpy as np
from numpy.typing import NDArray

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

MODEL_VER = "heuristic-v1"
# С коллаборативным сигналом (ALS) — суффикс к версии модели в показах.
CF_SUFFIX = "+cf"

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
    cf: Vec = field(default_factory=lambda: np.zeros((0, 0), dtype=np.float32))  # (n, k) ALS, нормирован
    has_cf: NDArray[np.bool_] = field(default_factory=lambda: np.zeros(0, dtype=bool))
    pos: dict[str, int] = field(default_factory=dict)

    def __post_init__(self) -> None:
        self.pos = {t: i for i, t in enumerate(self.ids)}
        if self.has_cf.shape[0] != len(self.ids):
            self.cf = np.zeros((len(self.ids), 1), dtype=np.float32)
            self.has_cf = np.zeros(len(self.ids), dtype=bool)

    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,
        raw_cf: list[Vec | None] | None = None,
    ) -> 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)
        cf_dim = next((len(v) for v in raw_cf or [] if v is not None), 0)
        cf = np.zeros((n, max(cf_dim, 1)), dtype=np.float32)
        has_cf = np.zeros(n, dtype=bool)
        for i, v in enumerate(raw_cf or []):
            if v is not None and len(v) == cf_dim:
                norm = float(np.linalg.norm(v))
                if norm > 0:
                    cf[i], has_cf[i] = v / norm, True
        return CatalogIndex(
            ids=ids, artist=artist, emb=emb, has_emb=has, features=feats, popularity=pop, cf=cf, has_cf=has_cf
        )


@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
    taste_cf: list[float] | None = None  # фактор пользователя ALS
    bandit: dict[str, list[float]] = field(
        default_factory=dict
    )  # источник → [α, β] (docs/recommendations.md#6.6)
    recent_likes: list[str] = field(default_factory=list)  # последние лайки, новые первыми (для item_cf)


@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
    features: list[float] | None = None  # вектор признаков в момент показа (для обучения ранжировщика)


# Источники кандидатов (one-hot в признаках ранжировщика; порядок фиксирован).
SOURCES = ["seed", "fresh", "audio_sim", "cf", "item_cf", "artist", "rediscover", "popular", "explore"]

# Признаки ранжировщика v2 (docs/recommendations.md#6.3). Порядок — часть формата модели:
# менять только вместе с FEATURES_VERSION (старые модели тогда не загружаются).
FEATURE_NAMES = [
    "sim_session",
    "sim_taste",
    "cf",
    "artist_affinity",
    "popularity",
    "v1_score",
    "context_fit",
    "is_new",
    "liked_recent",
    "bandit_mean",
    *[f"src_{s}" for s in SOURCES],
    "energy",
    "bpm",
    "danceability",
    "valence",
    "hour_sin",
    "hour_cos",
    "position",
    "skips_in_row",
    "has_activity",
    "has_mood",
]
FEATURES_VERSION = 1


class Ranker(Protocol):
    """Обученный ранжировщик: оценки кандидатов по матрице признаков (строки — кандидаты)."""

    version: str

    def predict(self, x: NDArray[np.float64]) -> NDArray[np.float64]: ...


@dataclass
class Context:
    """Контекст показа для признаков ранжировщика."""

    hour: float = field(default_factory=lambda: time.localtime().tm_hour + time.localtime().tm_min / 60)
    has_activity: bool = False
    has_mood: bool = False


# Вес бонуса бандита к оценке (θ − 0.5 ∈ [−0.5, 0.5]).
BANDIT_WEIGHT = 0.2

# Вклад реакции в сессионный вектор: 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 cf_scores(index: CatalogIndex, profile: Profile) -> NDArray[np.float64] | None:
    """Косинус треков к фактору пользователя ALS; None — нет модели или пользователь в неё не попал."""
    if profile.taste_cf is None or not index.has_cf.any():
        return None
    u = np.array(profile.taste_cf, dtype=np.float32)
    if u.shape[0] != index.cf.shape[1]:
        return None
    u = u / (float(np.linalg.norm(u)) or 1.0)
    return np.where(index.has_cf, index.cf @ u, 0.0).astype(np.float64)


def model_version(index: CatalogIndex, profile: Profile, ranker: Ranker | None = None) -> str:
    if ranker is not None:
        return ranker.version
    return MODEL_VER + (CF_SUFFIX if cf_scores(index, profile) is not None else "")


def candidate_features(
    index: CatalogIndex,
    profile: Profile,
    session: Session,
    cand: NDArray[np.intp],
    sources: dict[int, str],
    sim_session: NDArray[np.float64],
    sim_taste: NDArray[np.float64],
    cf_score: NDArray[np.float64] | None,
    artist_aff: NDArray[np.float64],
    v1: NDArray[np.float64],
    context_fit: NDArray[np.float64],
    ctx: Context,
) -> NDArray[np.float64]:
    """Матрица признаков (кандидаты × FEATURE_NAMES)."""
    m = len(cand)
    x = np.zeros((m, len(FEATURE_NAMES)), dtype=np.float64)

    def feat(name: str) -> NDArray[np.float64]:
        v = index.features.get(name)
        if v is None:
            return np.full(m, -1.0)
        return np.nan_to_num(v[cand], nan=-1.0)

    col = {n: k for k, n in enumerate(FEATURE_NAMES)}
    x[:, col["sim_session"]] = sim_session[cand]
    x[:, col["sim_taste"]] = sim_taste[cand]
    x[:, col["cf"]] = cf_score[cand] if cf_score is not None else 0.0
    x[:, col["artist_affinity"]] = artist_aff[cand]
    x[:, col["popularity"]] = index.popularity[cand]
    x[:, col["v1_score"]] = v1
    x[:, col["context_fit"]] = context_fit
    x[:, col["is_new"]] = [index.ids[i] not in profile.played for i in cand]
    x[:, col["liked_recent"]] = [index.ids[i] in profile.liked_recent for i in cand]
    for j, i in enumerate(cand):
        src = sources[int(i)]
        ab = profile.bandit.get(src)
        x[j, col["bandit_mean"]] = ab[0] / (ab[0] + ab[1]) if ab else 0.5
        if src in SOURCES:
            x[j, col[f"src_{src}"]] = 1.0
    x[:, col["energy"]] = feat("energy")
    bpm = feat("bpm")
    x[:, col["bpm"]] = np.where(bpm > 0, bpm / 200.0, -1.0)
    x[:, col["danceability"]] = feat("danceability")
    x[:, col["valence"]] = feat("valence")
    x[:, col["hour_sin"]] = math.sin(2 * math.pi * ctx.hour / 24)
    x[:, col["hour_cos"]] = math.cos(2 * math.pi * ctx.hour / 24)
    x[:, col["position"]] = min(len(session.served) / 50.0, 1.0)
    x[:, col["skips_in_row"]] = min(session.skips_in_row / 3.0, 1.0)
    x[:, col["has_activity"]] = float(ctx.has_activity)
    x[:, col["has_mood"]] = float(ctx.has_mood)
    return x


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,
    ranker: Ranker | None = None,
    ctx: Context | None = None,
    friend_ids: list[str] | 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)
    cf_score = cf_scores(index, profile)

    # Привязанность к исполнителю — сглаженное среднее оценок его треков.
    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")
    friends = [index.pos[t] for t in (friend_ids or []) if t in index.pos and allowed[index.pos[t]]][:10]
    add(friends, "friends")
    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")
    # Коллективный опыт (ALS): похожие по вкусу других слушателей и «с этим слушают» к последним лайкам.
    if cf_score is not None:
        add(_top(cf_score, allowed & index.has_cf, 100), "cf")
        likes = [
            index.pos[t] for t in profile.recent_likes[:10] if t in index.pos and index.has_cf[index.pos[t]]
        ]
        if likes:
            item_sim = (index.cf @ index.cf[likes].T).max(axis=1).astype(np.float64)
            add(_top(np.where(index.has_cf, item_sim, -1.0), allowed & index.has_cf, 60), "item_cf")
    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 (она же — признак и запасной вариант для v2).
    score = (
        0.35 * sim_session[cand]
        + 0.25 * sim_taste[cand]
        + 0.15 * artist_aff[cand]
        + 0.05 * index.popularity[cand]
        + (0.10 * cf_score[cand] if cf_score is not None else 0.0)
        + 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
    # Бандит: семплирование Томпсона по источникам — удачные для пользователя источники получают бонус.
    # Источник без наблюдений — нейтральные 0.5 (без шума, иначе холодный старт ранжировался бы случайно).
    theta: dict[str, float] = {}
    for src in set(sources.values()):
        ab = profile.bandit.get(src)
        theta[src] = float(rng.beta(ab[0], ab[1])) if ab else 0.5
    score = score + BANDIT_WEIGHT * (np.array([theta[sources[int(i)]] for i in cand]) - 0.5)
    v1 = score.copy()

    # 4. Соответствие контексту; если отсеялось слишком много — порог ослабляется.
    f = np.ones(len(cand))
    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, v1 = cand[keep], score[keep], f[keep], v1[keep]
        score = score + 0.5 * f  # контекст важнее вкуса, когда выбран явно
    if len(cand) == 0:
        return []

    # Признаки кандидатов: пишутся с показом (обучение v2) и подаются обученному ранжировщику.
    ctx = ctx or Context()
    f64 = np.float64
    x = candidate_features(
        index,
        profile,
        session,
        cand,
        sources,
        np.asarray(sim_session, f64),
        np.asarray(sim_taste, f64),
        cf_score,
        np.asarray(artist_aff, f64),
        np.asarray(v1, f64),
        np.asarray(f, f64),
        ctx,
    )
    if ranker is not None:
        # v2: порядок задаёт модель; в MMR нужна шкала 0..1 — берём ранг предсказания.
        pred = ranker.predict(x)
        score = np.argsort(np.argsort(pred)) / max(len(pred) - 1, 1)

    # Social influence is small, context-sensitive, and survives the learned ranker.
    score = score + 0.06 * np.isin(cand, friends) * f
    # 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])
            if sources[i] == "friends" and any(p.source == "friends" for p in picks):
                continue
            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}
                social_chosen = any(p.source == "friends" for p in picks)
                rest = [j for j in order if index.artist[int(cand[j])] in chosen_artists
                        and not (sources[int(cand[j])] == "friends" and social_chosen)]
                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),
                features=x[best_j].tolist(),
            )
        )

    # Исследование: «новых» (неслышанных) меньше доли 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), features=x[j].tolist()
        )
    return picks