import polars as pl


def stat_count(
    value_col: str,
    *,
    result_aliases: list[str] = ["n", "missing_n", "missing_pct"],
) -> list[pl.Expr]:
    null_count = pl.col(value_col).null_count()
    return [
        pl.len().alias(result_aliases[0]),
        null_count.alias(result_aliases[1]),
        (null_count / pl.len()).alias(result_aliases[2]),
    ]


def statistics_count_n(value_cont: str, result_alias: str = "n") -> list[pl.Expr]:
    return [stat_count(value_cont, result_aliases=[result_alias, "missing_n", "missing_pct"])[0].alias(result_alias)]


def statistics_count_missing_n(value_cont: str, result_alias: str = "missing_n") -> list[pl.Expr]:
    return [stat_count(value_cont, result_aliases=["n", result_alias, "missing_pct"])[1]]


def statistics_count_missing_pct(value_cont: str, result_alias: str = "missing_pct") -> list[pl.Expr]:
    return [stat_count(value_cont, result_aliases=["n", "missing_n", result_alias])[2].alias(result_alias)]


def stat_arithmetic(
    value_col: str,
    *,
    result_aliases: list[str] = ["mean", "st", "sem", "mean_95_ci_lower", "mean_95_ci_upper"],
) -> list[pl.Expr]:
    n = pl.col(value_col).is_finite().sum()
    mean = pl.col(value_col).mean()
    sem = pl.col(value_col).std() / n.sqrt()
    return [
        mean.alias(result_aliases[0]),
        pl.col(value_col).std().alias(result_aliases[1]),
        sem.alias(result_aliases[2]),
        (mean - 1.96 * sem).alias(result_aliases[3]),
        (mean + 1.96 * sem).alias(result_aliases[4]),
    ]


def statistics_mean(value_cont: str, result_alias: str = "mean") -> list[pl.Expr]:
    return [
        stat_arithmetic(value_cont, result_aliases=[result_alias, "st", "sem", "mean_95_ci_lower", "mean_95_ci_upper"])[
            0
        ].alias(result_alias)
    ]


def statistics_st(value_cont: str, result_alias: str = "st") -> list[pl.Expr]:
    return [
        stat_arithmetic(
            value_cont, result_aliases=["mean", result_alias, "sem", "mean_95_ci_lower", "mean_95_ci_upper"]
        )[1].alias(result_alias)
    ]


def statistics_sem(value_cont: str, result_alias: str = "sem") -> list[pl.Expr]:
    return [
        stat_arithmetic(
            value_cont, result_aliases=["mean", "st", result_alias, "mean_95_ci_lower", "mean_95_ci_upper"]
        )[2].alias(result_alias)
    ]


def statistics_mean_95_ci_lower(value_cont: str, result_alias: str = "mean_95_ci_lower") -> list[pl.Expr]:
    return [
        stat_arithmetic(value_cont, result_aliases=["mean", "st", "sem", result_alias, "mean_95_ci_upper"])[3].alias(
            result_alias
        )
    ]


def statistics_mean_95_ci_upper(value_cont: str, result_alias: str = "mean_95_ci_upper") -> list[pl.Expr]:
    return [
        stat_arithmetic(value_cont, result_aliases=["mean", "st", "sem", "mean_95_ci_lower", result_alias])[4].alias(
            result_alias
        )
    ]


def stat_geometric(
    value_col: str,
    *,
    result_aliases: list[str] = ["geom_mean", "geom_mean_95_ci_lower", "geom_mean_95_ci_upper"],
) -> list[pl.Expr]:
    n = pl.col(value_col).is_finite().sum()
    log_mean = pl.col(value_col).log().mean()
    se = pl.col(value_col).log().std() / n.sqrt()
    return [
        log_mean.exp().alias(result_aliases[0]),
        (log_mean - 1.96 * se).exp().alias(result_aliases[1]),
        (log_mean + 1.96 * se).exp().alias(result_aliases[2]),
    ]


def statistics_geom_mean(value_cont: str, result_alias: str = "geom_mean") -> list[pl.Expr]:
    return [
        stat_geometric(value_cont, result_aliases=[result_alias, "geom_mean_95_ci_lower", "geom_mean_95_ci_upper"])[
            0
        ].alias(result_alias)
    ]


def statistics_geom_mean_95_ci_lower(value_cont: str, result_alias: str = "geom_mean_95_ci_lower") -> list[pl.Expr]:
    return [
        stat_geometric(value_cont, result_aliases=["geom_mean", result_alias, "geom_mean_95_ci_upper"])[1].alias(
            result_alias
        )
    ]


def statistics_geom_mean_95_ci_upper(value_cont: str, result_alias: str = "geom_mean_95_ci_upper") -> list[pl.Expr]:
    return [
        stat_geometric(value_cont, result_aliases=["geom_mean", "geom_mean_95_ci_lower", result_alias])[2].alias(
            result_alias
        )
    ]


