"""Внешние контрагенты в ОСВ (спека §5.3): сколько и кто, для каждой пары (компания, счёт).

Число `n` — отрицательное биномиальное (гамма-пуассоновская смесь с дисперсией ≈ 2 × матожидание), матожидание
`mean × w(счёт) / среднее_w`, так что среднее `n` по файлам счетов 58/60/62/66/67/76 ≈ `osv_counterparties_mean`.
Потолки: 58/60/62/76 — словарь без банков и компаний группы, 66/67 — банки. Недобор из-за потолка перераспределяется на
остальные файлы того же счёта (решение 9); если потолок всё же срезал выборку — `capped=True`.
"""

import math
import random
from dataclasses import dataclass, field

from synthgen.config import Config
from synthgen.dictionary import Org
from synthgen.group import Group
from synthgen.rng import rng_for

OSV_ACCOUNTS = ("01", "58", "60", "62", "66", "67", "76", "80", "90")
WEIGHTS = {"60": 1.6, "62": 1.6, "76": 1.0, "66": 0.1, "67": 0.05, "58": 0.1}
BANK_ACCOUNTS = ("66", "67")


@dataclass(frozen=True)
class CounterpartyPlan:
    by_file: dict[tuple[str, str], list[Org]] = field(default_factory=dict)  # (код компании, корень счёта) → контрагенты
    capped: bool = False  # потолок словаря срезал выборку хотя бы в одном файле


def _poisson(rng: random.Random, lam: float) -> int:
    if lam <= 0:
        return 0
    if lam > 30:  # нормальное приближение: точное для нашей задачи, а цикл Кнута на больших λ медленный
        return max(0, round(rng.gauss(lam, math.sqrt(lam))))
    limit, k, p = math.exp(-lam), 0, 1.0
    while True:
        p *= rng.random()
        if p <= limit:
            return k
        k += 1


def _draw_n(rng: random.Random, mu: float) -> int:
    """Отрицательное биномиальное с матожиданием `mu` и дисперсией 2·mu: Poisson(Gamma(mu, 1))."""
    if mu <= 0:
        return 0
    return _poisson(rng, rng.gammavariate(mu, 1.0))


def plan(cfg: Config, group: Group) -> CounterpartyPlan:
    codes = [m.code for m in group.members]
    avg_w = sum(WEIGHTS.values()) / len(WEIGHTS)
    capped = False
    by_file: dict[tuple[str, str], list[Org]] = {}
    for acct, w in WEIGHTS.items():
        pool = group.banks if acct in BANK_ACCOUNTS else group.externals
        cap = len(pool)
        mu = cfg.osv_counterparties_mean * w / avg_w
        want = [_draw_n(rng_for(cfg.seed, "cp-count", code, acct), mu) for code in codes]
        got = [min(n, cap) for n in want]
        shortfall = sum(want) - sum(got)
        if shortfall > 0:
            capped = True
            room = [i for i, n in enumerate(got) if n < cap]
            while shortfall > 0 and room:  # круговое доведение недобора до потолка файлов, где есть место
                nxt = []
                for i in room:
                    if shortfall == 0:
                        break
                    got[i] += 1
                    shortfall -= 1
                    if got[i] < cap:
                        nxt.append(i)
                room = nxt
        for code, n in zip(codes, got):
            by_file[(code, acct)] = rng_for(cfg.seed, "cp-set", code, acct).sample(pool, n) if n else []
    return CounterpartyPlan(by_file=by_file, capped=capped)