"""Quantiles and independent-ensemble-mean mixing diagnostics."""

import math
import statistics


def quantile(values, fraction):
    ordered = sorted(values)
    position = (len(ordered) - 1) * fraction
    lower = int(position)
    return ordered[lower] + (position - lower) * (
        ordered[min(lower + 1, len(ordered) - 1)] - ordered[lower]
    )


def summarize(runs):
    names = [
        "middle_mass_earth",
        "outer_mass_earth",
        "real_Zfree",
        "imag_Zfree",
        "middle_center_bjd",
        "outer_center_bjd",
        "middle_period_days",
        "outer_period_days",
        "jitter_days",
    ]
    all_samples = [row for run in runs for row in run["samples"]]
    result = {}
    for column, name in enumerate(names):
        chains = [[row[column] for row in run["samples"]] for run in runs]
        # Independent ensembles are the independent chains. Within one ensemble,
        # average walkers at each retained iteration before convergence/ESS work;
        # interleaved correlated walkers are not independent time samples.
        averaged = [
            [
                statistics.mean(chain[start : start + 30])
                for start in range(0, len(chain), 30)
            ]
            for chain in chains
        ]
        means = [statistics.mean(chain) for chain in averaged]
        variances = [statistics.variance(chain) for chain in averaged]
        within = statistics.mean(variances)
        length = len(averaged[0])
        between = (
            length * statistics.variance(means) if len(chains) > 1 else float("nan")
        )
        ratio = math.sqrt(((length - 1) / length * within + between / length) / within)
        autocorrelations = []
        for lag in range(1, min(length // 2, 100)):
            correlation = statistics.mean(
                sum(
                    (chain[index] - mean) * (chain[index + lag] - mean)
                    for index in range(length - lag)
                )
                / ((length - lag) * variance)
                for chain, mean, variance in zip(averaged, means, variances)
            )
            autocorrelations.append(correlation)
        positive_sum = 0.0
        for index in range(0, len(autocorrelations) - 1, 2):
            pair = autocorrelations[index] + autocorrelations[index + 1]
            if pair < 0:
                break
            positive_sum += pair
        effective_means = len(runs) * length / max(1, 1 + 2 * positive_sum)
        values = [row[column] for row in all_samples]
        result[name] = {
            "q16": quantile(values, 0.16),
            "median": quantile(values, 0.5),
            "q84": quantile(values, 0.84),
            "ensemble_mean_rhat": ratio,
            "effective_ensemble_means": effective_means,
        }
    return result