"""The recommender: a small matrix factorisation trained by alternating least squares.

It only ever sees (user, item, shown n times, clicked s times) for items that were actually
shown. Predicted score = item bias + user vector . item vector, trained to fit the click rate
s/n on the shown pairs. Items nobody has been shown have no rows at all, so their bias and
vector shrink to exactly zero: they are predicted to be worth nothing, which is the whole
feedback loop in one line.
"""
from __future__ import annotations

import numpy as np

from .rng import STREAMS, stream_key, uniforms


class ALS:
    def __init__(self, n_users: int, n_items: int, factors: int = 6, reg: float = 1.0, sweeps: int = 2, seed: int = 1):
        self.U, self.M, self.d, self.reg, self.sweeps = n_users, n_items, factors, reg, sweeps
        # user vectors: column 0 is a fixed 1 so that column 0 of the item vectors is the item bias
        self.P = np.zeros((n_users, factors + 1))
        self.P[:, 0] = 1.0
        key = stream_key(seed, STREAMS["als"])
        u = uniforms(key, np.arange(n_items * factors)).reshape(n_items, factors)
        self.Q = np.zeros((n_items, factors + 1))
        self.Q[:, 1:] = 0.01 * (2.0 * u - 1.0)      # tiny start, only to break the symmetry
        # what the system has seen: times shown and times clicked, per (user, item)
        self.N = np.zeros((n_users, n_items))
        self.S = np.zeros((n_users, n_items))

    def observe(self, users: np.ndarray, items: np.ndarray, clicked: np.ndarray) -> None:
        np.add.at(self.N, (users, items), 1.0)
        np.add.at(self.S, (users, items), clicked.astype(float))

    def scores(self) -> np.ndarray:
        """Predicted click rate for every (user, item)."""
        return self.P @ self.Q.T

    def fit(self) -> None:
        d, lam = self.d, self.reg
        for _ in range(self.sweeps):
            # users: solve (sum_i n_ui q_i q_i^T + lam I) p_u = sum_i (s_ui - n_ui b_i) q_i over the latent part
            Ql = self.Q[:, 1:]
            A = np.einsum("ui,id,ie->ude", self.N, Ql, Ql) + lam * np.eye(d)[None]
            b = (self.S - self.N * self.Q[:, 0][None, :]) @ Ql
            self.P[:, 1:] = np.linalg.solve(A, b[:, :, None])[:, :, 0]
            # items: solve (sum_u n_ui p_u p_u^T + lam I) q_i = sum_u s_ui p_u, bias included
            A = np.einsum("ui,ud,ue->ide", self.N, self.P, self.P) + lam * np.eye(d + 1)[None]
            b = self.S.T @ self.P
            self.Q = np.linalg.solve(A, b[:, :, None])[:, :, 0]