def _percentile_ci_lower(
    value_col: str,
    q: float,
    *,
    result_aliases: list[str] = ["p", "ci_lower"],
) -> pl.Expr:
    """Calculate the lower confidence interval for a given percentile.

    Formula (Conover, 1999):
        j = nq - 1.96 root(nq(1-q))

    for ci_lower / n in [0,1]
    """
    n = pl.col(value_col).is_finite().sum()
    se = (n * q * (1 - q)).sqrt()
    ci_lower = (n * q - 1.96 * se).ceil()

    return (
        pl.col(value_col)
        .quantile((ci_lower / n).clip(0, 1), interpolation="nearest")
        .alias(f"{result_aliases[0]}{int(q * 100)}_{result_aliases[1]}")
    )


def _percentile_ci_upper(
    value_col: str,
    q: float,
    *,
    result_aliases: list[str] = ["p", "ci_upper"],
) -> pl.Expr:
    """Calculate the upper confidence interval for a given percentile.

    Formula (Conover, 1999):
        k = nq + 1.96 root(nq(1-q))
    """
    n = pl.col(value_col).is_finite().sum()
    se = (n * q * (1 - q)).sqrt()
    ci_upper = (n * q + 1.96 * se).ceil()

    return (
        pl.col(value_col)
        .quantile((ci_upper / n).clip(0, 1), interpolation="nearest")
        .alias(f"{result_aliases[0]}{int(q * 100)}_{result_aliases[1]}")
    )


def stat_percentiles(
    value_col: str,
    quants: list[float] = [0.05, 0.1, 0.25, 0.5, 0.75, 0.9, 0.95],
    *,
    result_aliases: list[str] = ["p", "ci_lower", "ci_upper"],
) -> list[pl.Expr]:
    quantile_exprs = [pl.col(value_col).quantile(q).alias(f"{result_aliases[0]}{int(q * 100)}") for q in quants]
    quantile_ci_lower_exprs = [
        _percentile_ci_lower(value_col, q, result_aliases=[result_aliases[0], result_aliases[1]]) for q in quants
    ]
    quantile_ci_upper_exprs = [
        _percentile_ci_upper(value_col, q, result_aliases=[result_aliases[0], result_aliases[2]]) for q in quants
    ]
    return quantile_exprs + quantile_ci_lower_exprs + quantile_ci_upper_exprs


def statistics_percentiles_p5(
    value_cont: str, quants: list[float] = [0.05], result_aliases: str = "p5"
) -> list[pl.Expr]:
    return [
        stat_percentiles(value_cont, quants=quants, result_aliases=["p", "ci_lower", "ci_upper"])[0].alias(
            result_aliases
        )
    ]


def statistics_percentiles_p10(
    value_cont: str, quants: list[float] = [0.1], result_aliases: str = "p10"
) -> list[pl.Expr]:
    return [
        stat_percentiles(value_cont, quants=quants, result_aliases=["p", "ci_lower", "ci_upper"])[0].alias(
            result_aliases
        )
    ]


def statistics_percentiles_p25(
    value_cont: str, quants: list[float] = [0.25], result_aliases: str = "p25"
) -> list[pl.Expr]:
    return [
        stat_percentiles(value_cont, quants=quants, result_aliases=["p", "ci_lower", "ci_upper"])[0].alias(
            result_aliases
        )
    ]


def statistics_percentiles_p50(
    value_cont: str, quants: list[float] = [0.5], result_aliases: str = "p50"
) -> list[pl.Expr]:
    return [
        stat_percentiles(value_cont, quants=quants, result_aliases=["p", "ci_lower", "ci_upper"])[0].alias(
            result_aliases
        )
    ]


def statistics_percentiles_p75(
    value_cont: str, quants: list[float] = [0.75], result_aliases: str = "p75"
) -> list[pl.Expr]:
    return [
        stat_percentiles(value_cont, quants=quants, result_aliases=["p", "ci_lower", "ci_upper"])[0].alias(
            result_aliases
        )
    ]


def statistics_percentiles_p90(
    value_cont: str, quants: list[float] = [0.9], result_aliases: str = "p90"
) -> list[pl.Expr]:
    return [
        stat_percentiles(value_cont, quants=quants, result_aliases=["p", "ci_lower", "ci_upper"])[0].alias(
            result_aliases
        )
    ]


def statistics_percentiles_p95(
    value_cont: str, quants: list[float] = [0.95], result_aliases: str = "p95"
) -> list[pl.Expr]:
    return [
        stat_percentiles(value_cont, quants=quants, result_aliases=["p", "ci_lower", "ci_upper"])[0].alias(
            result_aliases
        )
    ]


