"""AC18 independent science: differentiable latent activity covariance.

Times are days. RV is m/s; activity is log RHK divided by 0.01 dex after
centering. This authored implementation follows Rajpaul1506.07304v1 Eq28 and
AppendixA.2; no external kernel implementation is copied. MIT, Sandbox Physics.
"""
import numpy as np
from scipy.linalg import cho_factor, cho_solve


def latent_covariance(time_first, time_second, evolution_days, recurrence_days, periodic_scale):
    """Return k, derivative in FIRST time, and mixed derivative FIRST/SECOND."""
    lag = np.subtract.outer(time_first, time_second)
    angle = 2 * np.pi * lag / recurrence_days
    covariance = np.exp(-lag ** 2 / (2 * evolution_days ** 2)
                        - np.sin(angle / 2) ** 2 / (2 * periodic_scale ** 2))
    logarithmic_derivative = (-lag / evolution_days ** 2
                              - np.pi * np.sin(angle)
                              / (2 * recurrence_days * periodic_scale ** 2))
    logarithmic_second = (-1 / evolution_days ** 2
                          - np.pi ** 2 * np.cos(angle)
                          / (recurrence_days ** 2 * periodic_scale ** 2))
    first_derivative = covariance * logarithmic_derivative
    mixed_derivative = -covariance * (logarithmic_second + logarithmic_derivative ** 2)
    return covariance, first_derivative, mixed_derivative


def activity_covariance(time_first, time_second, parameters):
    """RV=a*G+b*Gdot; normalized calcium=A*G. Positive A fixes latent gauge."""
    covariance, derivative_first, derivative_mixed = latent_covariance(
        time_first, time_second, parameters["evolution_days"],
        parameters["rotation_days"], parameters["periodic_scale"])
    coupling = parameters["rv_coupling_m_s"]
    derivative_coupling = parameters["rv_derivative_m_s_day"]
    calcium_amplitude = parameters["calcium_amplitude_scaled"]
    rv_rv = coupling ** 2 * covariance + derivative_coupling ** 2 * derivative_mixed
    rv_calcium = calcium_amplitude * (coupling * covariance + derivative_coupling * derivative_first)
    calcium_rv = calcium_amplitude * (coupling * covariance - derivative_coupling * derivative_first)
    calcium_calcium = calcium_amplitude ** 2 * covariance
    return np.block([[rv_rv, rv_calcium], [calcium_rv, calcium_calcium]])


def linear_gaussian_posterior(values, noise_covariance, design, prior_scales):
    """Integrate the complete proper zero-mean Gaussian coefficient prior.

    The returned evidence is conditional on covariance hyperparameters. It is
    not the peak likelihood; covariance retains all coefficient cross terms.
    """
    prior_variances = np.asarray(prior_scales) ** 2
    scaled_design = design * prior_variances
    marginal_covariance = noise_covariance + np.einsum("ik,jk->ij", scaled_design, design)
    factor = cho_factor(marginal_covariance, lower=True, check_finite=False)
    inverse_values = cho_solve(factor, values, check_finite=False)
    inverse_design = cho_solve(factor, scaled_design, check_finite=False)
    mean = np.einsum("ij,i->j", scaled_design, inverse_values)
    covariance = np.diag(prior_variances) - np.einsum("ik,ij->kj", scaled_design, inverse_design)
    covariance = (covariance + covariance.T) / 2
    log_density = -.5 * (np.sum(values * inverse_values)
                        + 2 * np.log(np.diag(factor[0])).sum()
                        + len(values) * np.log(2 * np.pi))
    return float(log_density), mean, covariance