"""
.. module:: BK14
    :synopsis: BK14 likelihood from http://arxiv.org/pdf/1510.09217.pdf, http://bicepkeck.org/bk14_2015_release.html

.. moduleauthor:: Thomas Tram <thomas.tram@port.ac.uk>
Last updated July 20, 2016. Based on the CosmoMC module.
"""
import numpy as np
import pandas as pd
import scipy.linalg as la
import montepython.io_mp as io_mp
import os
import functools
from montepython.likelihood_class import Likelihood_sn

T_CMB = 2.7255     #CMB temperature
h = 6.62606957e-34     #Planck's constant
kB = 1.3806488e-23     #Boltzmann constant
Ghz_Kelvin = h/kB*1e9  #GHz Kelvin conversion

class BK14(Likelihood_sn):

    def __init__(self, path, data, command_line):
        # Unusual construction, since the data files are not distributed
        # alongside BK14 (size problems)
        try:
            # Read the .dataset file specifying the data.
            super(BK14, self).__init__(path, data, command_line)
        except IOError:
            raise io_mp.LikelihoodError(
                "The BK14 data files were not found. Please download the "
                "following link "
                "http://bicepkeck.org/BK14_datarelease/BK14_cosmomc.tgz"
                ", extract it, and copy the BK14 folder inside"
                "`BK14_cosmomc/data/` to `your_montepython/data/`")

        # Require tensor modes from CLASS as well as nonlinear lensing.
        # Nonlinearities enhance the B-mode power spectrum by more than 6%
        # at l>100. (Even more at l>2000, but not relevant to BICEP.)
        # See http://arxiv.org/abs/astro-ph/0601594.
        arguments = {
            'output': 'tCl pCl lCl',
            'lensing': 'yes',
            'modes': 's, t',
            'l_max_scalars': 2000,
            'k_max_tau0_over_l_max': 7.0,
            'non linear':'HALOFIT' if self.do_nonlinear else '',
            'accurate_lensing':1,
            'l_max_tensors': self.cl_lmax}
        self.need_cosmo_arguments(data, arguments)

        map_names_used = self.map_names_used.split()
        map_fields = self.map_fields.split()
        map_names = self.map_names.split()
        self.map_fields_used = [maptype for i, maptype in enumerate(map_fields) if map_names[i] in map_names_used]

        nmaps = len(map_names_used)
        ncrossmaps = nmaps*(nmaps+1)//2
        nbins = int(self.nbins)

        ## This constructs a different flattening of triangular matrices.
        ## v = [m for n in range(nmaps) for m in range(n,nmaps)]
        ## w = [m for n in range(nmaps) for m in range(nmaps-n)]
        ## # Store the indices in a tuple of integer arrays for later use.
        ## self.flat_to_diag = (np.array(v),np.array(w))

        # We choose the tril_indices layout for flat indexing of the triangular matrix
        self.flat_to_diag = np.tril_indices(nmaps)
        self.diag_to_flat = np.zeros((nmaps,nmaps),dtype='int')
        # It is now easy to generate an array with the corresponding flattened indices. (We only fill the lower triangular part.)
        self.diag_to_flat[self.flat_to_diag] = list(range(ncrossmaps))

        # Read in bandpasses
        self.ReadBandpasses()

        # Read window bins
        self.window_data = np.zeros((int(self.nbins),int(self.cl_lmax),ncrossmaps))
        # Retrieve mask and index permutation of windows:
        indices, mask = self.GetIndicesAndMask(self.bin_window_in_order.split())
        for k in range(nbins):
            windowfile = os.path.join(self.data_directory, self.bin_window_files.replace('%u',str(k+1)))
            tmp = pd.read_table(windowfile,comment='#',sep=' ',header=None, index_col=0).to_numpy()
            # Apply mask
            tmp = tmp[:,mask]
            # Permute columns and store this bin
            self.window_data[k][:,indices] = tmp
        # print('window_data',self.window_data.shape)

        #Read covmat fiducial
        # Retrieve mask and index permutation for a single bin.
        indices, mask = self.GetIndicesAndMask(self.covmat_cl.split())
        # Extend mask and indices. Mask just need to be copied, indices needs to be increased:
        superindices = []
        supermask = []
        for k in range(nbins):
            superindices += [idx+k*ncrossmaps for idx in indices]
            supermask += list(mask)
        supermask = np.array(supermask)

        tmp = pd.read_table(os.path.join(self.data_directory, self.covmat_fiducial),comment='#',sep=' ',header=None,skipinitialspace=True).to_numpy()
        # Apply mask:
        tmp = tmp[:,supermask][supermask,:]
        print('Covmat read with shape',tmp.shape)
        # Store covmat in correct order
        self.covmat = np.zeros((nbins*ncrossmaps,nbins*ncrossmaps))
        for index_tmp, index_covmat in enumerate(superindices):
            self.covmat[index_covmat,superindices] = tmp[index_tmp,:]

        #Compute inverse and store
        self.covmat_inverse = la.inv(self.covmat)
        # print('covmat',self.covmat.shape)
        # print(self.covmat_inverse)

        nbins = int(self.nbins)
        # Read noise:
        self.cl_noise_matrix = self.ReadMatrix(self.cl_noise_file,self.cl_noise_order)

        # Read Chat and perhaps add noise:
        self.cl_hat_matrix = self.ReadMatrix(self.cl_hat_file,self.cl_hat_order)
        if not self.cl_hat_includes_noise:
            for k in range(nbins):
                self.cl_hat_matrix[k] += self.cl_noise_matrix[k]

        # Read cl_fiducial and perhaps add noise:
        self.cl_fiducial_sqrt_matrix = self.ReadMatrix(self.cl_fiducial_file,self.cl_fiducial_order)
        if not self.cl_fiducial_includes_noise:
            for k in range(nbins):
                self.cl_fiducial_sqrt_matrix[k] += self.cl_noise_matrix[k]
        # Now take matrix square root:
        for k in range(nbins):
            self.cl_fiducial_sqrt_matrix[k] = la.sqrtm(self.cl_fiducial_sqrt_matrix[k])


    def ReadMatrix(self, filename, crossmaps):
        """
        Read matrices for each ell-bin for all maps inside crossmaps and
        ordered in the same way as usedmaps. Returns list of matrices.

        """
        usedmaps = self.map_names_used.split()
        nmaps = len(usedmaps)
        # Get mask and indices
        indices, mask = self.GetIndicesAndMask(crossmaps.split())
        # Read matrix in packed format
        A = pd.read_table(os.path.join(self.data_directory, filename),comment='#',sep=' ',header=None, index_col=0).to_numpy()
        # Apply mask
        A = A[:,mask]

        # Create matrix for each bin and unpack A:
        Mlist = []
        # Loop over bins:
        for k in range(int(self.nbins)):
            M = np.zeros((nmaps,nmaps))
            Mflat = np.zeros(((nmaps*(nmaps+1))//2))
            Mflat[indices] = A[k,:]
            M[self.flat_to_diag] = Mflat
            # Symmetrise M and append to list:
            Mlist.append(M+M.T-np.diag(M.diagonal()))
        return Mlist

    def GetIndicesAndMask(self, crossmaplist):
        """
        Given a list of used maps and a list of available crossmaps, find a mask
        for the used crossmaps, and for each used crossmap, compute the falttened
        triangular index. We must allow map1 and map2 to be interchanged.
        If someone finds a nicer way to do this, please email me.
        """
        usedmaps = self.map_names_used.split()
        nmaps = len(usedmaps)
        mask = np.array([False for i in range(len(crossmaplist))])

        flatindex = []
        for i, crossmap in enumerate(crossmaplist):
            map1, map2 = crossmap.split('x')
            if map1 in usedmaps and map2 in usedmaps:
                index1 = usedmaps.index(map1)
                index2 = usedmaps.index(map2)
                # This calculates the flat index in a diagonal flattening:
                # if index1 > index2:
                #     flatindex.append((index1-index2)*(2*nmaps+1-index1+index2)/2+index2)
                # else:
                #     flatindex.append((index2-index1)*(2*nmaps+1-index2+index1)/2+index1)
                # This calculates the flat index in the standard numpy.tril_indices() way:
                if index1 > index2:
                    flatindex.append((index1*(index1+1))//2+index2)
                else:
                    flatindex.append((index2*(index2+1))//2+index1)
                mask[i] = True
        return flatindex, mask

    def ReadBandpasses(self):
        """
        Read bandpasses and compute some thermodynamic quantities.
        Everything stored in the dictionary self.bandpasses.
        """
        #Read bandpasses
        self.bandpasses = {}
        map_fields = self.map_fields.split()
        map_names = self.map_names.split()
        map_names_used = self.map_names_used.split()
        for key in map_names_used:
            self.bandpasses[key] = {'field':map_fields[map_names.index(key)],'filename':getattr(self, 'bandpass['+key+']')}

        for key, valdict in io_mp.dictitems(self.bandpasses):
            tmp = np.loadtxt(os.path.join(self.data_directory, valdict['filename']))
            #Frequency nu, response resp:
            valdict['nu'] = tmp[:,0]
            valdict['resp'] = tmp[:,1]
            valdict['dnu'] = np.gradient(valdict['nu'])

            # Calculate thermodynamic temperature conversion between this bandpass
            # and pivot frequencies 353 GHz (used for dust) and 23 GHz (used for
            # sync).
            th_int = np.sum(valdict['dnu']*valdict['resp']*valdict['nu']**4*np.exp(Ghz_Kelvin*valdict['nu']/T_CMB)/(np.exp(Ghz_Kelvin*valdict['nu']/T_CMB)-1.)**2)
            nu0=353.
            th0 = nu0**4*np.exp(Ghz_Kelvin*nu0/T_CMB) / (np.exp(Ghz_Kelvin*nu0/T_CMB) - 1.)**2
            valdict['th353'] = th_int / th0
            nu0=23.
            th0 = nu0**4*np.exp(Ghz_Kelvin*nu0/T_CMB) / (np.exp(Ghz_Kelvin*nu0/T_CMB) - 1.)**2
            valdict['th023'] = th_int / th0
            #print('th353:', valdict['th353'], 'th023:', valdict['th023'])


    def loglkl(self, cosmo, data):
        """
        Compute negative log-likelihood using the Hamimeche-Lewis formalism, see
        http://arxiv.org/abs/arXiv:0801.0554
        """
        # Define the matrix transform
        def MatrixTransform(C, Chat, CfHalf):
            # C is real and symmetric, so we can use eigh()
            D, U = la.eigh(C)
            D = np.abs(D)
            S = np.sqrt(D)
            # Now form B = C^{-1/2} Chat C^{-1/2}. I am using broadcasting to divide rows and columns
            # by the eigenvalues, not sure if it is faster to form the matmul(S.T, S) matrix.
            # B = U S^{-1} V^T Chat U S^{-1} U^T
            B = np.dot(np.dot(U,np.dot(np.dot(U.T,Chat),U)/S[:,None]/S[None,:]),U.T)
            # Now evaluate the matrix function g[B]:
            D, U = la.eigh(B)
            gD = np.sign(D-1.)*np.sqrt(2.*np.maximum(0.,D-np.log(D)-1.))
            # Final transformation. U*gD = U*gD[None,:] done by broadcasting. Collect chain matrix multiplication using reduce.
            M = functools.reduce(np.dot, [CfHalf,U*gD[None,:],U.T,CfHalf.T])
            #M = np.dot(np.dot(np.dot(CfHalf,U*gD[None,:]),U.T),Cfhalf.T)
            return M

        # Recover Cl_s from CLASS, which is a dictionary, with the method
        # get_cl from the Likelihood class, because it already makes the
        # conversion to uK^2.
        dict_Cls = self.get_cl(cosmo, self.cl_lmax)
        # Make short hand expressions and remove l=0.
        ell = dict_Cls['ell'][1:]
        DlEE = ell*(ell+1)*dict_Cls['ee'][1:]/(2*np.pi)
        DlBB = ell*(ell+1)*dict_Cls['bb'][1:]/(2*np.pi)
        # Update foreground model
        self.UpdateForegroundModel(cosmo, data)
        #Make names and fields into lists
        map_names = self.map_names_used.split()
        map_fields = self.map_fields_used
        nmaps = len(map_names)
        ncrossmaps = nmaps*(nmaps+1)//2
        nbins = int(self.nbins)
        # Initialise Cls matrix to zero:
        Cls = np.zeros((nbins,nmaps,nmaps))
        # Initialise the X vector:
        X = np.zeros((nbins*ncrossmaps))
        for i in range(nmaps):
            for j in range(i+1):
                #If EE or BB, add theoretical prediction including foreground:
                if map_fields[i]==map_fields[j]=='E' or map_fields[i]==map_fields[j]=='B':
                    map1 = map_names[i]
                    map2 = map_names[j]
                    dust = self.fdust[map1]*self.fdust[map2]
                    sync = self.fsync[map1]*self.fsync[map2]
                    dustsync = self.fdust[map1]*self.fsync[map2] + self.fdust[map2]*self.fsync[map1]
                    # if EE spectrum, multiply foregrounds by the EE/BB ratio:
                    if map_fields[i]=='E':
                        dust = dust * self.EEtoBB_dust
                        sync = sync * self.EEtoBB_sync
                        dustsync = dustsync * np.sqrt(self.EEtoBB_dust*self.EEtoBB_sync)
                        # Deep copy is important here, since we want to reuse DlXX for each map.
                        DlXXwithforegound = np.copy(DlEE)
                    else:
                        DlXXwithforegound = np.copy(DlBB)
                    # Finally add the foreground model:
                    DlXXwithforegound += (dust*self.dustcoeff+sync*self.synccoeff+dustsync*self.dustsynccoeff)
                    # Apply the binning using the window function:
                    for k in range(nbins):
                        Cls[k,i,j] = Cls[k,j,i] = np.dot(DlXXwithforegound,self.window_data[k,:,self.diag_to_flat[i,j]])
        # Add noise contribution:
        for k in range(nbins):
            Cls[k,:,:] += self.cl_noise_matrix[k]
            # Compute entries in X vector using the matrix transform
            T = MatrixTransform(Cls[k,:,:], self.cl_hat_matrix[k], self.cl_fiducial_sqrt_matrix[k])
            # Add flat version of T to the X vector
            X[k*ncrossmaps:(k+1)*ncrossmaps] = T[self.flat_to_diag]
        # Compute chi squared
        chi2 = np.dot(X.T,np.dot(self.covmat_inverse,X))
        return -0.5*chi2


    def UpdateForegroundModel(self, cosmo, data):
        """
        Update the foreground model.
        """
        # Function to compute f_dust
        def DustScaling(beta, Tdust, bandpass):
            # Calculates greybody scaling of dust signal defined at 353 GHz to specified bandpass.
            nu0 = 353 #Pivot frequency for dust (353 GHz).
            # Integrate greybody scaling and thermodynamic temperature conversion across experimental bandpass.
            gb_int = np.sum(bandpass['dnu']*bandpass['resp']*bandpass['nu']**(3+beta)/(np.exp(Ghz_Kelvin*bandpass['nu']/Tdust) - 1))
            # Calculate values at pivot frequency.
            gb0 = nu0**(3+beta) / (np.exp(Ghz_Kelvin*nu0/Tdust) - 1)
            # Calculate and return dust scaling fdust.
            return ((gb_int / gb0) / bandpass['th353'])

        # Function to compute f_sync
        def SyncScaling(beta, bandpass):
            #Calculates power-law scaling of synchrotron signal defined at 150 GHz to specified bandpass.
            nu0 = 23.0 # Pivot frequency for sync (23 GHz).
            # Integrate power-law scaling and thermodynamic temperature conversion across experimental bandpass.
            pl_int = np.sum( bandpass['dnu']*bandpass['resp']*bandpass['nu']**(2+beta))
            # Calculate values at pivot frequency.
            pl0 = nu0**(2+beta)
            # Calculate and return dust scaling fsync.
            return ((pl_int / pl0) / bandpass['th023'])


        ellpivot = 80.
        ell = np.arange(1,int(self.cl_lmax)+1)

        # Convenience variables: store the nuisance parameters in short named variables
        # for parname in self.use_nuisance:
        #     evalstring = parname+" = data.mcmc_parameters['"+parname+"']['current']*data.mcmc_parameters['"+parname+"']['scale']"
        #     print(evalstring)
        BBdust = data.mcmc_parameters['BBdust']['current']*data.mcmc_parameters['BBdust']['scale']
        BBsync = data.mcmc_parameters['BBsync']['current']*data.mcmc_parameters['BBsync']['scale']
        BBalphadust = data.mcmc_parameters['BBalphadust']['current']*data.mcmc_parameters['BBalphadust']['scale']
        BBbetadust = data.mcmc_parameters['BBbetadust']['current']*data.mcmc_parameters['BBbetadust']['scale']
        BBTdust = data.mcmc_parameters['BBTdust']['current']*data.mcmc_parameters['BBTdust']['scale']
        BBalphasync = data.mcmc_parameters['BBalphasync']['current']*data.mcmc_parameters['BBalphasync']['scale']
        BBbetasync = data.mcmc_parameters['BBbetasync']['current']*data.mcmc_parameters['BBbetasync']['scale']
        BBdustsynccorr = data.mcmc_parameters['BBdustsynccorr']['current']*data.mcmc_parameters['BBdustsynccorr']['scale']

        # Store current EEtoBB conversion parameters.
        self.EEtoBB_dust = data.mcmc_parameters['EEtoBB_dust']['current']*data.mcmc_parameters['EEtoBB_dust']['scale']
        self.EEtoBB_sync = data.mcmc_parameters['EEtoBB_sync']['current']*data.mcmc_parameters['EEtoBB_sync']['scale']

        # Compute fdust and fsync for each bandpass
        self.fdust = {}
        self.fsync = {}
        for key, bandpass in io_mp.dictitems(self.bandpasses):
            self.fdust[key] = DustScaling(BBbetadust, BBTdust, bandpass)
            self.fsync[key] = SyncScaling(BBbetasync, bandpass)

        # Computes coefficients such that the foreground model is simply
        # dust*self.dustcoeff+sync*self.synccoeff+dustsync*self.dustsynccoeff
        # These coefficients are independent of the map used,
        # so we save some time by computing them here.
        self.dustcoeff = BBdust*(ell/ellpivot)**BBalphadust
        self.synccoeff = BBsync*(ell/ellpivot)**BBalphasync
        self.dustsynccoeff = BBdustsynccorr*np.sqrt(BBdust*BBsync)*(ell/ellpivot)**(0.5*(BBalphadust+BBalphasync))
