"""R script export for fitted priors."""

from __future__ import annotations

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


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

    if prior.family == DistributionFamily.NORMAL:
        return f"{name} <- rnorm(n, mean = {p['mu']}, sd = {p['sigma']})"

    if prior.family == DistributionFamily.TRUNCATED_NORMAL:
        return (
            f"# Requires: library(truncnorm)\n"
            f"{name} <- rtruncnorm(n, a = {p['a']}, b = {p['b']}, "
            f"mean = {p['mu']}, sd = {p['sigma']})"
        )

    if prior.family == DistributionFamily.GAMMA:
        return f"{name} <- rgamma(n, shape = {p['alpha']}, scale = {p['scale']})"

    if prior.family == DistributionFamily.LOGNORMAL:
        return f"{name} <- rlnorm(n, meanlog = {p['mu']}, sdlog = {p['sigma']})"

    if prior.family == DistributionFamily.BETA:
        return (
            f"{name}_01 <- rbeta(n, shape1 = {p['alpha']}, shape2 = {p['beta']})\n"
            f"{name} <- {name}_01 * ({p['upper']} - {p['lower']}) + {p['lower']}"
        )

    if prior.family == DistributionFamily.UNIFORM:
        return f"{name} <- runif(n, min = {p['lower']}, max = {p['upper']})"

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


def export_r(batch: BatchResult, n_samples: int = 10000) -> str:
    """Export batch results as an R script."""
    lines = [
        "# Prior distributions generated by Distribird",
        "# https://github.com/HUN-REN-AI1Science/Distribird",
        "",
        "library(truncnorm)",
        "",
        f"n <- {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_r(prior))
        lines.append("")

    return "\n".join(lines)


def export_single_r(result: PipelineResult) -> str:
    """Export a single parameter as R code."""
    return _prior_to_r(result.prior)
