"""ALS для неявной обратной связи (Hu, Koren, Volinsky, 2008) на numpy.

На масштабе проекта (~100 пользователей × десятки тысяч треков) решение каждой строки — система 64×64:
весь проход — секунды, без компилируемых зависимостей. Положительная оценка — предпочтение 1 с уверенностью
c = 1 + α·score; отрицательная — предпочтение 0 с повышенной уверенностью (модель учит «не это»).
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np
from numpy.typing import NDArray

F32 = NDArray[np.float32]


@dataclass
class ALSModel:
    users: list[str]
    items: list[str]
    user_factors: F32  # (n_users, k)
    item_factors: F32  # (n_items, k)

    def user_vec(self, user: str) -> F32 | None:
        try:
            i = self.users.index(user)
        except ValueError:
            return None
        vec: F32 = self.user_factors[i]
        return vec

    def scores(self, user: str) -> NDArray[np.float64]:
        u = self.user_vec(user)
        if u is None:
            return np.zeros(len(self.items))
        return (self.item_factors @ u).astype(np.float64)


def _solve_side(
    fixed: F32,
    rows: list[tuple[NDArray[np.intp], NDArray[np.float64], NDArray[np.float64]]],
    reg: float,
) -> F32:
    """Пересчёт одной стороны: для каждой строки (индексы, уверенность, предпочтение) — закрытое решение
    (YᵀY + Yᵀ(C−I)Y + λI) x = YᵀCp, где YᵀY общий для всех строк."""
    k = fixed.shape[1]
    yty = fixed.T @ fixed
    out = np.zeros((len(rows), k), dtype=np.float32)
    eye = reg * np.eye(k)
    for r, (idx, conf, pref) in enumerate(rows):
        if len(idx) == 0:
            continue
        y = fixed[idx].astype(np.float64)
        a = yty + (y.T * (conf - 1.0)) @ y + eye
        b = (y.T * conf) @ pref
        out[r] = np.linalg.solve(a, b)
    return out


def train_als(
    scores: dict[str, dict[str, float]],
    factors: int = 64,
    reg: float = 0.1,
    alpha: float = 20.0,
    iterations: int = 12,
    seed: int = 0,
) -> ALSModel | None:
    """Обучает ALS по implicit_score. None — данных нет."""
    users = sorted(u for u, row in scores.items() if row)
    items = sorted({t for row in scores.values() for t in row})
    if not users or not items:
        return None
    upos = {u: i for i, u in enumerate(users)}
    ipos = {t: i for i, t in enumerate(items)}
    by_user: list[list[tuple[int, float]]] = [[] for _ in users]
    by_item: list[list[tuple[int, float]]] = [[] for _ in items]
    for u, row in scores.items():
        if u not in upos:
            continue
        for t, s in row.items():
            by_user[upos[u]].append((ipos[t], s))
            by_item[ipos[t]].append((upos[u], s))

    def prep(
        lists: list[list[tuple[int, float]]],
    ) -> list[tuple[NDArray[np.intp], NDArray[np.float64], NDArray[np.float64]]]:
        out = []
        for lst in lists:
            idx = np.array([i for i, _ in lst], dtype=np.intp)
            s = np.array([v for _, v in lst], dtype=np.float64)
            conf = 1.0 + alpha * np.abs(s) * np.where(s > 0, 1.0, 0.5)
            pref = (s > 0).astype(np.float64)
            out.append((idx, conf, pref))
        return out

    u_rows, i_rows = prep(by_user), prep(by_item)
    k = min(factors, max(2, len(items) - 1))
    rng = np.random.default_rng(seed)
    x = (rng.standard_normal((len(users), k)) * 0.01).astype(np.float32)
    y = (rng.standard_normal((len(items), k)) * 0.01).astype(np.float32)
    for _ in range(iterations):
        x = _solve_side(y, u_rows, reg)
        y = _solve_side(x, i_rows, reg)
    return ALSModel(users=users, items=items, user_factors=x, item_factors=y)


def normalize(v: F32) -> F32:
    """L2-нормировка строк (косинусное сходство в волне)."""
    n = np.linalg.norm(v, axis=-1, keepdims=True)
    return (v / np.where(n > 0, n, 1.0)).astype(np.float32)