"""
routes/cochran_routes.py
------------------------
Endpoints for Cochran sample size calculations.
Completely stateless — no shared data.
"""
from fastapi import APIRouter
from pydantic import BaseModel, Field

from functions.cochran import z_score_from_confidence, cochran_sample_size

router = APIRouter()


# ── Request models ────────────────────────────────────────────────────────────

class ConfidenceRequest(BaseModel):
    confidence: float = Field(..., gt=0, lt=1, description="Confidence level, e.g. 0.95")


class CochranRequest(BaseModel):
    Z: float = Field(..., description="Z-score OR confidence level (if <= 1 it's auto-converted)")
    p: float = Field(0.5, ge=0, le=1, description="Estimated proportion (0.5 = max variance)")
    E: float = Field(..., gt=0, description="Margin of error, e.g. 0.05 for ±5 %")
    N: int | None = Field(None, gt=0, description="Population size (optional, for finite correction)")


# ── Endpoints ─────────────────────────────────────────────────────────────────

@router.post("/z-score")
def get_z_score(payload: ConfidenceRequest):
    """Convert a confidence level to its corresponding Z-score."""
    return {"z_score": z_score_from_confidence(payload.confidence)}


@router.post("/sample-size")
def calculate_sample_size(payload: CochranRequest):
    """
    Calculate Cochran sample size.

    - If N is omitted, returns the infinite-population formula result.
    - If N is supplied, also returns the finite-population corrected value.
    """
    return cochran_sample_size(
        Z=payload.Z,
        p=payload.p,
        E=payload.E,
        N=payload.N,
    )