"""Proper joint posterior with DE, truncated Gibbs and jitter slice updates."""

import math
import random
from timing_data import ephemeris
from timing_priors import physical, log_probability
from timing_conditionals import conditional_update, jitter_slice_update
from posterior_diagnostics import summarize


def run_ensemble(
    records, grid, chopping, seed, iterations=12000, eccentricity_scale=0.03
):
    generator = random.Random(seed)
    inner = [row for row in records if row["planet"] == "Kepler-289d"]
    outer = [row for row in records if row["planet"] == "Kepler-289c"]
    epoch_inner, period_inner = ephemeris(inner)
    epoch_outer, period_outer = ephemeris(outer)
    center = [
        math.log(6),
        math.log(150),
        0.0,
        0.0,
        epoch_inner + 10 * period_inner,
        epoch_outer + 5 * period_outer,
        period_inner,
        period_outer,
        math.log(10 / 1440),
    ]
    spreads = [0.4, 0.35, 0.01, 0.01, 0.025, 0.003, 0.004, 0.0008, 0.4]
    walkers = []
    while len(walkers) < 30:
        trial = [
            value + generator.gauss(0, scale) for value, scale in zip(center, spreads)
        ]
        score = log_probability(trial, records, grid, chopping, eccentricity_scale)
        if math.isfinite(score):
            walkers.append([trial, score])
    samples = []
    accepted = 0
    best = max(walkers, key=lambda item: item[1])
    for iteration in range(iterations):
        for index in range(len(walkers)):
            choices = [other for other in range(len(walkers)) if other != index]
            first, second = generator.sample(choices, 2)
            block = [4, 5, 6, 7] if iteration % 2 == 0 else list(range(9))
            gamma = (
                1.0 if generator.random() < 0.1 else 2.38 / math.sqrt(2 * len(block))
            )
            proposal = [
                (
                    value + gamma * (left - right) + generator.gauss(0, scale * 1e-6)
                    if component in block
                    else value
                )
                for component, (value, left, right, scale) in enumerate(
                    zip(
                        walkers[index][0],
                        walkers[first][0],
                        walkers[second][0],
                        spreads,
                    )
                )
            ]
            score = log_probability(
                proposal, records, grid, chopping, eccentricity_scale
            )
            if math.log(generator.random()) < score - walkers[index][1]:
                walkers[index] = [proposal, score]
                accepted += 1
                if score > best[1]:
                    best = [proposal[:], score]
            if iteration % 2 == 0:
                updated = conditional_update(
                    walkers[index][0],
                    records,
                    grid,
                    chopping,
                    eccentricity_scale,
                    generator,
                )
                updated = jitter_slice_update(
                    updated, records, grid, chopping, generator
                )
                updated_score = log_probability(
                    updated, records, grid, chopping, eccentricity_scale
                )
                if math.isfinite(updated_score):
                    walkers[index] = [updated, updated_score]
                    if updated_score > best[1]:
                        best = [updated[:], updated_score]
        if iteration >= iterations // 3 and iteration % 20 == 0:
            samples.extend(physical(walker[0]) for walker in walkers)
    return {
        "seed": seed,
        "iterations": iterations,
        "warmup": iterations // 3,
        "walkers": 30,
        "sampler": "ac19-de-gibbs-slice-1: differential evolution, truncated conditional Gaussian Gibbs, log-jitter slice",
        "acceptance": accepted / (iterations * 30),
        "samples": samples,
        "best": physical(best[0]),
        "best_log_density": best[1],
    }