package recommendation

import (
	"context"

	"github.com/jackc/pgx/v5/pgtype"

	"git.fabledsword.com/bvandeusen/minstrel/internal/db/dbq"
)

// contextAffinityShrinkage is the pseudo-count that pulls a low-play artist's
// cell-share toward the user's baseline, so an artist with one or two plays
// can't swing its affinity to ±1 on noise. At k plays the estimate sits
// halfway between the raw cell-share and the baseline.
const contextAffinityShrinkage = 5.0

// ContextAffinity is the read-side map of per-artist time-of-day/weekday
// affinity for the CURRENT context (#1531): artist_id → score in [-1, +1].
// Absent artists (no play history) score 0, so cold-start candidates stay
// neutral. The zero value is a valid all-neutral affinity.
type ContextAffinity struct {
	byArtist map[pgtype.UUID]float64
}

// Affinity returns the artist's current-context affinity, or 0 if unknown.
func (c ContextAffinity) Affinity(artistID pgtype.UUID) float64 {
	return c.byArtist[artistID]
}

// LoadContextAffinity computes each artist's affinity for the user's CURRENT
// daypart × weekday cell. For every artist with completed plays in the window
// it compares the share of that artist's plays that fall in the current cell
// against the user's overall baseline share, shrinking sparse artists toward
// the baseline. Returns an empty (all-neutral) affinity when the user has no
// plays.
func LoadContextAffinity(
	ctx context.Context, q *dbq.Queries, userID pgtype.UUID,
) (ContextAffinity, error) {
	rows, err := q.ListArtistContextPlayCountsForUser(ctx, userID)
	if err != nil {
		return ContextAffinity{}, err
	}
	var totalPlays, cellPlays int64
	for _, r := range rows {
		totalPlays += r.TotalPlays
		cellPlays += r.CellPlays
	}
	out := ContextAffinity{byArtist: make(map[pgtype.UUID]float64, len(rows))}
	if totalPlays == 0 {
		return out, nil
	}
	baseline := float64(cellPlays) / float64(totalPlays)
	for _, r := range rows {
		out.byArtist[r.ArtistID] = contextAffinity(
			float64(r.CellPlays), float64(r.TotalPlays), baseline, contextAffinityShrinkage)
	}
	return out, nil
}

// contextAffinity returns an artist's shrunk cell-share minus the user's
// baseline share, clamped to [-1, 1]. The shrinkage pseudo-count k pulls
// low-play artists toward the baseline (→ 0 affinity) so noise can't dominate;
// a heavily-played artist keeps close to its raw over/under-representation.
func contextAffinity(cellPlays, totalPlays, baseline, k float64) float64 {
	if totalPlays == 0 {
		return 0
	}
	shrunk := (cellPlays + baseline*k) / (totalPlays + k)
	return clampUnit(shrunk - baseline)
}