"""Least squares smooting and interpolation.

   Implements a penalised least-squares fit and associated interpolations.

   The panalty term (or smoothing term) is controlled by the smoothing
   parameter alpha.
   With a value of alpha=0, the fit function will attempt
   to interpolate as closely as possible in the least-squares sense.
   With values alpha > 0, a certain amount of smoothing will be applied.
   A positive alpha is essential in cases where there are too few
   data points.
   A negative alpha is not allowed.
   A typical value of alpha is 1.0e-6
          

   Ole Nielsen, Stephen Roberts, Duncan Gray, Christopher Zoppou
   Geoscience Australia, 2004.   
"""

#FIXME: Current implementation uses full matrices and a general solver.
#Later on we may consider using sparse techniques

from general_mesh import General_Mesh
from Numeric import zeros, array, Float, Int, dot, transpose
from LinearAlgebra import solve_linear_equations

try:
    from util import gradient

except ImportError, e: 
    #FIXME reduce the dependency of modules in pyvolution
    # Have util in a dir, working like load_mesh, and get rid of this
    def gradient(x0, y0, x1, y1, x2, y2, q0, q1, q2):
        """
        """
    
        det = (y2-y0)*(x1-x0) - (y1-y0)*(x2-x0)            
        a = (y2-y0)*(q1-q0) - (y1-y0)*(q2-q0)
        a /= det

        b = (x1-x0)*(q2-q0) - (x2-x0)*(q1-q0)
        b /= det            

        return a, b

ALPHA_INITIAL = 0.001

def fit_to_mesh_file(mesh_file, point_file, mesh_output_file, alpha= ALPHA_INITIAL):
    """
    Given a mesh file (tsh) and a point attribute file (xya), fit
    point attributes to the mesh and write a mesh file with the
    results.
    """
    from load_mesh.loadASCII import mesh_file_to_mesh_dictionary, \
                 load_xya_file, export_trianglulation_file
    # load in the .tsh file
    mesh_dic = mesh_file_to_mesh_dictionary(mesh_file)
    vertex_coordinates = mesh_dic['generatedpointlist']
    triangles = mesh_dic['generatedtrianglelist']
        
    # load in the .xya file
    point_dict = load_xya_file(point_file)
    point_coordinates = point_dict['pointlist']
    point_attributes = point_dict['pointattributelist']
    point_attributes = point_dict['pointattributelist']
        
        
    f = fit_to_mesh(vertex_coordinates,
                    triangles,
                    point_coordinates,
                    point_attributes,
                    alpha = alpha)
    # convert array to list of lists
    mesh_dic['generatedpointattributelist'] = f.tolist()
    export_trianglulation_file(mesh_output_file, mesh_dic)
        

def fit_to_mesh(vertex_coordinates,
                              triangles,
                              point_coordinates,
                              point_attributes,
                              alpha = ALPHA_INITIAL):
    """
    Fit a smooth surface to a trianglulation,
    given data points with attributes.

          
        Inputs:
        
          vertex_coordinates: List of coordinate pairs [xi, eta] of points
          constituting mesh (or a an m x 2 Numeric array)
        
          triangles: List of 3-tuples (or a Numeric array) of
          integers representing indices of all vertices in the mesh.

          point_coordinates: List of coordinate pairs [x, y] of data points
          (or an nx2 Numeric array)

          alpha: Smoothing parameter.

          point_attributes: Vector or array of data at the point_coordinates.
    """
    interp = Interpolation(vertex_coordinates,
                           triangles,
                           point_coordinates,
                           alpha = alpha)
    
    vertex_attributes = interp.fit(point_attributes)
    return vertex_attributes


