"""Support-preserving conditional Gibbs and log-jitter slice updates."""

import math
import statistics
from chopping_model import predict, linear_design
from timing_priors import physical


def conditional_update(state, records, grid, chopping, eccentricity_scale, generator):
    """Exact Gaussian conditional updates for Z and the positive linear masses.

    This Gibbs move traverses curved mass/eccentricity ridges that an otherwise
    valid differential-evolution sampler can explore too slowly.
    """
    parameters = physical(state)
    design = linear_design(parameters, records, grid, chopping)
    inverse_prior = 1 / eccentricity_scale**2
    precision_aa, precision_bb, precision_ab = inverse_prior, inverse_prior, 0.0
    right_a, right_b = 0.0, 0.0
    for row, (baseline, index, forcing, real, imaginary) in zip(records, design):
        weight = 1 / (row["sigma_days"] ** 2 + parameters[8] ** 2)
        first = parameters[index] * real
        second = parameters[index] * imaginary
        residual = row["observed_reduced_bjd"] - baseline - parameters[index] * forcing
        precision_aa += weight * first * first
        precision_bb += weight * second * second
        precision_ab += weight * first * second
        right_a += weight * first * residual
        right_b += weight * second * residual
    determinant = precision_aa * precision_bb - precision_ab**2
    covariance_aa = precision_bb / determinant
    covariance_ab = -precision_ab / determinant
    covariance_bb = precision_aa / determinant
    mean_a = covariance_aa * right_a + covariance_ab * right_b
    mean_b = covariance_ab * right_a + covariance_bb * right_b
    root_a = math.sqrt(covariance_aa)
    root_b = math.sqrt(max(0, covariance_bb - covariance_ab**2 / covariance_aa))
    for attempt in range(1000):
        first = generator.gauss(0, 1)
        real = mean_a + root_a * first
        imaginary = (
            mean_b + covariance_ab / root_a * first + root_b * generator.gauss(0, 1)
        )
        if abs(real) < 0.15 and abs(imaginary) < 0.15:
            parameters[2:4] = [real, imaginary]
            break
    for mass_index, bounds in [(0, (0.05, 60)), (1, (5, 600))]:
        numerator, denominator = 0.0, 0.0
        for row, (baseline, index, forcing, real, imaginary) in zip(records, design):
            if index != mass_index:
                continue
            coefficient = forcing + parameters[2] * real + parameters[3] * imaginary
            weight = 1 / (row["sigma_days"] ** 2 + parameters[8] ** 2)
            numerator += weight * coefficient * (row["observed_reduced_bjd"] - baseline)
            denominator += weight * coefficient * coefficient
        mean = numerator / denominator
        deviation = 1 / math.sqrt(denominator)
        normal = statistics.NormalDist(mean, deviation)
        lower, upper = normal.cdf(bounds[0]), normal.cdf(bounds[1])
        if upper > lower and upper > 1e-15 and lower < 1 - 1e-15:
            probability = lower + generator.random() * (upper - lower)
            if 0 < probability < 1:
                parameters[mass_index] = normal.inv_cdf(probability)
    return [
        math.log(parameters[0]),
        math.log(parameters[1]),
        *parameters[2:8],
        math.log(parameters[8]),
    ]


def jitter_slice_update(state, records, grid, chopping, generator):
    """Univariate slice step for log jitter, including its physical-prior Jacobian."""
    parameters = physical(state)
    predictions = predict(parameters, records, grid, chopping)
    residuals = [
        (row["observed_reduced_bjd"] - model) ** 2
        for row, model in zip(records, predictions)
    ]
    variances = [row["sigma_days"] ** 2 for row in records]
    minimum, maximum = math.log(1e-5), math.log(0.2)

    def score(value):
        if not minimum < value < maximum:
            return -math.inf
        squared = math.exp(2 * value)
        return (
            -0.5
            * sum(
                residual / (variance + squared) + math.log(variance + squared)
                for residual, variance in zip(residuals, variances)
            )
            - 0.5 * squared / (30 / 1440) ** 2
            + value
        )

    level = score(state[8]) + math.log(1 - generator.random())
    left = state[8] - generator.random()
    right = left + 1
    while left > minimum and score(left) > level:
        left -= 1
    while right < maximum and score(right) > level:
        right += 1
    for attempt in range(200):
        trial = generator.uniform(left, right)
        if score(trial) >= level:
            state = state[:]
            state[8] = trial
            return state
        if trial < state[8]:
            left = trial
        else:
            right = trial
    raise RuntimeError("Log-jitter slice failed to find a finite point")