import numpy as np 
import sympy as sp



class calcTensor:
    '''
    ***Requires numpy as 'np' and sympy as 'sp' installation (import) to load.***

    **Class:** Functions to compute tensor identities.

    **Functions:**
        * coordinateGenerator
        * metricGenerator
        * deltaFunction
        * chrisUDD
        * riemannUDDD
        * ricciDD
        * ricciScalar
        * retrieveAll
        * classInfo
    '''
    def __init__(self):
        """
        Initialize the class.
        """
        pass


    def write_dict(
            array: list[list]
            ) -> dict:
        """
        Write a dictionary with all the non-zero elements of an array and their respective incides.
        """
        nonzero_dict = {}

        if len(array.shape) == 2:
            N = array.shape[0]
            M = array.shape[1]

            for i in range(N):
                for j in range(M):
                    if array[i][j] != 0:
                        nonzero_dict[(i, j)] = array[i][j]


        if len(array.shape) == 3:
            N = array.shape[0]
            M = array.shape[1]
            K = array.shape[2]

            for i in range(N):
                for j in range(M):
                    for l in range(K):
                            if array[i][j][l] != 0:
                                nonzero_dict[(i, j, l)] = array[i][j][l]

        if len(array.shape) == 4:
            N = array.shape[0]
            M = array.shape[1]
            K = array.shape[2]
            J = array.shape[3]

            for i in range(N):
                for j in range(M):
                    for l in range(K):
                        for o in range(J):
                            if array[i][j][l][o] != 0:
                                nonzero_dict[(i, j, l, o)] = array[i][j][l][o]

        return "\n".join(['%s : %s' % (key, str(value)) for key, value in nonzero_dict.items()])
    
    sp.init_printing()

    
    def coordinateGenerator(
            vars: list[str]
            ) -> list[str]:
        """
        **Returns:** dtype = numpy.array. List of coordinate variables in N-vector contravariant tensor notation. 

        *vars:* List of symbolic strings generated by sympy.symbols.
        """
        return np.array(vars, dtype=object)
    

    def metricGenerator(
            vars: list[str]
            ) -> tuple[list[str]]:
        """
        **Returns:** dtype = numpy.array. Metric tensor lower 'g_dndn' and upper 'g_upup' for symbolic calculation (tuple). Dimensionality out is equivalent to dimensionality in.

        *vars:* Input array for 'g_dndn' metric of dimensionality NxN of symbolic strings generated by sympy.symbols. 
        """
        dims = np.shape(vars)
        if dims[0] != dims[1]:
            raise ValueError('Input array-list *vars* must be of dimensionality NxN.')
        
        gdndn = np.array(vars, dtype=object)

        gupup = sp.Matrix(gdndn).inv()
        gupup = np.array(gupup, dtype=object)
        return gdndn, gupup
    

    def deltaFunction(
            dimension: int
            ) -> list[int]:
        """
        **Returns:** dtype = numpy.array. Kronecker Delta function of dimensionality N. Iterable through by tensor indices.

        *dimension:* Input integer to specify dimension output array. 
        """
        if type(dimension) != int:
            raise ValueError('Input variable *dimension* must be an integer.')

        N = dimension
        indexlist = np.arange(N)
        arr = np.zeros((N, N))
        for n in indexlist:
            arr[n][n] = 1

        return arr
    

    def chrisUDD(
            g_upup: list[str],
            g_dndn: list[str], 
            xup: list[str]
            ) -> list[str]:
        """
        **Returns:** dtype = numpy.array. NxNxN dimensional array of symbolic Christoffel symbols of tensor type 'up-down-down' generated from upstairs and downstairs metric tensors. 

        *g_upup:* dtype = numpy.array. NxN upstairs metric tensor of symbolic sympy values.

        *g_dndn:* dtype = numpy.array. NxN downstairs metric tensor of symbolic sympy values.

        *xup:* dtype = numpy.array. Nx1 vector list of symbolic variables generated from sympy.symbols.
        """
        N = len(xup)
        indexlist = np.arange(N)

        vdotfill = np.ones(N)
        Chrisupdndn = np.empty((N,N,N), dtype=object)
        contract = np.empty((N), dtype=object)

        for mu in indexlist:
            for nu in indexlist:
                for la in indexlist:
                    for si in indexlist: 
                        contract[si] = 0.5*(g_upup[mu][si])*(sp.diff(g_dndn[si][la], xup[nu]) + sp.diff(g_dndn[nu][si], xup[la]) - sp.diff(g_dndn[nu][la], xup[si]))
                    Chrisupdndn[mu][nu][la] = sp.simplify(contract.dot(vdotfill))  


        Chrisupdndn = sp.simplify(Chrisupdndn)
        
        return Chrisupdndn, calcTensor.write_dict(Chrisupdndn)
    


    def riemannUDDD(
            christoffel: list[str], 
            xup: list[str]
            ) -> list[str]:
        """
        **Returns:** dtype = numpy.array. NxNxNxN dimensional array of symbolic Riemann tensor of type 'up-down-down-down' generated from Christoffel symbols. 

        *christoffel:* dtype = numpy.array. NxNxN Christoffel symbol tensor of symbolic sympy values.

        *xup:* dtype = numpy.array. 1xN vector list of symbolic variables generated from sympy.symbols.
        """
        N = len(xup)
        indexlist = np.arange(N)

        christoffel = np.array(christoffel)

        vdotfill = np.ones(N)
        Riemupdndndn = np.empty((N, N, N, N), dtype=object)
        contract = np.empty((N), dtype=object)

        for rho in indexlist:
            for si in indexlist:
                for mu in indexlist:
                    for nu in indexlist:
                        for la in indexlist:    
                            contract[la] = (christoffel[la][si][nu]*christoffel[rho][la][mu]) - (christoffel[la][si][mu]*christoffel[rho][la][nu])
                        Riemupdndndn[rho][si][mu][nu] = contract.dot(vdotfill) + sp.diff(christoffel[rho][si][nu], xup[mu]) - sp.diff(christoffel[rho][si][mu], xup[nu])  
        
        Riemupdndndn = sp.simplify(Riemupdndndn)

        return Riemupdndndn, calcTensor.write_dict(Riemupdndndn)


    def ricciDD(
            riemann: list[str]
            ) -> list[str]:
        """
        **Returns:** dtype = numpy.array. NxN dimensional array of symbolic Ricci tensor of type 'down-down' generated from Riemann tensor. 

        *riemann:* dtype = numpy.array. NxNxNxN Riemann tensor of symbolic sympy values.
        """
        N = np.shape(riemann)[0]
        indexlist = np.arange(N)

        riemann = np.array(riemann)

        vdotfill = np.ones(N)
        Riccdndn = np.empty((N, N), dtype = object)
        contract = np.empty(N, dtype = object)
        
        for mu in indexlist:
            for nu in indexlist:
                for la in indexlist:        
                    contract[la] = riemann[la][mu][nu][la]
                Riccdndn[mu][nu] = sp.simplify(contract.dot(vdotfill)) 

        return Riccdndn, calcTensor.write_dict(Riccdndn)
    
    def ricciScalar(
            ricciTensor: list[str],
            g_upup: list[str]
            ) -> str:
        """
        **Returns:** dtype = string. Symbolic scalar value of type sympy variable. Contraction along two axes of Ricci tensor type 'down-down'.

        *ricciTensor:* dtype = numpy array. NxN Ricci curvature tensor of symbolic sympy values.

        *g_upup:* dtype = numpy array. NxN metric tensor of symbolic sympy values of type 'up-up' indexing.
        """
        N = np.shape(g_upup)[0]
        indexlist = np.arange(N)

        vdotfill = np.ones(N)
        contract2 = np.empty(N, dtype=object)
        contract1 = np.empty(N, dtype=object)
        for mu in indexlist:
            for nu in indexlist:
                contract1[nu] = ricciTensor[mu][nu] * g_upup[mu][nu]     #determine the  nu-th term
            contract2[mu] = sp.simplify(contract1.dot(vdotfill))  #contract over nu 
        RicciScalar = sp.simplify(contract2.dot(vdotfill))  #now contract over mu 

        return RicciScalar


    def retrieveAll(
            metric_vars: list[str], 
            coordinates: list[str], 
            return_type: str
            ) -> None:
        """
        **Returns:** Any chosen value determine from previous functions. 

        *metric_vars:* numpy.array. Input array for 'g_dndn' metric of dimensionality NxN of symbolic strings generated by sympy.symbols.

        *coordinates:* numpy.array. List of symbolic strings generated by sympy.symbols.

        *return_type:* string. 'CC' for Christoffel symbols up-down-down. 'RMT' for Riemann tensor up-down-down-down. 'RIT' for Ricci tensor down-down. 'RS' for Ricci Scalar.
        """

        xup = calcTensor.coordinateGenerator(coordinates)
        gdndn, gupup = calcTensor.metricGenerator(metric_vars)

        christoffels = calcTensor.chrisUDD(gupup, gdndn, xup)

        riemanntensor = calcTensor.riemannUDDD(christoffels[0], xup)

        riccitensor = calcTensor.ricciDD(riemanntensor[0])
        ricciscalar = calcTensor.ricciScalar(riccitensor[0], gupup)


        for retry in range(5):
            if return_type == 'MT':
                print('g_dndn: {}'.format(gdndn))
                print('g_upup: {}'.format(gupup))
                break 
            if return_type == 'CC':
                print(christoffels[1])
                break
            elif return_type == 'RMT':
                print(riemanntensor[1])
                break
            elif return_type == 'RIT':
                print(riccitensor[1])
                break
            elif return_type == 'RIS':
                print(ricciscalar)
                break
            elif return_type == 'ALL':
                print('----------')
                print('xup: {}'.format(xup))
                print('g_dndn: {}'.format(gdndn))
                print('g_upup: {}'.format(gupup))
                print('----------')
                print('christoffel: {}'.format(christoffels[1]))
                print('----------')
                print('riemann tensor: {}'.format(riemanntensor[1]))
                print('----------')
                print('ricci tensor: {}'.format(riccitensor[1]))
                print('----------')
                print('ricci scalar: {}'.format(ricciscalar))
                print('----------')
                break
            else:
                return_type = input('Please enter a valid return_type string (MT, CC, RMT, RIT, RIS, ALL):')
        else: 
            print('You keep making invalid choices. Now exiting.')



    def classInfo() -> any:
        """
        Retrieve class 'calcTensor' documentation.
        """
        print(help(calcTensor))