#!/usr/bin/python3  -O
from utils import ProgressBar
from config_manager import ConfigReader
import pandas
import numpy as np
from sklearn.covariance import GraphicalLassoCV, LedoitWolf
from sklearn import cluster, covariance, manifold
from math import *
import dill
import sys
if not sys.warnoptions:
    import warnings
    warnings.simplefilter("ignore")
    
class DistList:
    def __init__(self, covindex,chunks,covlist,ilist):
        self._covindex = covindex
        self._chunks = chunks
        self._covlist = covlist
        self._ilist = ilist    
        self._index = None
        
    def hellinger(self,n,x,y,smooth,**kwargs):        
        x1 = x.reshape(n,n+1)
        x2 = y.reshape(n,n+1)
        mu1,sigma1=x1[:,0],x1[:,1:]
        mu2,sigma2=x2[:,0],x2[:,1:]
       
        sign1, ld1 = np.linalg.slogdet(sigma1)
        sign2, ld2 = np.linalg.slogdet(sigma2)
       
        dd = 0.5*(sigma1+sigma2)
        
        signdd, ldd = np.linalg.slogdet(dd)
        
        signt = (sign1 * sign2)/signdd
        if signt <= 0:
            print ('negative')
        s_1 = np.linalg.inv(dd)
        e = -0.125*(mu1-mu2).T.dot(s_1).dot(mu1-mu2)
        #d = pow(d1,.25)*pow(d2,.25)/pow(det1,.5)
        d = .25*(ld1+ld2-2*ldd)
        d = exp(d)
        es = 1-d*exp(e)
        #if es <=0:
        #    es = 0
        if smooth:
            r = atanh(es)
        else:
            r = sqrt(es)
        return r    
     
    def vdist(self,x,centroids,sm):
        n = self._covlist.shape[1]
        V = np.full((len(centroids)),0.)
        for i,u in enumerate(centroids):
            V[i] = self.hellinger(n,x,self._covlist[u[1]],sm)    
        return V         


class Distributions:
    def __init__(self,varfact=1.1,eps=.2, window=90, overlap=60): 
        self._varfact=varfact
        self.eps=eps
        self._window=window
        self._overlap=overlap
        
    def compute_gaussians(self,window,overlap):
        i = 0
        r = None
        k = 0
        # i= 
        ret = None
        rw=[]
        avdi = np.full((self._N,),0.)
        ilist=[]
        pr = ProgressBar(self._logret.shape[0]/(window-overlap), prefix='Covariances Computed', suffix='Complete')
        while i < self._logret.shape[0]:
            iend = min(i+window,self._logret.shape[0])                
            chunk = self._logret[i:iend,:]#*100.
            ilist.append((i,iend))
            rw.append(chunk)
            pr.iterate()
            
            mu = chunk.mean(axis=0)
            model = GraphicalLassoCV(cv=3)
            model.fit(chunk.astype('float64'))
            cov = model.covariance_
            di = np.diag(cov)
            avdi += di
            np.fill_diagonal(cov,di*self._varfact)                
            # check spd 
            #    ch = np.linalg.cholesky(cov) 
            r = None
            r_ = np.hstack((mu.reshape(-1,1),cov))
            if ret is None:
                ret = r_.reshape(1,r_.shape[0],r_.shape[1])          
            else:
                ret = np.vstack((ret,r_.reshape(1,r_.shape[0],r_.shape[1])))
            if iend < self._logret.shape[0]:
                i = iend-overlap        
            else:
                break
            
        dlist  = DistList( np.arange(ret.shape[0]), rw,ret,ilist)
        return dlist
    
    def load(self,mfile,mfields):
        pd = pandas.read_csv(mfile,sep=',')
        self._index = pd['Date'].values
        self._logret = pd[mfields].values
        self._N = self._logret.shape[1]
       
def main():
    creader = ConfigReader('regimes.ini')
    mfile = creader.distrmastersource()
    mfields = creader.distrfields()
    D = Distributions(varfact=1.1,eps=.2, window=90, overlap=60)
    D.load(mfile, mfields)
    window = creader.distrws()
    overlap = creader.distroverlap()
    dlist = D.compute_gaussians(window, overlap)
    dlist._index=D._index
    
    mcovf = creader.mastercovlist()
    print()
    print('list of covariances saved in..')
    with open(mcovf, 'wb') as f:
        dill.dump(dlist,f)    
    print (mcovf)
        
if __name__ == "__main__":
    main()