import numpy as np
import copy
import tensorflow as tf
from collections.abc import Iterable, Iterator
import itertools



def poisson_matrix(m,n):
    '''
    Generates the matrix A to express the Poisson equation in the form Ax=b for an m-by-n grid
    
    The matrix returned shall be (m-2)*(n-2)-by-(m-2)*(n-2) in size
    
    YOU MUST RESHAPE THE RESULT FROM (i.e. solution = inv(poisson_matrix) * right_hand_size) FROM (...,(m-2)*(n-2)) TO (...,m-2,n-2) BY USING FORTRAN COLUMN-MAJOR ORDERING!!!!!!!!
    '''
    m = m-2
    n = n-2
    
    D = np.zeros((m,m), dtype = np.float64)
    i,j = np.indices(D.shape)
    D[i==j] = 4.0
    D[i==j-1] = -1.0
    D[i==j+1] = -1.0
    
    S = -np.eye(D.shape[0], dtype = np.float64)
    
    P = np.zeros((m*n,m*n), dtype = np.float64)
    ind = np.arange(0,m*(n+1), m)
    
    for i in range(len(ind)-1):
        P[ind[i]:ind[i+1], ind[i]:ind[i+1]] = D
        try:
            P[ind[i+1]:ind[i+2], ind[i]:ind[i+1]] = S
        except:
            pass
        try:
            P[ind[i-1]:ind[i], ind[i]:ind[i+1]] = S
        except:
            pass
    return P

    

def poisson_RHS(F, boundaries = None, h = None, rho = None):
    '''
    Generates the RHS vector b of a discretized Poisson problem in the form Ax=b.
    h = grid spacing. if not supplied, the RHS will be returned without the (delta x)^2 factor. necessary for nonuniform grid spacing.
    boundaries = dict containing entries 'top', 'bottom', 'right' and 'left' which correspond to the Dirichlet BCs at these boundaries. Each entry must be a vector of length m or n, where m and n are defined as in te function poisson_matrix
    F = an m by n matrix containing the RHS values of the Poisson equation
    
    (i.e. this function merely takes the BC information and the (structured) RHS array to provide the RHS for the matrix eq. form)
    '''
    #print(F)
    if isinstance(F, list):
        boundaries = F[1]
        try:
            h = F[2]
        except:
            pass
        try:
            rho = F[3]
        except:
            pass
        F = F[0]
    
    boundaries = copy.deepcopy(boundaries)
    F = copy.deepcopy(F)
    
    for key in boundaries.keys():
        if len(boundaries[key].shape) > 1:
            boundaries[key] = tf.squeeze(boundaries[key])
        if len(boundaries[key].shape) == 1:
            boundaries[key] = itertools.repeat(boundaries[key], F.shape[0])
    
    if (h is not None) and isinstance(h, float):
        h = itertools.repeat(h, F.shape[0])
    
    for i in range(F.shape[0]):
        if h is not None:
            try:
                dx = next(h)
            except:
                dx = h[i]
            F[i] = -dx**2 * F[i]
                
        
        if isinstance(boundaries['top'], Iterator):
            top = next(boundaries['top'])[1:-1] 
        else:
            top = boundaries['top'][i,1:-1]
            
        if isinstance(boundaries['bottom'], Iterator):
            bottom = next(boundaries['bottom'])[1:-1]
        else:
            bottom = boundaries['bottom'][i,1:-1]
            
        if isinstance(boundaries['left'], Iterator):
            left = next(boundaries['left'])[1:-1]
        else:
            left = boundaries['left'][i,1:-1]
            
        if isinstance(boundaries['right'], Iterator):
            right = next(boundaries['right'])[1:-1]
        else:
            right = boundaries['right'][i,1:-1]
            
        if rho is not None:
            top = top/rho[i,...,1:-1,1]
            bottom = bottom/rho[i,...,1:-1,-2]
            left = left/rho[i,...,1,1:-1]
            right = right/rho[i,...,-2,1:-1]
        
        F[i,...,1:-1,1] += bottom
        F[i,...,1:-1,-2] += top
        F[i,...,1,1:-1] += left
        F[i,...,-2,1:-1] += right

    return F[...,1:-1,1:-1].reshape(list(F[...,1:-1,1:-1].shape[:-2]) + [np.prod(F[...,1:-1,1:-1].shape[-2:])], order = 'C')
 

def cholesky_poisson_solve(rhses, boundaries, h, system_matrix = None, system_matrix_is_decomposed = False):
    '''
    Solves the Poisson equation for the given RHSes.
    
    rhses: tf.Tensor representing the RHS functions of the Poisson equation, defined across the last 2 dimensions
    boundaries: boundary conditions of the outputs; see poisson_RHS documentation
    h: grid spacing of the outputs
    system_matrix: Poisson equation LHS matrix generated by poisson_matrix OR its Cholesky decomposition computed by tf.linalg.cholesky. If the Cholesky decomposition is supplied, set system_matrix_is_decomposed to True.
    system_matrix_is_decomposed: Flag to declare if a provided system_matrix is already Cholesky decomposed. Setting this to True when the system_matrix isn't decomposed will lead to undefined behaviour. Ignored if system_matrix is None.
    
    Outputs a tf.Tensor of identical shape to rhses.
    
    Note: Not tested if this function works on CPU.
    '''
    try: #handle spurious 1 dimensions
        rhses = tf.squeeze(rhses, axis = 1)
    except:
        pass
    boundaries = copy.deepcopy(boundaries)
    for boundary in boundaries.keys():
        try:
            boundaries[boundary] = tf.squeeze(boundaries[boundary], axis = 1)
        except:
            pass
    
    #generate poisson matrix, or use the provided one
    if system_matrix == None:
        system_matrix = tf.cast(poisson_matrix(int(rhses.shape[-2]), int(rhses.shape[-1])), tf.keras.backend.floatx())
        system_matrix_is_decomposed = False
    #import pdb
    #pdb.set_trace()
    if not system_matrix_is_decomposed:
        system_matrix_chol = tf.expand_dims(tf.linalg.cholesky(system_matrix), axis=0)
    else:
        system_matrix_chol = system_matrix
    
    def chol(r):
        return tf.linalg.cholesky_solve(system_matrix_chol, tf.expand_dims(r, axis = 0))
    
    try:
        @tf.contrib.eager.defun
        def chol_solve(rhs_arr):
            return tf.map_fn(chol, rhs_arr)
    except:
        @tf.function
        def chol_solve(rhs_arr):
            return tf.map_fn(chol, rhs_arr)
    
    #put problem into Ax=b format
    try:
        rhs_vectors = tf.expand_dims(tf.transpose(tf.squeeze(poisson_RHS([np.array(rhses), boundaries, h])), (0,1)),axis=2)
    except:
        rhs_vectors = tf.expand_dims(tf.expand_dims(tf.squeeze(poisson_RHS([np.array(rhses), boundaries, h])),axis=1), axis=0)
    
    #solve and reshape
    z = tf.reshape(chol_solve(rhs_vectors), list(rhses.shape[:-2]) + [int(rhses.shape[-1])-2, int(rhses.shape[-2])-2])
    
    soln = np.zeros(rhses.shape, dtype = np.float64)
    soln[...,:,-1] = boundaries['top']
    soln[...,:,0] = boundaries['bottom']
    soln[...,0,:] = boundaries['left']
    soln[...,-1,:] = boundaries['right']
    soln[...,1:-1,1:-1] = z
    
    return tf.expand_dims(soln, axis = 1)

