"""Accounting of rounds, messages and transmitted numbers/bytes for the protocols (star topology)."""
from __future__ import annotations


def protocol_accounting(K: int, p: int, r: int = 1, q: int = 1, hessian_mode: str = "exact", bytes_per_scalar: int = 8,
                        send_local_hessian: bool = False, header_bytes: int = 0):
    """Numbers and bytes moved for: local pilot (phase 1) + q derivative phases.

    Phase 1 uplink per worker: p (corrected local estimate) [+ p(p+1)/2 if the local Hessian is shipped: stale/AEE].
    Each derivative phase: downlink broadcast of p to every worker, uplink of p (score) [+ p(p+1)/2 Hessian if exact].
    hessian_mode: 'exact' | 'stale' | 'single' (no Hessian uplink; machine 0 solves) | 'none' (q = 0).
    Returns dict with rounds, messages, uplink/downlink numbers and bytes, local fits per worker.
    """
    h = p * (p + 1) // 2
    up1 = p + (h if (send_local_hessian or hessian_mode == "stale") else 0)
    if q == 0:
        up_d, down_d = 0, 0
    elif hessian_mode == "exact" or hessian_mode == "quantized":
        up_d, down_d = q * (p + h), q * p
    elif hessian_mode in ("stale", "single"):
        up_d, down_d = q * p, q * p
    else:
        raise ValueError(hessian_mode)
    uplink_numbers = K * (up1 + up_d)
    downlink_numbers = K * down_d
    rounds = 1 + q                      # synchronized exchange phases (uplink of estimates, then q derivative phases)
    messages = K * (1 + 2 * q)          # K uploads + per phase: K broadcasts + K uploads
    local_fits = 2 ** r - 1             # splits (1,2,4,...): 1 + 2 + 4 + ...
    local_data_passes = r               # each level touches every local observation once
    return {"rounds": rounds, "messages": messages, "uplink_numbers": uplink_numbers, "downlink_numbers": downlink_numbers,
            "uplink_bytes": uplink_numbers * bytes_per_scalar + messages * header_bytes,
            "downlink_bytes": downlink_numbers * bytes_per_scalar,
            "local_fits_per_worker": local_fits, "local_data_passes": local_data_passes}


def sufficient_exponent(r: int, q: int) -> float:
    """alpha(r,q) = 1 - 1/(2^(q+1) r): K = o(N^alpha) is the sufficient regime of Theorem 2 (exact Newton)."""
    return 1.0 - 1.0 / (2 ** (q + 1) * r)


def sufficient_exponent_stale(r: int, q: int) -> float:
    """Sufficient exponent with a fixed (stale) averaged local Hessian: effective depth r + q (Theorem 4, case (b))."""
    if q == 0:
        return sufficient_exponent(r, 0)
    return 1.0 - 1.0 / (2.0 * (r + q))


def sufficient_exponent_single(r: int, q: int) -> float:
    """Sufficient exponent with a single machine's Hessian (CSL-type curvature): effective depth r + q/2."""
    if q == 0:
        return sufficient_exponent(r, 0)
    return 1.0 - 1.0 / (2.0 * r + q)