from scipy.stats import norm
import math


def z_score_from_confidence(confidence: float) -> float:
    return round(norm.ppf(1 - (1 - confidence) / 2), 3)


def cochran(Z: float, p: float, E: float, N: int | None = None):
    Z = z_score_from_confidence(Z)

    n0 = ((Z ** 2) * p * (1 - p)) / (E ** 2)

    if N is None:
        return {
            "n0": n0,
            "n0_round": math.ceil(n0)
        }

    finite_corrections = n0 / (1 + ((n0 - 1) / N))

    return {
        "n0": n0,
        "n0_round": math.ceil(n0),
        "finite_corrections": finite_corrections,
        "finite_corrections_round": math.ceil(finite_corrections)
    }