"""Independent normalized Gaussian and separated multimodal target benchmarks."""
import json
import numpy as np
from scipy.special import logsumexp
from rv_importance import fit_defensive_proposal,importance_diagnostics


def main():
    summaries=[]
    for multimodal in [False,True]:
        means=np.array([[0.,0.]]) if not multimodal else np.array([[-4.,0.],[4.,0.],[0.,5.]])
        masses=np.array([1.]) if not multimodal else np.array([.4,.35,.25])
        covariance=np.array([[.6,.22],[.22,.8]])
        inverse=np.linalg.inv(covariance)
        normalization=-np.log(2*np.pi)-.5*np.linalg.slogdet(covariance)[1]
        def target(point):
            residual=means-point
            return float(logsumexp(np.log(masses)+normalization-.5*np.einsum('ni,ij,nj->n',residual,inverse,residual)))
        expected_mean=np.einsum('n,ni->i',masses,means)
        centered=means-expected_mean
        expected_covariance=covariance+np.einsum('n,ni,nj->ij',masses,centered,centered)
        independent=[]
        for seed in [1803,1804,1805]:
            random=np.random.default_rng(seed)
            starts=np.array([[first,second] for first in [-4.,0.,4.] for second in [-4.,0.,4.]])
            proposal,fits=fit_defensive_proposal(target,2,random,probe_count=256,maximum_modes=6,exploration_starts=starts)
            points=proposal.sample(131072,random)
            residual_grid=means[None,:,:]-points[:,None,:]
            target_values=logsumexp(np.log(masses)[None,:]+normalization-.5*np.einsum('nki,ij,nkj->nk',residual_grid,inverse,residual_grid),axis=1)
            log_weights=target_values-proposal.log_density(points)
            weights=np.exp(log_weights-logsumexp(log_weights))
            mean=np.einsum('n,ni->i',weights,points)
            residual=points-mean
            inferred_covariance=np.einsum('n,ni,nj->ij',weights,residual,residual)
            diagnostic=importance_diagnostics(log_weights)
            mean_error=float(np.max(np.abs(mean-expected_mean)))
            covariance_error=float(np.max(np.abs(inferred_covariance-expected_covariance)))
            evidence_error=float(abs(np.exp(logsumexp(log_weights)-np.log(len(points)))-1))
            assert mean_error<.045 and covariance_error<.18 and evidence_error<.025
            assert diagnostic['effective_samples']>7000
            independent.append(dict(seed=seed,mean_error=mean_error,covariance_error=covariance_error,
                normalized_evidence_error=evidence_error,proposal_centers=len(proposal.centers),**diagnostic))
        summaries.append(dict(target='separated3Gaussianmixture' if multimodal else 'correlatedGaussian',runs=independent))
    print(json.dumps(summaries,indent=2))


if __name__=='__main__':main()