"""Deck & Agol 1411.0004v1 Eqs.10–15 plus only the free-e resonant term.

Geometric coordinate times. Leading-order resonant Zfree convention; the model
does not resolve the separate individual free eccentricities. No Engine imports.
"""

import math

TAU = 2 * math.pi
EARTHS_PER_STAR = 332946.0487 * 1.08


def laplace(alpha, harmonic, intervals=512):
    # Periodic trapezoid; differentiated integrand supplies D_alpha b.
    coefficient = 0.0
    derivative = 0.0
    for index in range(intervals):
        angle = TAU * index / intervals
        cosine = math.cos(angle)
        denominator = 1 - 2 * alpha * cosine + alpha * alpha
        factor = math.cos(harmonic * angle)
        coefficient += factor / math.sqrt(denominator)
        derivative += factor * alpha * (cosine - alpha) / denominator**1.5
    return 2 * coefficient / intervals, 2 * derivative / intervals


def coefficients(period_ratio, harmonics=8, intervals=512):
    alpha = period_ratio ** (-2 / 3)
    inner = []
    outer = []
    for harmonic in range(1, harmonics + 1):
        value, derivative = laplace(alpha, harmonic, intervals)
        beta = harmonic * (1 - 1 / period_ratio)
        kappa = harmonic * (period_ratio - 1)
        indirect_inner = alpha * (beta * beta + 2 * beta + 3) if harmonic == 1 else 0
        indirect_outer = (
            (kappa * kappa - 2 * kappa + 3) / alpha**2 if harmonic == 1 else 0
        )
        inner.append(
            -alpha
            * (
                harmonic * (beta * beta + 3) * value
                + 2 * beta * derivative
                - indirect_inner
            )
            / (beta * beta * (beta * beta - 1))
        )
        outer.append(
            (
                harmonic * (kappa * kappa + 3) * value
                + 2 * kappa * (derivative + value)
                - indirect_outer
            )
            / (kappa * kappa * (kappa * kappa - 1))
        )
    return inner, outer


class CoefficientGrid:
    def __init__(self):
        self.lower = 1.89
        self.step = 0.0001
        self.values = [
            coefficients(self.lower + index * self.step) for index in range(301)
        ]

    def at(self, ratio):
        coordinate = (ratio - self.lower) / self.step
        if not 0 <= coordinate < len(self.values) - 1:
            raise ValueError("Harmonic inference outside frozen ratio grid")
        index = int(coordinate)
        fraction = coordinate - index
        return [
            [
                a + fraction * (b - a)
                for a, b in zip(
                    self.values[index][planet], self.values[index + 1][planet]
                )
            ]
            for planet in range(2)
        ]


def linear_design(parameters, records, grid, chopping=True):
    (
        mass_inner,
        mass_outer,
        real_free,
        imaginary_free,
        center_inner,
        center_outer,
        period_inner,
        period_outer,
        _,
    ) = parameters
    ratio = period_outer / period_inner
    delta = ratio / 2 - 1
    inner_coefficients, outer_coefficients = grid.at(ratio)
    offsets = [center_inner - 10 * period_inner, center_outer - 5 * period_outer]
    periods = [period_inner, period_outer]
    design = []
    for row in records:
        planet = 0 if row["planet"] == "Kepler-289d" else 1
        linear_time = offsets[planet] + periods[planet] * row["planet_epoch"]
        lambda_inner = TAU * (linear_time - offsets[0]) / period_inner
        lambda_outer = TAU * (linear_time - offsets[1]) / period_outer
        synodic_angle = lambda_inner - lambda_outer
        slow_angle = 2 * lambda_outer - lambda_inner
        terms = inner_coefficients if planet == 0 else outer_coefficients
        selected = (
            range(1, len(terms) + 1) if chopping else ([2] if planet == 0 else [1])
        )
        timing = (
            periods[planet]
            / TAU
            / EARTHS_PER_STAR
            * sum(
                terms[harmonic - 1] * math.sin(harmonic * synodic_angle)
                for harmonic in selected
            )
        )
        free_factor = periods[planet] / math.pi / EARTHS_PER_STAR * 1.5 / delta**2
        multiplier = -free_factor / 2 ** (2 / 3) if planet == 0 else free_factor / 2
        design.append(
            (
                linear_time,
                1 - planet,
                timing,
                multiplier * math.sin(slow_angle),
                -multiplier * math.cos(slow_angle),
            )
        )
    return design


def predict(parameters, records, grid, chopping=True):
    return [
        baseline
        + parameters[mass_index]
        * (forcing + parameters[2] * real + parameters[3] * imaginary)
        for baseline, mass_index, forcing, real, imaginary in linear_design(
            parameters, records, grid, chopping
        )
    ]