"""Original fixed-weight leave-one-out residual correlation matrices."""
from pathlib import Path
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np


def covariance(a,v):
    a,v=np.asarray(a,dtype=float),np.asarray(v,dtype=float)
    assert np.isclose(a.sum(),1) and np.all(a>0) and np.all(a<1) and np.all(v>0)
    p=np.eye(len(a))-np.outer(np.ones(len(a)),a)
    transform=np.linalg.solve(np.diag(1-a),p)
    cov=transform@np.diag(v)@transform.T
    s=np.sum(a*a*v)
    predicted=(s-a[:,None]*v[:,None]-a[None,:]*v[None,:])/np.outer(1-a,1-a)
    np.fill_diagonal(predicted,(v*(1-2*a)+s)/(1-a)**2)
    assert np.allclose(cov,predicted)
    assert np.allclose(cov@(a*(1-a)),0)
    assert np.linalg.matrix_rank(cov)==len(a)-1
    assert np.linalg.eigvalsh(cov).min()>-1e-12
    return cov


def main():
    cases=[np.repeat(1/3,3),np.array([.05,.05,.9])]
    matrices=[]
    for a in cases:
        cov=covariance(a,np.ones(3))
        matrices.append(cov/np.sqrt(np.outer(np.diag(cov),np.diag(cov))))
    assert np.allclose(matrices[0][~np.eye(3,dtype=bool)],-.5)
    assert np.isclose(covariance(cases[1],np.ones(3))[0,1],286/361)
    v=np.array([1.,4.,16.]);a=(1/v)/(1/v).sum()
    cov=covariance(a,v)
    assert np.all(cov[~np.eye(3,dtype=bool)]<0)
    # Independent arithmetic check for the held-out weighted-mean example.
    values=np.array([10.,12.,8.]);u=np.array([1.,2.,4.]);w=1/u**2
    mean=(w@values)/w.sum()
    assert np.isclose(mean,72/7) and np.isclose(.5**2+1/w.sum(),85/84)
    fig,axes=plt.subplots(1,2,figsize=(10,4.2),dpi=100,facecolor='white')
    fig.subplots_adjust(left=.08,right=.88,bottom=.20,top=.77,wspace=.50)
    titles=['Equal coefficients','Third result dominates benchmarks']
    for ax,r,a,title in zip(axes,matrices,cases,titles):
        im=ax.imshow(r,vmin=-1,vmax=1,cmap='RdBu_r')
        for i in range(3):
            for j in range(3):
                ax.text(j,i,f'{r[i,j]:+.2f}',ha='center',va='center',fontsize=12,
                        color='white'if abs(r[i,j])>.55 else '#202020')
        ax.set_xticks(range(3),['1','2','3']);ax.set_yticks(range(3),['1','2','3'])
        ax.set_xlabel('Laboratory');ax.set_ylabel('Laboratory')
        ax.set_title(title,fontsize=11,pad=10)
        ax.text(.5,-.25,'a = ('+', '.join(f'{x:.2f}'for x in a)+')',transform=ax.transAxes,ha='center',fontsize=10)
    cax=fig.add_axes([.91,.28,.016,.43]);bar=fig.colorbar(im,cax=cax);bar.set_label('Residual correlation',fontsize=10)
    fig.suptitle('Leave-one-out residuals need not have negative pairwise correlation',fontsize=13,y=.97)
    fig.text(.5,.85,'Original models: three independent observations, each with variance 1',ha='center',fontsize=10)
    fig.savefig(Path.cwd()/'leave-one-out-residual-correlation-patterns.png',dpi=100,facecolor='white')
    plt.close(fig)


if __name__=='__main__':
    main()