def statistics_percentiles_p5_ci_lower(
    value_cont: str, quants: list[float] = [0.05], result_aliases: str = "p5_ci_lower"
) -> list[pl.Expr]:
    return [
        stat_percentiles(value_cont, quants=quants, result_aliases=["p", "ci_lower", "ci_upper"])[1].alias(
            result_aliases
        )
    ]


def statistics_percentiles_p10_ci_lower(
    value_cont: str, quants: list[float] = [0.1], result_aliases: str = "p10_ci_lower"
) -> list[pl.Expr]:
    return [
        stat_percentiles(value_cont, quants=quants, result_aliases=["p", "ci_lower", "ci_upper"])[1].alias(
            result_aliases
        )
    ]


def statistics_percentiles_p25_ci_lower(
    value_cont: str, quants: list[float] = [0.25], result_aliases: str = "p25_ci_lower"
) -> list[pl.Expr]:
    return [
        stat_percentiles(value_cont, quants=quants, result_aliases=["p", "ci_lower", "ci_upper"])[1].alias(
            result_aliases
        )
    ]


def statistics_percentiles_p50_ci_lower(
    value_cont: str, quants: list[float] = [0.5], result_aliases: str = "p50_ci_lower"
) -> list[pl.Expr]:
    return [
        stat_percentiles(value_cont, quants=quants, result_aliases=["p", "ci_lower", "ci_upper"])[1].alias(
            result_aliases
        )
    ]


def statistics_percentiles_p75_ci_lower(
    value_cont: str, quants: list[float] = [0.75], result_aliases: str = "p75_ci_lower"
) -> list[pl.Expr]:
    return [
        stat_percentiles(value_cont, quants=quants, result_aliases=["p", "ci_lower", "ci_upper"])[1].alias(
            result_aliases
        )
    ]


def statistics_percentiles_p90_ci_lower(
    value_cont: str, quants: list[float] = [0.9], result_aliases: str = "p90_ci_lower"
) -> list[pl.Expr]:
    return [
        stat_percentiles(value_cont, quants=quants, result_aliases=["p", "ci_lower", "ci_upper"])[1].alias(
            result_aliases
        )
    ]


def statistics_percentiles_p95_ci_lower(
    value_cont: str, quants: list[float] = [0.95], result_aliases: str = "p95_ci_lower"
) -> list[pl.Expr]:
    return [
        stat_percentiles(value_cont, quants=quants, result_aliases=["p", "ci_lower", "ci_upper"])[1].alias(
            result_aliases
        )
    ]


def statistics_percentiles_p5_ci_upper(
    value_cont: str, quants: list[float] = [0.05], result_aliases: str = "p5_ci_upper"
) -> list[pl.Expr]:
    return [
        stat_percentiles(value_cont, quants=quants, result_aliases=["p", "ci_lower", "ci_upper"])[2].alias(
            result_aliases
        )
    ]


def statistics_percentiles_p10_ci_upper(
    value_cont: str, quants: list[float] = [0.1], result_aliases: str = "p10_ci_upper"
) -> list[pl.Expr]:
    return [
        stat_percentiles(value_cont, quants=quants, result_aliases=["p", "ci_lower", "ci_upper"])[2].alias(
            result_aliases
        )
    ]


def statistics_percentiles_p25_ci_upper(
    value_cont: str, quants: list[float] = [0.25], result_aliases: str = "p25_ci_upper"
) -> list[pl.Expr]:
    return [
        stat_percentiles(value_cont, quants=quants, result_aliases=["p", "ci_lower", "ci_upper"])[2].alias(
            result_aliases
        )
    ]


def statistics_percentiles_p50_ci_upper(
    value_cont: str, quants: list[float] = [0.5], result_aliases: str = "p50_ci_upper"
) -> list[pl.Expr]:
    return [
        stat_percentiles(value_cont, quants=quants, result_aliases=["p", "ci_lower", "ci_upper"])[2].alias(
            result_aliases
        )
    ]


def statistics_percentiles_p75_ci_upper(
    value_cont: str, quants: list[float] = [0.75], result_aliases: str = "p75_ci_upper"
) -> list[pl.Expr]:
    return [
        stat_percentiles(value_cont, quants=quants, result_aliases=["p", "ci_lower", "ci_upper"])[2].alias(
            result_aliases
        )
    ]


def statistics_percentiles_p90_ci_upper(
    value_cont: str, quants: list[float] = [0.9], result_aliases: str = "p90_ci_upper"
) -> list[pl.Expr]:
    return [
        stat_percentiles(value_cont, quants=quants, result_aliases=["p", "ci_lower", "ci_upper"])[2].alias(
            result_aliases
        )
    ]


def statistics_percentiles_p95_ci_upper(
    value_cont: str, quants: list[float] = [0.95], result_aliases: str = "p95_ci_upper"
) -> list[pl.Expr]:
    return [
        stat_percentiles(value_cont, quants=quants, result_aliases=["p", "ci_lower", "ci_upper"])[2].alias(
            result_aliases
        )
    ]
