"""Proper Gaussian marginal evidence for a saved stationary harmonic record.

This is a declared observation model, not a periodogram FAP or a spot-truth readout.
"""
import numpy as np


class HarmonicLikelihood:
    def __init__(self, time_days, observed_flux, sigma_flux, exposure_days=None, harmonics=4,
                 baseline_sigma=.1, harmonic_sigma=.05, trend_sigma=.03):
        self.times = np.asarray(time_days, dtype=float)
        self.values = np.asarray(observed_flux, dtype=float) - 1.
        self.sigma = np.broadcast_to(np.asarray(sigma_flux, dtype=float), self.times.shape).copy()
        self.exposures = np.zeros_like(self.times) if exposure_days is None else np.broadcast_to(exposure_days, self.times.shape).copy()
        assert len(self.times) >= 12 and np.ptp(self.times) > 0 and np.all(self.sigma > 0)
        assert np.all(np.isfinite(self.values)) and np.all(self.exposures >= 0)
        self.centered_times = self.times - np.mean(self.times)
        self.scaled_time = self.centered_times / np.ptp(self.times)
        self.harmonics = harmonics
        self.prior_variances = np.array([baseline_sigma ** 2, trend_sigma ** 2] + [value for k in range(1, harmonics + 1) for value in [(harmonic_sigma / k) ** 2] * 2])
        self.weights = self.sigma ** -2
        self.noise_logdet = np.sum(np.log(self.sigma ** 2))
        self.prior_logdet = np.sum(np.log(self.prior_variances))
        self.raw_quadratic = np.sum(self.values ** 2 * self.weights)
        self.evaluations = 0

    def design(self, frequencies):
        frequencies = np.atleast_1d(frequencies)
        shape = (len(frequencies), len(self.times))
        columns = [np.ones(shape), np.broadcast_to(self.scaled_time, shape)]
        for harmonic in range(1, self.harmonics + 1):
            phase = 2 * np.pi * harmonic * frequencies[:, None] * self.centered_times[None, :]
            response = np.sinc(harmonic * frequencies[:, None] * self.exposures[None, :])
            columns.extend([np.sin(phase) * response, np.cos(phase) * response])
        return np.stack(columns, axis=-1)

    def evaluate(self, frequencies, diagnostics=False):
        frequencies = np.atleast_1d(frequencies)
        rows = []
        for start in range(0, len(frequencies), 128):
            design = self.design(frequencies[start:start + 128])
            normal = np.einsum('fni,n,fnj->fij', design, self.weights, design)
            normal += np.diag(1 / self.prior_variances)[None, :, :]
            right = np.einsum('fni,n,n->fi', design, self.weights, self.values)
            coefficients = np.linalg.solve(normal, right[..., None])[..., 0]
            signs, logdet = np.linalg.slogdet(normal)
            assert np.all(signs == 1)
            log_evidence = -.5 * (self.raw_quadratic - np.einsum('fi,fi->f', right, coefficients) + self.noise_logdet + self.prior_logdet + logdet + len(self.times) * np.log(2 * np.pi))
            self.evaluations += len(design)
            if diagnostics:
                residual = self.values[None, :] - np.einsum('fni,fi->fn', design, coefficients)
                rows.extend([{'log_evidence': float(value), 'residual_chi_square': float(np.sum(errors ** 2 * self.weights)), 'coefficients': vector.tolist()} for value, errors, vector in zip(log_evidence, residual, coefficients)])
            else:
                rows.extend(log_evidence)
        return rows if diagnostics else np.array(rows)

    def constant_log_evidence(self):
        design = np.column_stack((np.ones(len(self.times)), self.scaled_time))
        variances = self.prior_variances[:2]
        normal = (design.T * self.weights) @ design + np.diag(1 / variances)
        right = (design.T * self.weights) @ self.values
        return float(-.5 * (self.raw_quadratic - right @ np.linalg.solve(normal, right) + self.noise_logdet + np.sum(np.log(variances)) + np.linalg.slogdet(normal)[1] + len(self.times) * np.log(2 * np.pi)))