class Interpolation:

    def __init__(self,
                 vertex_coordinates,
                 triangles,
                 point_coordinates = None,
                 alpha = ALPHA_INITIAL):

        
        """ Build interpolation matrix mapping from
        function values at vertices to function values at data points

        Inputs:
        
          vertex_coordinates: List of coordinate pairs [xi, eta] of points
          constituting mesh (or a an m x 2 Numeric array)
        
          triangles: List of 3-tuples (or a Numeric array) of
          integers representing indices of all vertices in the mesh.

          point_coordinates: List of coordinate pairs [x, y] of data points
          (or an nx2 Numeric array)

          alpha: Smoothing parameter
          
        """


        #Convert input to Numeric arrays
        vertex_coordinates = array(vertex_coordinates).astype(Float)
        triangles = array(triangles).astype(Int)                
        
        #Build underlying mesh
        self.mesh = General_Mesh(vertex_coordinates, triangles)

        #Smoothing parameter
        self.alpha = alpha

        #Build coefficient matrices
        self.build_coefficient_matrix_B(point_coordinates)    


        
    def build_coefficient_matrix_B(self, point_coordinates=None):
        """Build final coefficient matrix
        """

        self.build_smoothing_matrix_D()
        
        if point_coordinates:
            self.build_interpolation_matrix_A(point_coordinates)
            AtA = dot(self.At, self.A)
            self.B = AtA + self.alpha*self.D

        
    def build_interpolation_matrix_A(self, point_coordinates):
        """Build n x m interpolation matrix, where
        n is the number of data points and
        m is the number of basis functions phi_k (one per vertex)
        """ 

        #Convert input to Numeric arrays
        point_coordinates = array(point_coordinates).astype(Float)
        
        #Build n x m interpolation matrix        
        m = self.mesh.coordinates.shape[0] #Nbr of basis functions (1/vertex)
        n = point_coordinates.shape[0]     #Nbr of data points         

        self.A = zeros((n,m), Float)

        #Compute matrix elements
        for i in range(n):
            #For each data_coordinate point

            x = point_coordinates[i]
            for k in range(len(self.mesh)):
                #For each triangle (brute force)
                #FIXME: Real algorithm should only visit relevant triangles

                #Get the three vertex_points
                xi0 = self.mesh.get_vertex_coordinate(k, 0)
                xi1 = self.mesh.get_vertex_coordinate(k, 1)
                xi2 = self.mesh.get_vertex_coordinate(k, 2)                 

                #Get the three normals
                n0 = self.mesh.get_normal(k, 0)
                n1 = self.mesh.get_normal(k, 1)
                n2 = self.mesh.get_normal(k, 2)                

                #Compute interpolation
                sigma2 = dot((x-xi0), n2)/dot((xi2-xi0), n2)
                sigma0 = dot((x-xi1), n0)/dot((xi0-xi1), n0)
                sigma1 = dot((x-xi2), n1)/dot((xi1-xi2), n1)

                #FIXME: Maybe move out to test or something
                epsilon = 1.0e-6
                assert abs(sigma0 + sigma1 + sigma2 - 1.0) < epsilon

                #Check that this triangle contains data point
                if sigma0 >= 0 and sigma1 >= 0 and sigma2 >= 0:

                    #Assign values to matrix A
                    j = self.mesh.triangles[k,0] #Global vertex id
                    self.A[i, j] = sigma0

                    j = self.mesh.triangles[k,1] #Global vertex id
                    self.A[i, j] = sigma1

                    j = self.mesh.triangles[k,2] #Global vertex id
                    self.A[i, j] = sigma2

    
        #Precompute
        self.At = transpose(self.A)


        #FIXME: Remember to re-introduce the 1/n factor in the
        #interpolation term
        
    def build_smoothing_matrix_D(self):
        """Build m x m smoothing matrix, where
        m is the number of basis functions phi_k (one per vertex)

        The smoothing matrix is defined as

        D = D1 + D2

        where

        [D1]_{k,l} = \int_\Omega
           \frac{\partial \phi_k}{\partial x}
           \frac{\partial \phi_l}{\partial x}\,
           dx dy

        [D2]_{k,l} = \int_\Omega
           \frac{\partial \phi_k}{\partial y}
           \frac{\partial \phi_l}{\partial y}\,
           dx dy


        The derivatives \frac{\partial \phi_k}{\partial x},
        \frac{\partial \phi_k}{\partial x} for a particular triangle
        are obtained by computing the gradient a_k, b_k for basis function k
        """

        #FIXME: algorithm might be optimised by computing local 9x9
        #"element stiffness matrices:

        
        m = self.mesh.coordinates.shape[0] #Nbr of basis functions (1/vertex)

        self.D = zeros((m,m), Float)

        #For each triangle compute contributions to D = D1+D2        
        for i in range(len(self.mesh)):

            #Get area
            area = self.mesh.areas[i]

            #Get global vertex indices
            v0 = self.mesh.triangles[i,0]
            v1 = self.mesh.triangles[i,1]
            v2 = self.mesh.triangles[i,2]
            
            #Get the three vertex_points
            xi0 = self.mesh.get_vertex_coordinate(i, 0)
            xi1 = self.mesh.get_vertex_coordinate(i, 1)
            xi2 = self.mesh.get_vertex_coordinate(i, 2)                 

            #Compute gradients for each vertex
            a0, b0 = gradient(xi0[0], xi0[1], xi1[0], xi1[1], xi2[0], xi2[1],
                              1, 0, 0)

            a1, b1 = gradient(xi0[0], xi0[1], xi1[0], xi1[1], xi2[0], xi2[1],
                              0, 1, 0)

            a2, b2 = gradient(xi0[0], xi0[1], xi1[0], xi1[1], xi2[0], xi2[1],
                              0, 0, 1)            

            #Compute diagonal contributions
            self.D[v0,v0] += (a0*a0 + b0*b0)*area
            self.D[v1,v1] += (a1*a1 + b1*b1)*area
            self.D[v2,v2] += (a2*a2 + b2*b2)*area            

            #Compute contributions for basis functions sharing edges
            e01 = (a0*a1 + b0*b1)*area
            self.D[v0,v1] += e01
            self.D[v1,v0] += e01

            e12 = (a1*a2 + b1*b2)*area
            self.D[v1,v2] += e12
            self.D[v2,v1] += e12

            e20 = (a2*a0 + b2*b0)*area
            self.D[v2,v0] += e20
            self.D[v0,v2] += e20             
            
            
    def fit(self, z):
        """Fit a smooth surface to given data points z.

        The smooth surface is computed at each vertex in the underlying
        mesh using the formula given in the module doc string. 

        Pre Condition:
          self.A, self.At and self.B have been initialised
          
        Inputs:
          z: Vector or array of data at the point_coordinates.
          If z is an array, smoothing will be done for each column
        """

        #Convert input to Numeric arrays
        z = array(z).astype(Float)
        
        #Compute right hand side based on data
        Atz = dot(self.At, z)

        #Check sanity
        n, m = self.A.shape
        if n<m and self.alpha == 0.0:
            msg = 'ERROR (least_squares): Too few data points\n'
            msg += 'There only %d data points. Need at least %d\n' %(n,m)
            msg += 'Alternatively, increase smoothing parameter alpha' 
            raise msg


        #Solve and return
        return solve_linear_equations(self.B, Atz)

        #FIXME: Should we store the result here for later use? (ON)
    

    def interpolate(self, f):
        """Compute predicted values at data points implied in self.A. 

        The mesh values representing a smooth surface are
        assumed to be specified in f. This argument could,
        for example have been obtained from the method self.fit()
        
        Pre Condition:
          self.A has been initialised
        
        Inputs:
          f: Vector or array of data at the mesh vertices.
          If f is an array, interpolation will be done for each column
        """
        
        return dot(self.A, f)

     
    #FIXME: We will need a method 'evaluate(self):' that will interpolate
    #a computed surface living on the mesh onto a collection of
    #arbitrary data points
    #
    #Precondition: self.fit(z) has stored its result in self.f.
    #
    #Input: data_points
    #Algorithm:
    #  1 Build a new temporary A matrix based on mesh and new data points
    #  2 Apply it to self.f (return A*self.f)
    #
    # ON



#-------------------------------------------------------------
if __name__ == "__main__":
    """
    Load in a mesh and data points with attributes.
    Fit the attributes to the mesh.
    Save a new mesh file.
    """
    import os, sys
    usage = "usage: %s mesh_input.tsh point.xya  mesh_output.tsh alpha" %         os.path.basename(sys.argv[0])

    if len(sys.argv) < 4:
        print usage
    else:
        mesh_file = sys.argv[1]
        point_file = sys.argv[2]
        mesh_output_file = sys.argv[3]
        if len(sys.argv) > 4:
            alpha = sys.argv[4]
        else:
            alpha = ALPHA_INITIAL
        fit_to_mesh_file(mesh_file, point_file, mesh_output_file, alpha)