'''
Version del codigo completa para obtener los coeficientes
PASOS:
    1. Creamos un grid icosahedrico, y proyectamos sobre el 
    2. Sobre esta proyeccion que "conserva" el radio de los datos de simulaciones se realiza expansion por arm.es
    3. De alm se convierte a al
    FALTA. Se eligen los modos del l = 2 a l = 6, y se realiza ajuste por la funcion (ecuacion 7)
    
    NOTA: No es necesario especifiar Lmin, parte del 0 el programa
    
'''

# %% LIBRERIAS

import sys,os

from itertools import combinations
from tqdm import tqdm

import logging

import numpy as np

import orthopoly

import MDAnalysis as mda
from MDAnalysis.analysis import align

from scipy import spatial
from scipy.optimize import least_squares
from scipy.optimize import lsq_linear


# %% FUNCIONES PREVIAS

def public_c2s(posicionescart):
    '''
    De cartesianas a esfericas (c2s) de forma publica
    '''
    x,y,z       = posicionescart.T
    r,theta,phi = orthopoly.spherical_harmonic.cart2sph(-x,-y,z) #Bug in library
    return np.asarray([r,theta,phi]).T

def public_s2c(posicionesesf):
    '''
    De esfericas a cartesianas (s2c) de forma publica
    '''
    r,theta,phi = posicionesesf.T
    x,y,z       = orthopoly.spherical_harmonic.sph2cart(r,theta,phi)
    return np.asarray([x,y,z]).T

def msph_mcart(matrizsph):
    tiempo = matrizsph.shape[0]
    matrizcart = (public_s2c(matrizsph.reshape(-1, 3))).reshape(tiempo, -1, 3)
    return matrizcart

def mcart_msph(matrizcart):
    tiempo = matrizcart.shape[0]
    matrizsph = (public_c2s(matrizcart.reshape(-1, 3))).reshape(tiempo, -1, 3)
    return matrizsph

# %% DEFINICION DE LA CLASE

'''
NOTA: phi [0,2pi]
    theta [0,pi]
'''

