"""The true distribution p* over K knowledge items (blueprint 2.1, 2.7).

``p*`` is fixed and deliberately heavy-tailed: most mass on common ("head") items, a
long thin tail of rare items whose loss *is* model collapse. Items are partitioned into
R disjoint regions; regions are what grounding and specialisation target.

Region design: each region is an identical, independently-normalised block carrying
mass 1/R, so regions are symmetric and every region has its own head and tail. This
makes region-matched grounding (E3) well posed and, at R=1, reduces to a single global
Zipf identical to the reference used by the scientific-validation suite.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np


@dataclass(frozen=True)
class TrueDist:
    """The fixed ground truth for a lineage.

    Attributes:
        p_star (np.ndarray): Length-K probability vector (sums to 1).
        regions (np.ndarray): Length-K int array; region index of each item.
        tail_mask (np.ndarray): Length-K bool array; True for designated tail items.
    """

    p_star: np.ndarray
    regions: np.ndarray
    tail_mask: np.ndarray


def make_true_distribution(
    K: int,
    R: int,
    tail: str,
    tail_frac: float,
    zipf_s: float,
    seed: int,
    *,
    tail_threshold: float = 1e-3,
) -> TrueDist:
    """Construct the true distribution, region assignment, and tail mask.

    Args:
        K (int): Number of knowledge items. Must be divisible by R.
        R (int): Number of regions (disjoint contiguous blocks of size K/R).
        tail (str): Tail family, ``"zipf"`` or ``"twocomponent"``.
        tail_frac (float): For ``twocomponent``, the fraction of each region's items
            placed in the low-mass tail component. Unused for ``zipf``.
        zipf_s (float): Zipf exponent (larger -> heavier head, thinner tail).
        seed (int): Present for interface symmetry; the construction is deterministic,
            so the seed only matters if a randomised item ordering is added later.
        tail_threshold (float): Items with ``p*_i < tail_threshold`` are the tail
            (blueprint 2.3). Keyword-only.

    Returns:
        TrueDist: p_star, regions, tail_mask.

    Raises:
        ValueError: If ``K`` is not divisible by ``R`` or ``tail`` is unknown.
    """
    if K % R != 0:
        raise ValueError(f"K={K} must be divisible by R={R} (contiguous equal regions).")
    per = K // R
    regions = np.repeat(np.arange(R), per)

    if tail == "zipf":
        block = 1.0 / np.arange(1, per + 1, dtype=float) ** zipf_s
        block = block / block.sum()                     # each region normalised to 1
    elif tail == "twocomponent":
        n_tail = max(1, int(round(tail_frac * per)))
        block = np.ones(per, dtype=float)
        # Tail items sit safely below threshold; heads carry the rest. Per-region mass
        # is 1/R after the global normalisation below.
        block[per - n_tail:] = tail_threshold * 0.5 * R
        block = block / block.sum()
    else:
        raise ValueError(f"unknown tail family {tail!r} (expected zipf|twocomponent)")

    p_star = np.tile(block, R) / R                      # R blocks, total mass 1
    p_star = p_star / p_star.sum()                      # guard float drift
    tail_mask = p_star < tail_threshold
    return TrueDist(p_star=p_star, regions=regions, tail_mask=tail_mask)


def uniform_init(K: int) -> np.ndarray:
    """The maximum-entropy initial distribution p_0 = 1/K (blueprint: the fresh base).

    The default lineage start. The neutral-decay law E[H_t]=H_0(1-1/n)^t holds from any
    start, so this choice fixes H_0 = 1 - 1/K without loss of generality.
    """
    return np.full(K, 1.0 / K)