"""The simulated world: a catalogue of items with topics and names, and users with tastes.

Everything is drawn from the "world" stream of rng.py in a fixed order, so the browser
version builds exactly the same catalogue and the same users from the same seed.
"""
from __future__ import annotations

from dataclasses import dataclass, field

import numpy as np

from .rng import Seq

# Eight topics. Each item's name is "<first> <joiner> <second>" from its topic's two word lists,
# so there are 64 possible names per topic and no two items in a topic share one.
TOPICS = [
    {"name": "Travel", "join": "of",
     "a": ["Slow trains", "Night markets", "Hidden beaches", "Old ferries", "Street food", "Mountain huts", "Back roads", "Winter light"],
     "b": ["Hokkaido", "Lisbon", "Patagonia", "the Hebrides", "Oaxaca", "Hanoi", "the Karoo", "Tbilisi"]},
    {"name": "Food", "join": "the",
     "a": ["Ramen", "Flatbread", "Dumplings", "Risotto", "Chilli", "Pho", "Focaccia", "Lentils"],
     "b": ["lazy way", "slow way", "Sunday way", "weeknight way", "one-pan way", "campfire way", "budget way", "proper way"]},
    {"name": "Science", "join": "",
     "a": ["Why glaciers", "Why bees", "Why stars", "Why cells", "Why magnets", "Why tides", "Why crystals", "Why volcanoes"],
     "b": ["hum", "remember", "flicker", "drift", "fail", "sing", "lie", "wait"]},
    {"name": "Music", "join": "from",
     "a": ["Lost demos", "B-sides", "Basslines", "Choirs", "Synths", "Drum breaks", "Live tapes", "Duets"],
     "b": ["1973", "Lagos", "a bedroom", "Detroit", "a cathedral", "late nights", "the Baltic", "Kingston"]},
    {"name": "Sport", "join": "of the",
     "a": ["Final laps", "Last-minute goals", "Comebacks", "Training diaries", "Great saves", "Long breakaways", "Tactics", "Underdogs"],
     "b": ["Tour", "World Cup", "marathon", "Six Nations", "Olympics", "derby", "Giro", "Ashes"]},
    {"name": "Film", "join": "in",
     "a": ["Long takes", "Jump cuts", "Last lines", "Opening shots", "Silent scenes", "Title cards", "Train scenes", "Final frames"],
     "b": ["noir", "westerns", "Ghibli", "the French new wave", "80s horror", "musicals", "heist movies", "silent cinema"]},
    {"name": "Garden", "join": "for",
     "a": ["Tomatoes", "Roses", "Moss", "Dahlias", "Herbs", "Peonies", "Bulbs", "Wildflowers"],
     "b": ["a north-facing wall", "a tiny balcony", "clay soil", "dry summers", "late frosts", "a shaded yard", "a windowsill", "salty air"]},
    {"name": "Money", "join": "",
     "a": ["Index funds", "Rent or buy", "Compound interest", "Bond yields", "Budgets", "Pensions", "Inflation", "Tax brackets"],
     "b": ["in five minutes", "with one chart", "for freelancers", "without the jargon", "with a spreadsheet", "for students", "in plain words", "with real numbers"]},
]