class SHD:

    def __cart2sph(self,cartPositions):
        x,y,z       = cartPositions.T
        r,theta,phi = orthopoly.spherical_harmonic.cart2sph(-x,-y,z) #Bug in library
        return np.asarray([r,theta,phi]).T

    def __sph2cart(self,sphPositions):

        r,theta,phi = sphPositions.T
        x,y,z       = orthopoly.spherical_harmonic.sph2cart(r,theta,phi)
        return np.asarray([x,y,z]).T
            
    def __degreeOfTruncationToLM(self,n):
        A = orthopoly.spherical_harmonic.Tnm(self.__Lmax)
        l,m = A[0][n],A[1][n]
        return l,m

    def __init__(self,name="SHD",Lmin=2,Lmax=10,debug=False,
                 radiusMode    = "geometric",
                 expansionMode = "abs"):

        self.__availableRadiusModes    = ["expansion","geometric"]
        self.__availableExpansionModes = ["fluct","abs"]

        self.__DEBUG = debug

        self.__Lmin = Lmin
        self.__Lmax = Lmax
        self.__name = name

        self.logger = logging.getLogger("sphAnalysis")
        #Clean up the logger
        for handler in self.logger.handlers[:]:
            self.logger.removeHandler(handler)
        self.logger.setLevel(logging.DEBUG)
        #Formatter, ascii time to day, month, year, hour, minute, second
        formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s',"%d-%b-%Y %H:%M:%S")

        #Set logger to console, with level INFO
        ch = logging.StreamHandler()
        ch.setLevel(logging.INFO)
        ch.setFormatter(formatter)

        #Create other logger. Log to file, with level DEBUG
        fh = logging.FileHandler(f'{self.__name}.log')
        fh.setLevel(logging.DEBUG)
        fh.setFormatter(formatter)

        self.logger.addHandler(ch)
        self.logger.addHandler(fh)

        ######################################################

        if radiusMode not in self.__availableRadiusModes:
            self.logger.error("Radius mode {} not available. Available modes are: {}".format(radiusMode,self.__availableRadiusModes))
            raise ValueError
        else:
            self.__radiusMode = radiusMode

        if expansionMode not in self.__availableExpansionModes:
            self.logger.error("Traj mode {} not available. Available modes are: {}".format(expansionMode,self.__availableExpansionModes))
            raise ValueError
        else:
            self.__expansionMode = expansionMode

        if self.__radiusMode == "expansion" and self.__expansionMode == "fluct":
            #Not compatible
            self.logger.error("Radius mode {} and traj mode {} not compatible".format(self.__radiusMode,self.__expansionMode))
            raise ValueError
    
   
    def __setGrid(self):
        self.gridSize = self.thetaGrid.shape[0]
        self.grid = np.asarray([self.thetaGrid,self.phiGrid]).T

        r=[1.0]*self.gridSize
        self.gridVectors = self.__sph2cart(np.asarray([r,self.thetaGrid,self.phiGrid]).T)
        self.gridVectorssph = self.__cart2sph(self.gridVectors)

        self.gridKDtree = spatial.KDTree(self.gridVectors)
     
    def generateIcosahedralGrid(self,n):
        
        print("----------")
        self.logger.info("Generating icosahedral grid with {} subdivisions ...".format(n))

        self.thetaGrid,self.phiGrid=orthopoly.spherical_harmonic.grid_icosahedral(n)
        self.__setGrid()
        

    def distributeTrajPointsAlongGrid(self):
        """
        Distribute the trajectory points along the grid.
        """
        self.traj = self.traj2 # --> self.traj se deberia obtener desde loadsimulation
        
        self.time = int(self.traj.shape[0]) # Numero de "tiempos"
        
        self.trajSph = self.__cart2sph(self.traj.reshape(-1,3)).reshape(self.time,-1,3)
        
        #Check if the trajectory has been loaded
        if self.traj is None:
            self.logger.error("Trying to distribute trajectory points along grid, but no trajectory has been loaded.")
            raise RuntimeError("No trajectory loaded.")

        #Check if the grid has been generated
        if self.grid is None:
            self.logger.error("Trying to distribute trajectory points along grid, but no grid has been generated.")
            raise RuntimeError("No grid generated.")

        self.logger.info("Distributing trajectory points along grid ...")

        r,theta,phi = self.__cart2sph(self.traj.reshape(-1,3)).T

        #Projects points over sphere
        r=np.asarray([1.0]*r.shape[0])
        trajOverSphere = self.__sph2cart(np.asarray([r,theta,phi]).T)

        self.trajBeadsGridPosition = self.gridKDtree.query(trajOverSphere)[1].reshape(self.time,-1)

        self.trajGridSph = np.empty((self.time,self.gridSize,3))
        for f,[framePos,frameGrid] in enumerate(tqdm(zip(self.trajSph,self.trajBeadsGridPosition),total=self.time)):
            for i in range(self.gridSize):
                mask = (frameGrid == i)

                if not np.any(mask):
                    self.logger.error("Grid point with no data. Try to decrease the number of grid points.")
                    raise Exception("Grid point with no data")

                r,theta,phi = framePos[mask].T
                
                # f es cada tiempo, i es cada fila del grid que se calcula con el valor medio de los particulas
                # mas cercanas al grid

                self.trajGridSph[f,i] = np.asarray([np.mean(r),self.thetaGrid[i],self.phiGrid[i]])
        
        self.trajGrid = self.__sph2cart(self.trajGridSph)

        if self.__DEBUG:

            self.logger.debug("Distributed trajectory points along grid. Writing to file ...")

            with open("gridPos.sp","w") as f:
                for framePos,frameGrid in zip(self.traj,self.trajBeadsGridPosition):
                    f.write("#\n")
                    for p,g in zip(framePos,frameGrid):
                        f.write("{} {} {} 1.0 {}\n".format(p[0],p[1],p[2],g))

            with open("gridTraj.sp","w") as f:
                for frame in self.trajGridSph:
                    f.write("#\n")
                    for p in frame:
                        x,y,z = self.__sph2cart(np.asarray([p[0],p[1],p[2]]).T).T
                        f.write("{} {} {} 1.0 0\n".format(x,y,z))
                        
    def sphericalHarmonicExpansion(self,TGS):
        
        '''
        TGS (esf): Son los datos que queremos desarrollar, entonces deberian ser los datos que nos de onlytheta_mitdpi matriz 3D
        Objetivo --> Calcular los coeficientes
        '''
      
        if self.trajGridSph is None:
            self.logger.error("Trying to compute spherical harmonic expansion, but trajectory has not been distributed along grid.")
            raise RuntimeError("No trajectory distributed along grid.")

        self.logger.info("Performing spherical harmonic expansion ...")
        
        v1 = TGS.shape[1]
        v2 = int((self.__Lmax+1)**2)
        
        MATRIXY = np.empty((self.time,v1,v2))
        
        self.aCoeffComputed      = []
        
        # THETA Y PHI NO CAMBIAN EN EL TIEMPO (cambia solo r)
        
        th1 = TGS[0,:,1]
        ph1 = TGS[0,:,2]
        
        infoY = orthopoly.spherical_harmonic.sph_har_T_matrix(th1,ph1,self.__Lmax)
        Y = infoY[0] # Matriz Ylm, como le conocemos theta y phi conocemos esta matriz
        
        for i in range(self.time):
            
            f,theta,phi = TGS[i].T

            aFluct = lsq_linear(Y,f,max_iter=1000,tol=0.000001,lsq_solver='lsmr') # Sistema de ecuaciones R = Ylm*A donde conocemos R e Ylm
            self.aCoeffComputed.append(aFluct.x)
        
        self.aCoeffComputed = np.asarray(self.aCoeffComputed)
        self.MatrizY = Y;
        
    def valoresmedios(self,M1):
        
        # Mensaje de advertencia
        
        if self.aCoeffComputed is None:
            self.logger.error("No se puede calcular valores medios porque no esta definido self.aCoeffComputed")
            self.logger.error("--> Comprobar función: sphericalHarmonicExpansion")
            raise RuntimeError("No se hizo nada")
        self.logger.info("Realizando el calculo de valores medios...")
        
        '''
        Tenemos los armonicos esfericos en formato de matriz de dimensiones (tiempo x (lmax+1)**2)
        En esta función calculamos:
            1) <alm**2> : Elevamos la matriz al cuadrado, despues promediamos temporalmente
            2) <alm> : Promediamos temporalmente (por columnas) # dimensiones de Lmax+1
        
        Calcularemos el radio promedio tambien (espacial y temporal)
        --> M1 es la matriz (en esfericas) de dimensiones tiempoxnumero de particulas o numero puntos grid x3 
        matriz proyectada para calcular R0
        '''
        # CALCULO DE VALORES MEDIOS
        
        Vini = self.aCoeffComputed;
        
        self.VM_CoefsCuadrado = np.mean(Vini**2,axis = 0) # Valor medio a lo largo de las filas (tiempo)
        self.VM_Coefs = np.mean(Vini,axis = 0)
        self.VMCuadrado_Coefs = self.VM_Coefs**2
        
        # CALCULO DE LA LIBRERIA (nos permite conocer l y m a partir de n)
        
        N = orthopoly.spherical_harmonic.T2nharm(self.__Lmax)
        
        lm2n = {}
        
        for n in range(N):
            l,m = self.__degreeOfTruncationToLM(n)
            lm2n[(l,m)] = n
        
        # CALCULO DE <al**2> (ya no depende de m porque sumamos)
        
        L   = []
        a2l = []
        
        for l in range(self.__Lmax+1):
            
            L.append(l)
            
            a2 = 0
            for m in range(-l,l+1):

                a2       += self.VM_CoefsCuadrado[lm2n[(l,m)]]/(2*l+1)

            a2l.append(a2)
            
        va2l = np.asarray(a2l)
        self.VM_AlCuadrado = va2l # ESTO ES LO QUE NOS INTERESA
        
        radial = M1[:,:,0] # dimensiones de tiempo x particulas
        self.R0 = np.mean(radial)