"""Defensive multi-center importance sampling in Gaussian-prior coordinates.

The proposal retains an explicit standard-normal prior component and normalized
Student-t components. Adaptation only chooses the proposal. Every final weight
uses target/proposal, so a proposal covariance is never a posterior covariance.
Independently authored public scientific analysis, MIT, Sandbox Physics AC18.
"""
import numpy as np
from scipy.special import gammaln,logsumexp
from scipy.optimize import minimize


def standard_normal_log_density(points):
    return -.5*(np.sum(points**2,axis=-1)+points.shape[-1]*np.log(2*np.pi))


class DefensiveProposal:
    def __init__(self,dimension,centers,covariances,prior_fraction=.15,degree=5.,component_masses=None):
        self.dimension=dimension;self.centers=np.asarray(centers)
        self.covariances=np.asarray(covariances);self.prior_fraction=prior_fraction;self.degree=degree
        self.component_masses=np.ones(len(centers))/len(centers) if component_masses is None else np.asarray(component_masses)/np.sum(component_masses)
        self.factors=[np.linalg.cholesky(covariance) for covariance in covariances]
        self.inverses=[np.linalg.inv(covariance) for covariance in covariances]
        self.log_normalizations=[gammaln((degree+dimension)/2)-gammaln(degree/2)
            -.5*(dimension*np.log(degree*np.pi)+np.linalg.slogdet(covariance)[1]) for covariance in covariances]

    def log_density(self,points):
        points=np.atleast_2d(points)
        components=[np.log(self.prior_fraction)+standard_normal_log_density(points)]
        for mass,center,inverse,normalization in zip(self.component_masses,self.centers,self.inverses,self.log_normalizations):
            residual=points-center
            squared=np.einsum('ni,ij,nj->n',residual,inverse,residual)
            components.append(np.log((1-self.prior_fraction)*mass)+normalization
                -(self.degree+self.dimension)/2*np.log1p(squared/self.degree))
        return logsumexp(np.array(components),axis=0)

    def sample(self,count,random):
        probabilities=[self.prior_fraction]+list((1-self.prior_fraction)*self.component_masses)
        choices=random.choice(len(probabilities),size=count,p=probabilities)
        points=random.normal(size=(count,self.dimension))
        for index,(center,factor) in enumerate(zip(self.centers,self.factors),1):
            selected=np.flatnonzero(choices==index)
            if not len(selected):continue
            normal=random.normal(size=(len(selected),self.dimension))
            student=normal/np.sqrt(random.chisquare(self.degree,size=len(selected))[:,None]/self.degree)
            points[selected]=center+np.einsum('ij,nj->ni',factor,student)
        return points

    def metadata(self):
        return dict(prior_fraction=self.prior_fraction,student_degrees=self.degree,
            centers=self.centers.tolist(),covariances=self.covariances.tolist(),
            mixture_component_weights=[self.prior_fraction]+list((1-self.prior_fraction)*self.component_masses))


def numerical_hessian(objective,point,step=.01):
    dimensions=len(point);value=objective(point);hessian=np.empty((dimensions,dimensions))
    basis=np.eye(dimensions)*step
    for first in range(dimensions):
        hessian[first,first]=(objective(point+basis[first])-2*value+objective(point-basis[first]))/step**2
        for second in range(first):
            entry=(objective(point+basis[first]+basis[second])-objective(point+basis[first]-basis[second])
                -objective(point-basis[first]+basis[second])+objective(point-basis[first]-basis[second]))/(4*step**2)
            hessian[first,second]=entry;hessian[second,first]=entry
    return hessian


