"""Python script export for fitted priors."""

from __future__ import annotations

from distribird.models import BatchResult, DistributionFamily, FittedPrior, PipelineResult


def _prior_to_python(prior: FittedPrior) -> str:
    """Generate Python/scipy code for sampling from a fitted prior."""
    p = prior.params
    name = prior.parameter_name.replace(" ", "_").replace("-", "_")

    if prior.family == DistributionFamily.NORMAL:
        return f"{name} = stats.norm(loc={p['mu']}, scale={p['sigma']})"

    if prior.family == DistributionFamily.TRUNCATED_NORMAL:
        a_std = (p["a"] - p["mu"]) / p["sigma"]
        b_std = (p["b"] - p["mu"]) / p["sigma"]
        return (
            f"{name} = stats.truncnorm("
            f"a={a_std:.6g}, b={b_std:.6g}, "
            f"loc={p['mu']}, scale={p['sigma']})"
        )

    if prior.family == DistributionFamily.GAMMA:
        return f"{name} = stats.gamma(a={p['alpha']}, scale={p['scale']})"

    if prior.family == DistributionFamily.LOGNORMAL:
        return f"{name} = stats.lognorm(s={p['sigma']}, scale=np.exp({p['mu']}))"

    if prior.family == DistributionFamily.BETA:
        return (
            f"# Beta on [{p['lower']}, {p['upper']}]\n"
            f"{name}_01 = stats.beta(a={p['alpha']}, b={p['beta']})\n"
            f"# To sample: {name}_01.rvs(n) * "
            f"({p['upper']} - {p['lower']}) + {p['lower']}"
        )

    if prior.family == DistributionFamily.UNIFORM:
        scale = p["upper"] - p["lower"]
        return f"{name} = stats.uniform(loc={p['lower']}, scale={scale})"

    return f"# Unsupported distribution: {prior.family.value}"


def export_python(batch: BatchResult, n_samples: int = 10000) -> str:
    """Export batch results as a Python script."""
    lines = [
        '"""Prior distributions generated by Distribird."""',
        "",
        "import numpy as np",
        "from scipy import stats",
        "",
        f"N_SAMPLES = {n_samples}",
        "",
    ]

    for result in batch.results:
        prior = result.prior
        lines.append(f"# {result.parameter.name}: {prior.confidence.value} confidence")
        lines.append(f"# {prior.reason}")
        if prior.evidence:
            for e in prior.evidence:
                doi_str = f" DOI: {e.doi}" if e.doi else ""
                lines.append(f"#   Source: {e.title} ({e.year}){doi_str}")
        lines.append(_prior_to_python(prior))
        lines.append("")

    return "\n".join(lines)


def export_single_python(result: PipelineResult) -> str:
    """Export a single parameter as Python code."""
    return _prior_to_python(result.prior)
