"""
functions/cochran.py
--------------------
Pure calculation helpers for Cochran sample size formula.
No shared state — all inputs come from request, all outputs returned directly.
"""
import math
from scipy.stats import norm


def z_score_from_confidence(confidence: float) -> float:
    """Convert a confidence level (e.g. 0.95) to a Z-score."""
    return round(float(norm.ppf(1 - (1 - confidence) / 2)), 4)


def cochran_sample_size(Z: float, p: float, E: float, N: int | None = None) -> dict:
    """
    Cochran formula for sample size.

    Parameters
    ----------
    Z : float
        Z-score (e.g. 1.96 for 95 % confidence).
        If <= 1 it is treated as a confidence level and converted first.
    p : float
        Estimated proportion of the attribute (use 0.5 for maximum variance).
    E : float
        Margin of error (e.g. 0.05 for ±5 %).
    N : int | None
        Population size. When supplied, the finite-population correction is applied.

    Returns
    -------
    dict with n0, n0_rounded, and (when N given) finite_corrected, finite_corrected_rounded.
    """
    # Allow confidence level input
    if Z <= 1:
        Z = z_score_from_confidence(Z)

    n0 = (Z ** 2 * p * (1 - p)) / (E ** 2)
    result = {
        "z_used": round(Z, 4),
        "n0": round(n0, 4),
        "n0_rounded": math.ceil(n0),
    }

    if N is not None and N > 0:
        nf = n0 / (1 + (n0 - 1) / N)
        result["population_size"] = N
        result["finite_corrected"] = round(nf, 4)
        result["finite_corrected_rounded"] = math.ceil(nf)
        result["recommended_n"] = math.ceil(nf)
    else:
        result["recommended_n"] = math.ceil(n0)

    return result