def fit_defensive_proposal(log_target,dimension,random,probe_count=512,maximum_modes=8,exploration_starts=None):
    probes=random.normal(size=(probe_count,dimension))
    log_targets=np.array([log_target(probe) for probe in probes])
    starts=[]
    for index in np.argsort(log_targets)[::-1]:
        point=probes[index]
        if not starts or min(np.linalg.norm(point-previous) for previous in starts)>.8:
            starts.append(point)
        if len(starts)==maximum_modes:break
    if exploration_starts is not None:
        starts.extend(np.asarray(exploration_starts))
    centers=[];covariances=[];diagnostics=[]
    objective=lambda point:-float(log_target(point))
    for start in starts:
        fit=minimize(objective,start,method='L-BFGS-B',bounds=[(-5.,5.)]*dimension,
                     options=dict(maxiter=250,ftol=1e-10,gtol=1e-5,maxls=35))
        point=fit.x
        record=dict(start=start.tolist(),point=point.tolist(),objective=float(fit.fun),
            success=bool(fit.success),message=str(fit.message),evaluations=int(fit.nfev))
        if centers and min(np.linalg.norm(point-center) for center in centers)<.15:
            record['duplicate_center']=True;diagnostics.append(record);continue
        hessian=numerical_hessian(objective,point)
        eigenvalues,eigenvectors=np.linalg.eigh(hessian)
        record['proposal_hessian_eigenvalues']=eigenvalues.tolist()
        # This regularization changes only q; the exact q density appears in all weights.
        inverse_values=1/np.clip(eigenvalues,.1,1e4)
        covariance=np.einsum('ik,k,jk->ij',eigenvectors,inverse_values,eigenvectors)*2.25
        centers.append(point);covariances.append(covariance);diagnostics.append(record)
    if not centers:raise ValueError('no finite importance proposal centers')
    return DefensiveProposal(dimension,centers,covariances),diagnostics


def importance_diagnostics(log_weights):
    weights=np.exp(log_weights-logsumexp(log_weights))
    ordered=np.sort(weights)[::-1]
    return dict(effective_samples=float(1/np.sum(weights**2)),maximum_weight=float(ordered[0]),
        top_ten_weight=float(ordered[:10].sum()),top_one_percent_weight=float(ordered[:max(1,len(weights)//100)].sum()))


def adapt_defensive_proposal(points,log_weights,random,maximum_centers=8):
    """Weighted k-means chooses a new q; pilot points never become final samples."""
    weights=np.exp(log_weights-logsumexp(log_weights))
    effective=1/np.sum(weights**2)
    count=min(maximum_centers,max(4,int(effective/12)))
    mean=np.einsum('n,ni->i',weights,points)
    residual=points-mean
    global_covariance=np.einsum('n,ni,nj->ij',weights,residual,residual)+np.eye(points.shape[1])*1e-4
    scale=np.sqrt(np.diag(global_covariance))
    centers=[points[random.choice(len(points),p=weights)]]
    for _ in range(count-1):
        squared=np.min([np.sum(((points-center)/scale)**2,axis=1) for center in centers],axis=0)
        probability=weights*squared
        centers.append(points[random.choice(len(points),p=probability/probability.sum())])
    centers=np.array(centers)
    for _ in range(30):
        distances=np.array([np.sum(((points-center)/scale)**2,axis=1) for center in centers])
        labels=np.argmin(distances,axis=0)
        updated=[]
        for index in range(count):
            selected=labels==index;mass=weights[selected].sum()
            updated.append(np.einsum('n,ni->i',weights[selected],points[selected])/mass if mass>1e-8 else centers[index])
        updated=np.array(updated)
        if np.max(np.abs(updated-centers))<1e-5:break
        centers=updated
    kept_centers=[];covariances=[];masses=[]
    for index,center in enumerate(centers):
        selected=labels==index;mass=float(weights[selected].sum())
        if mass<1e-5:continue
        residual=points[selected]-center
        covariance=np.einsum('n,ni,nj->ij',weights[selected]/mass,residual,residual)
        covariance=(.9*covariance+.1*global_covariance)*1.44+np.eye(points.shape[1])*1e-4
        kept_centers.append(center);covariances.append(covariance);masses.append(mass)
    return DefensiveProposal(points.shape[1],kept_centers,covariances,prior_fraction=.15,degree=7.,component_masses=masses)