def title(topic: int, combo: int) -> str:
    t = TOPICS[topic]
    a, b = t["a"][combo // 8], t["b"][combo % 8]
    return a + " " + (t["join"] + " " if t["join"] else "") + b


@dataclass
class World:
    seed: int
    n_items: int
    n_users: int
    k: int
    topic: np.ndarray          # (M,) topic of each item
    quality: np.ndarray        # (M,) in [0.35, 1]
    style: np.ndarray          # (M,) in [-1, 1]
    titles: list
    taste: np.ndarray          # (U, T) each user's liking for each topic, max 1
    pref: np.ndarray           # (U,) style preference in [-1, 1]
    top_topic: np.ndarray      # (U,)
    new_item: int = -1         # index of the held-back item, or -1
    affinity: np.ndarray = field(default=None, repr=False)   # (U, M) true click probability at the top slot, first showing

    @property
    def n_topics(self) -> int:
        return len(TOPICS)


def make_world(seed: int = 1, n_items: int = 300, n_users: int = 400, k: int = 10,
               new_item: bool = False, new_item_topic: int = 3, new_item_quality: float = 0.9,
               visitor_ratings: dict | None = None) -> World:
    """Build the catalogue and the users. With new_item=True the last item is held back and
    released during the run (see sim.run); with visitor_ratings ({item: True/False}) one extra
    user is appended whose taste is guessed from those ratings."""
    T = len(TOPICS)
    rng = Seq(seed, 0)
    M = n_items
    topic = np.arange(M) % T
    quality = np.zeros(M)
    style = np.zeros(M)
    titles = []
    used = [set() for _ in range(T)]
    for i in range(M):
        quality[i] = 0.35 + 0.65 * rng.uniform()
        style[i] = 2.0 * rng.uniform() - 1.0
        while True:
            c = int(np.floor(rng.uniform() * 64))
            if c not in used[topic[i]]:
                used[topic[i]].add(c)
                break
        titles.append(title(int(topic[i]), c))

    U = n_users + (1 if visitor_ratings is not None else 0)
    taste = np.zeros((U, T))
    pref = np.zeros(U)
    for u in range(n_users):
        t1 = int(np.floor(rng.uniform() * T))
        t2 = int(np.floor(rng.uniform() * (T - 1)))
        if t2 >= t1:
            t2 += 1
        t3 = int(np.floor(rng.uniform() * (T - 2)))
        if t3 >= min(t1, t2):
            t3 += 1
        if t3 >= max(t1, t2):
            t3 += 1
        w = np.array([0.15 * rng.uniform() for _ in range(T)])
        w[t1] += 1.0
        w[t2] += 0.55
        w[t3] += 0.3
        taste[u] = w / w.max()
        pref[u] = 2.0 * rng.uniform() - 1.0

    new_idx = -1
    if new_item:
        new_idx = M - 1
        topic[new_idx] = new_item_topic
        quality[new_idx] = new_item_quality
        style[new_idx] = 0.0
        titles[new_idx] = title(new_item_topic, 63) if title(new_item_topic, 63) not in titles else "New: " + title(new_item_topic, 63)

    if visitor_ratings is not None:
        ups = np.zeros(T)
        n = np.zeros(T)
        for i, liked in visitor_ratings.items():
            n[topic[int(i)]] += 1
            ups[topic[int(i)]] += 1.0 if liked else 0.0
        w = (ups + 0.25) / (n + 0.5)
        taste[U - 1] = w / w.max()
        pref[U - 1] = 0.0

    top_topic = taste.argmax(axis=1)
    affinity = taste[:, topic] * quality[None, :] * (1.0 - 0.15 * np.abs(style[None, :] - pref[:, None]))
    return World(seed, M, U, k, topic, quality, style, titles, taste, pref, top_topic, new_idx, affinity)


def rating_items(world: World, count: int = 12) -> list:
    """Items for the visitor to rate: for topic t (cycling), its highest-quality item not yet used.
    The last item is never offered, because it is the one held back when a run releases a new item."""
    T = world.n_topics
    out = []
    for j in range(count):
        t = j % T
        cands = [i for i in range(world.n_items - 1) if world.topic[i] == t and i not in out]
        cands.sort(key=lambda i: (-world.quality[i], i))
        out.append(cands[0])
    return out


def oracle(world: World, active: np.ndarray) -> np.ndarray:
    """Mean true affinity of each user's k best active items: the best a feed could do."""
    a = np.where(active[None, :], world.affinity, -1.0)
    top = -np.partition(-a, world.k - 1, axis=1)[:, :world.k]
    return top.mean(axis=1)