"""Builds the trader population for a run: randomized assignment of policies to trader slots,
stratified by wealth, information and latency (see docs/identification.md).

Assignment is a function of (seed, config) only, so matched pairs that differ in `llm_share`
share the same slot endowments, latencies and signal precisions: the LLM treatment *replaces* the
policy in a subset of slots, chosen by a dedicated RNG stream, holding everything else fixed."""
from __future__ import annotations

import numpy as np

from ..agents.base import AgentSpec
from .config import RunConfig


def _stratum(wealth_mult: float, precision: float, latency: float) -> str:
    w = "rich" if wealth_mult > 1.0 else "poor"
    i = "informed" if precision > 0 else "uninformed"
    l_ = "fast" if latency < 0.3 else "slow"
    return f"{w}/{i}/{l_}"


def build_population(cfg: RunConfig, seed_streams: dict[str, np.random.Generator]) -> list[AgentSpec]:
    pop = cfg.population
    N = pop.n_traders
    rng_slots = seed_streams["slots"]          # endowments, latencies, information
    rng_assign = seed_streams["assign"]        # policy assignment to slots
    rng_llm = seed_streams["llm_assign"]       # which slots become LLM-controlled

    # 1. slot characteristics (independent of policies)
    wealth = np.exp(rng_slots.standard_normal(N) * pop.endowment.wealth_dispersion)
    wealth /= wealth.mean()
    latency = rng_slots.uniform(pop.latency.min, pop.latency.max, N)
    informed_draw = rng_slots.random(N)

    # 2. classical policy for each slot (stratified: sort slots by wealth tercile and deal types evenly)
    names = list(pop.mix)
    counts = np.floor(np.array([pop.mix[n] for n in names]) * N).astype(int)
    while counts.sum() < N:
        counts[int(np.argmax(np.array([pop.mix[n] for n in names]) * N - counts))] += 1
    deck = np.concatenate([np.full(c, i) for i, c in enumerate(counts)])
    rng_assign.shuffle(deck)
    order = np.argsort(wealth)             # stratify by wealth: deal the shuffled deck down the wealth ranking
    policy_idx = np.empty(N, dtype=int)
    policy_idx[order] = deck

    # 3. LLM slots: replace a fraction of the (non market-maker) slots, stratified by policy
    n_llm = int(round(pop.llm_share * N))
    eligible = [i for i in range(N) if not (pop.llm_keep_market_makers and names[policy_idx[i]] == "market_maker"
                                             and pop.llm_replace != "market_maker")]
    if pop.llm_replace != "proportional":
        eligible = [i for i in eligible if names[policy_idx[i]] == pop.llm_replace]
    n_llm = min(n_llm, len(eligible))
    llm_slots: set[int] = set()
    if n_llm > 0:
        # stratified draw: proportional within each policy type, remainder at random
        by_type: dict[int, list[int]] = {}
        for i in eligible:
            by_type.setdefault(int(policy_idx[i]), []).append(i)
        quota = {k: int(np.floor(n_llm * len(v) / len(eligible))) for k, v in by_type.items()}
        for k, v in by_type.items():
            chosen = rng_llm.choice(v, size=min(quota[k], len(v)), replace=False)
            llm_slots.update(int(c) for c in chosen)
        rest = [i for i in eligible if i not in llm_slots]
        extra = n_llm - len(llm_slots)
        if extra > 0 and rest:
            llm_slots.update(int(c) for c in rng_llm.choice(rest, size=min(extra, len(rest)), replace=False))

    # 4. LLM provider / prompt / cohort assignment
    llm = pop.llm
    providers = list(llm.provider_mix)
    pw = np.array([llm.provider_mix[p] for p in providers], dtype=float)
    pw = pw / pw.sum()
    llm_list = sorted(llm_slots)
    prov_of: dict[int, str] = {}
    if llm_list:
        # deterministic proportional allocation then shuffle
        alloc = np.floor(pw * len(llm_list)).astype(int)
        while alloc.sum() < len(llm_list):
            alloc[int(np.argmax(pw * len(llm_list) - alloc))] += 1
        plist = np.concatenate([np.full(c, i) for i, c in enumerate(alloc)])
        rng_llm.shuffle(plist)
        prov_of = {s: providers[int(p)] for s, p in zip(llm_list, plist)}
    prompts = ["value", "momentum", "contrarian", "liquidity", "risk_averse", "macro", "quant", "generalist"]
    slots = max(1, min(llm.stagger_slots, llm.decision_interval))
    offsets = (rng_llm.integers(0, slots, size=max(len(llm_list), 1)) * llm.decision_interval) // slots
    cohort_of = {s: (k // max(llm.cohort_size, 1)) for k, s in enumerate(llm_list)} if llm.channel else {}

    specs: list[AgentSpec] = []
    for i in range(N):
        pname = names[policy_idx[i]]
        kw = dict(pop.policy_kwargs.get(pname, {}))
        prec = 0.0
        if pname == "fundamentalist" and informed_draw[i] < pop.informed_share:
            prec = pop.precision
        lev = pop.leverage_cap if pname == "leveraged" else 1.0
        cash = int(pop.endowment.cash_cents * wealth[i])
        shares = int(round(pop.endowment.shares * wealth[i]))
        positions = {a: shares for a in range(cfg.assets.n_assets)}
        if pname == "market_maker":
            cash *= 5
            positions = {a: shares * 5 for a in range(cfg.assets.n_assets)}
        if i in llm_slots:
            prov = prov_of[i]
            k = llm_list.index(i)
            prompt_id = prompts[k % len(prompts)] if llm.prompt_set == "diverse" else "generalist"
            llm_prec = prec if llm.informed_share is None else (llm.precision if informed_draw[i] < llm.informed_share else 0.0)
            specs.append(AgentSpec(
                agent_id=i, policy_name="llm", policy_kwargs={"replaced": pname}, cash=cash, positions=positions,
                latency=float(latency[i] if pop.latency.llm_latency is None else pop.latency.llm_latency),
                activity=1.0, decision_interval=llm.decision_interval,
                decision_offset=int(offsets[k]) if llm.stagger else 0, signal_precision=llm_prec,
                leverage_cap=lev, short_allowed=False, stratum=_stratum(wealth[i], llm_prec, latency[i]),
                provider=prov, model=llm.models.get(prov, ""), prompt_id=prompt_id, cohort=cohort_of.get(i, -1)))
        else:
            specs.append(AgentSpec(
                agent_id=i, policy_name=pname, policy_kwargs=kw, cash=cash, positions=positions,
                latency=float(latency[i]), activity=pop.activity.get(pname, 0.1), signal_precision=prec,
                leverage_cap=lev, short_allowed=False, stratum=_stratum(wealth[i], prec, latency[i])))
    return specs