#!/usr/bin/env python

#TEST
import sys
import unittest
from math import sqrt


from interpolate import *
from Numeric import allclose, array, transpose

from coordinate_transforms.geo_reference import Geo_reference

def distance(x, y):
    return sqrt( sum( (array(x)-array(y))**2 ))

def linear_function(point):
    point = array(point)
    return point[:,0]+point[:,1]


class Test_Interpolate(unittest.TestCase):

    def setUp(self):
        pass

    def tearDown(self):
        pass

    def test_datapoint_at_centroid(self):
        a = [0.0, 0.0]
        b = [0.0, 2.0]
        c = [2.0,0.0]
        points = [a, b, c]
        vertices = [ [1,0,2] ]   #bac

        data = [ [2.0/3, 2.0/3] ] #Use centroid as one data point

        interp = Interpolate(points, vertices)
        assert allclose(interp._build_interpolation_matrix_A(data).todense(),
                        [[1./3, 1./3, 1./3]])


    def test_quad_tree(self):
        p0 = [-10.0, -10.0]
        p1 = [20.0, -10.0]
        p2 = [-10.0, 20.0]
        p3 = [10.0, 50.0]
        p4 = [30.0, 30.0]
        p5 = [50.0, 10.0]
        p6 = [40.0, 60.0]
        p7 = [60.0, 40.0]
        p8 = [-66.0, 20.0]
        p9 = [10.0, -66.0]

        points = [p0, p1, p2, p3, p4, p5, p6, p7, p8, p9]
        triangles = [ [0, 1, 2],
                      [3, 2, 4],
                      [4, 2, 1],
                      [4, 1, 5],
                      [3, 4, 6],
                      [6, 4, 7],
                      [7, 4, 5],
                      [8, 0, 2],
                      [0, 9, 1]]

        data = [ [4,4] ]
        interp = Interpolate(points, triangles,
                               max_vertices_per_cell = 4)
        #print "PDSG - interp.get_A()", interp.get_A()
        answer =  [ [ 0.06666667,  0.46666667,  0.46666667,  0.,
                      0., 0. , 0., 0., 0., 0.]]

        
        assert allclose(interp._build_interpolation_matrix_A(data).todense(),
                        answer)
        #interp.set_point_coordinates([[-30, -30]]) #point outside of mesh
        #print "PDSG - interp.get_A()", interp.get_A()
        data = [[-30, -30]]
        answer =  [ [ 0.0,  0.0,  0.0,  0.,
                      0., 0. , 0., 0., 0., 0.]]
        
        assert allclose(interp._build_interpolation_matrix_A(data).todense(),
                        answer)


        #point outside of quad tree root cell
        #interp.set_point_coordinates([[-70, -70]])
        #print "PDSG - interp.get_A()", interp.get_A()
        data = [[-70, -70]]
        answer =  [ [ 0.0,  0.0,  0.0,  0.,
                      0., 0. , 0., 0., 0., 0.]]
        assert allclose(interp._build_interpolation_matrix_A(data).todense(),
                        answer)

    def test_datapoints_at_vertices(self):
        """Test that data points coinciding with vertices yield a diagonal matrix
        """

        a = [0.0, 0.0]
        b = [0.0, 2.0]
        c = [2.0,0.0]
        points = [a, b, c]
        vertices = [ [1,0,2] ]   #bac

        data = points #Use data at vertices

        interp = Interpolate(points, vertices)
        answer = [[1., 0., 0.],
                   [0., 1., 0.],
                   [0., 0., 1.]]
        assert allclose(interp._build_interpolation_matrix_A(data).todense(),
                        answer)


    def test_datapoints_on_edge_midpoints(self):
        """Try datapoints midway on edges -
        each point should affect two matrix entries equally
        """

        a = [0.0, 0.0]
        b = [0.0, 2.0]
        c = [2.0,0.0]
        points = [a, b, c]
        vertices = [ [1,0,2] ]   #bac

        data = [ [0., 1.], [1., 0.], [1., 1.] ]
        answer =  [[0.5, 0.5, 0.0],  #Affects vertex 1 and 0
                    [0.5, 0.0, 0.5],  #Affects vertex 0 and 2
                    [0.0, 0.5, 0.5]]
        interp = Interpolate(points, vertices, data)

        assert allclose(interp._build_interpolation_matrix_A(data).todense(),
                        answer)

    def test_datapoints_on_edges(self):
        """Try datapoints on edges -
        each point should affect two matrix entries in proportion
        """

        a = [0.0, 0.0]
        b = [0.0, 2.0]
        c = [2.0,0.0]
        points = [a, b, c]
        vertices = [ [1,0,2] ]   #bac

        data = [ [0., 1.5], [1.5, 0.], [1.5, 0.5] ]
        answer =  [[0.25, 0.75, 0.0],  #Affects vertex 1 and 0
                   [0.25, 0.0, 0.75],  #Affects vertex 0 and 2
                   [0.0, 0.25, 0.75]]

        interp = Interpolate(points, vertices, data)

        assert allclose(interp._build_interpolation_matrix_A(data).todense(),
                        answer)


    def test_arbitrary_datapoints(self):
        """Try arbitrary datapoints
        """

        from Numeric import sum

        a = [0.0, 0.0]
        b = [0.0, 2.0]
        c = [2.0,0.0]
        points = [a, b, c]
        vertices = [ [1,0,2] ]   #bac

        data = [ [0.2, 1.5], [0.123, 1.768], [1.43, 0.44] ]

        interp = Interpolate(points, vertices, data)
        #print "interp.get_A()", interp.get_A()
        results = interp._build_interpolation_matrix_A(data).todense()
        assert allclose(sum(results, axis=1), 1.0)

        #FIXME - have to change this test to check default info
    def NO_test_arbitrary_datapoints_some_outside(self):
        """Try arbitrary datapoints one outside the triangle.
        That one should be ignored
        """

        from Numeric import sum

        a = [0.0, 0.0]
        b = [0.0, 2.0]
        c = [2.0,0.0]
        points = [a, b, c]
        vertices = [ [1,0,2] ]   #bac

        data = [ [0.2, 1.5], [0.123, 1.768], [1.43, 0.44], [5.0, 7.0]]


        interp = Interpolate(points, vertices, data, precrop = True)
        
        results = interp._build_interpolation_matrix_A(data).todense()
        assert allclose(sum(results, axis=1), 1.0)

        interp = Interpolate(points, vertices, data, precrop = False)
        results = interp._build_interpolation_matrix_A(data).todense()
        assert allclose(sum(results, axis=1), [1,1,1,0])



    # this causes a memory error in scipy.sparse
    def test_more_triangles(self):

        a = [-1.0, 0.0]
        b = [3.0, 4.0]
        c = [4.0,1.0]
        d = [-3.0, 2.0] #3
        e = [-1.0,-2.0]
        f = [1.0, -2.0] #5

        points = [a, b, c, d,e,f]
        triangles = [[0,1,3],[1,0,2],[0,4,5], [0,5,2]] #abd bac aef afc

        #Data points
        data = [ [-3., 2.0], [-2, 1], [0.0, 1], [0, 3], [2, 3], [-1.0/3,-4./3] ]
        interp = Interpolate(points, triangles)

        answer = [[0.0, 0.0, 0.0, 1.0, 0.0, 0.0],    #Affects point d
                  [0.5, 0.0, 0.0, 0.5, 0.0, 0.0],    #Affects points a and d
                  [0.75, 0.25, 0.0, 0.0, 0.0, 0.0],  #Affects points a and b
                  [0.0, 0.5, 0.0, 0.5, 0.0, 0.0],    #Affects points a and d
                  [0.25, 0.75, 0.0, 0.0, 0.0, 0.0],  #Affects points a and b
                  [1./3, 0.0, 0.0, 0.0, 1./3, 1./3]] #Affects points a, e and f


        A = interp._build_interpolation_matrix_A(data).todense()
        for i in range(A.shape[0]):
            for j in range(A.shape[1]):
                if not allclose(A[i,j], answer[i][j]):
                    print i,j,':',A[i,j], answer[i][j]


        #results = interp._build_interpolation_matrix_A(data).todense()

        assert allclose(A, answer)


    def test_interpolate_attributes_to_points(self):
        v0 = [0.0, 0.0]
        v1 = [0.0, 5.0]
        v2 = [5.0, 0.0]

        vertices = [v0, v1, v2]
        triangles = [ [1,0,2] ]   #bac

        d0 = [1.0, 1.0]
        d1 = [1.0, 2.0]
        d2 = [3.0, 1.0]
        point_coords = [ d0, d1, d2]

        interp = Interpolate(vertices, triangles, point_coords)
        f = linear_function(vertices)
        z = interp.interpolate(f, point_coords)
        answer = linear_function(point_coords)


        assert allclose(z, answer)



    def test_interpolate_attributes_to_pointsII(self):
        a = [-1.0, 0.0]
        b = [3.0, 4.0]
        c = [4.0, 1.0]
        d = [-3.0, 2.0] #3
        e = [-1.0, -2.0]
        f = [1.0, -2.0] #5

        vertices = [a, b, c, d,e,f]
        triangles = [[0,1,3], [1,0,2], [0,4,5], [0,5,2]] #abd bac aef afc


        point_coords = [[-2.0, 2.0],
                        [-1.0, 1.0],
                        [0.0, 2.0],
                        [1.0, 1.0],
                        [2.0, 1.0],
                        [0.0, 0.0],
                        [1.0, 0.0],
                        [0.0, -1.0],
                        [-0.2, -0.5],
                        [-0.9, -1.5],
                        [0.5, -1.9],
                        [3.0, 1.0]]

        interp = Interpolate(vertices, triangles)
        f = linear_function(vertices)
        z = interp.interpolate(f, point_coords)
        answer = linear_function(point_coords)
        #print "z",z
        #print "answer",answer
        assert allclose(z, answer)

    def test_interpolate_attributes_to_pointsIII(self):
        """Test linear interpolation of known values at vertices to
        new points inside a triangle
        """
        a = [0.0, 0.0]
        b = [0.0, 5.0]
        c = [5.0, 0.0]
        d = [5.0, 5.0]

        vertices = [a, b, c, d]
        triangles = [ [1,0,2], [2,3,1] ]   #bac, cdb

        #Points within triangle 1
        d0 = [1.0, 1.0]
        d1 = [1.0, 2.0]
        d2 = [3.0, 1.0]

        #Point within triangle 2
        d3 = [4.0, 3.0]

        #Points on common edge
        d4 = [2.5, 2.5]
        d5 = [4.0, 1.0]

        #Point on common vertex
        d6 = [0., 5.]
        
        point_coords = [d0, d1, d2, d3, d4, d5, d6]

        interp = Interpolate(vertices, triangles)

        #Known values at vertices
        #Functions are x+y, x+2y, 2x+y, x-y-5
        f = [ [0., 0., 0., -5.],        # (0,0)
              [5., 10., 5., -10.],      # (0,5)
              [5., 5., 10.0, 0.],       # (5,0)
              [10., 15., 15., -5.]]     # (5,5)

        z = interp.interpolate(f, point_coords)
        answer = [ [2., 3., 3., -5.],   # (1,1)
                   [3., 5., 4., -6.],   # (1,2)
                   [4., 5., 7., -3.],   # (3,1)
                   [7., 10., 11., -4.], # (4,3)
                   [5., 7.5, 7.5, -5.], # (2.5, 2.5)
                   [5., 6., 9., -2.],   # (4,1)
                   [5., 10., 5., -10.]]  # (0,5)

        #print "***********"
        #print "z",z
        #print "answer",answer
        #print "***********"

        #Should an error message be returned if points are outside
        # of the mesh? Not currently.  

        assert allclose(z, answer)


    def test_interpolate_point_outside_of_mesh(self):
        """Test linear interpolation of known values at vertices to
        new points inside a triangle
        """
        a = [0.0, 0.0]
        b = [0.0, 5.0]
        c = [5.0, 0.0]
        d = [5.0, 5.0]

        vertices = [a, b, c, d]
        triangles = [ [1,0,2], [2,3,1] ]   #bac, cdb

        #Far away point
        d7 = [-1., -1.]
        
        point_coords = [ d7]
        interp = Interpolate(vertices, triangles)

        #Known values at vertices
        #Functions are x+y, x+2y, 2x+y, x-y-5
        f = [ [0., 0., 0., -5.],        # (0,0)
              [5., 10., 5., -10.],      # (0,5)
              [5., 5., 10.0, 0.],       # (5,0)
              [10., 15., 15., -5.]]     # (5,5)

        z = interp.interpolate(f, point_coords)
        answer = [ [0., 0., 0., 0.]] # (-1,-1)

        #print "***********"
        #print "z",z
        #print "answer",answer
        #print "***********"

        #Should an error message be returned if points are outside
        # of the mesh? Not currently.  

        assert allclose(z, answer)

    def test_interpolate_attributes_to_pointsIV(self):
        a = [-1.0, 0.0]
        b = [3.0, 4.0]
        c = [4.0, 1.0]
        d = [-3.0, 2.0] #3
        e = [-1.0, -2.0]
        f = [1.0, -2.0] #5

        vertices = [a, b, c, d,e,f]
        triangles = [[0,1,3], [1,0,2], [0,4,5], [0,5,2]] #abd bac aef afc


        point_coords = [[-2.0, 2.0],
                        [-1.0, 1.0],
                        [0.0, 2.0],
                        [1.0, 1.0],
                        [2.0, 1.0],
                        [0.0, 0.0],
                        [1.0, 0.0],
                        [0.0, -1.0],
                        [-0.2, -0.5],
                        [-0.9, -1.5],
                        [0.5, -1.9],
                        [3.0, 1.0]]

        interp = Interpolate(vertices, triangles)
        f = array([linear_function(vertices),2*linear_function(vertices) ])
        f = transpose(f)
        #print "f",f
        z = interp.interpolate(f, point_coords)
        answer = [linear_function(point_coords),
                  2*linear_function(point_coords) ]
        answer = transpose(answer)
        #print "z",z
        #print "answer",answer
        assert allclose(z, answer)


    def test_interpolate_blocking(self):
        a = [-1.0, 0.0]
        b = [3.0, 4.0]
        c = [4.0, 1.0]
        d = [-3.0, 2.0] #3
        e = [-1.0, -2.0]
        f = [1.0, -2.0] #5

        vertices = [a, b, c, d,e,f]
        triangles = [[0,1,3], [1,0,2], [0,4,5], [0,5,2]] #abd bac aef afc


        point_coords = [[-2.0, 2.0],
                        [-1.0, 1.0],
                        [0.0, 2.0],
                        [1.0, 1.0],
                        [2.0, 1.0],
                        [0.0, 0.0],
                        [1.0, 0.0],
                        [0.0, -1.0],
                        [-0.2, -0.5],
                        [-0.9, -1.5],
                        [0.5, -1.9],
                        [3.0, 1.0]]

        interp = Interpolate(vertices, triangles)
        f = array([linear_function(vertices),2*linear_function(vertices) ])
        f = transpose(f)
        #print "f",f
        for blocking_max in range(len(point_coords)+2):
        #if True:
         #   blocking_max = 5
            z = interp.interpolate(f, point_coords,
                                   start_blocking_len=blocking_max)
            answer = [linear_function(point_coords),
                      2*linear_function(point_coords) ]
            answer = transpose(answer)
            #print "z",z
            #print "answer",answer
            assert allclose(z, answer)

    def test_interpolate_reuse(self):
        a = [-1.0, 0.0]
        b = [3.0, 4.0]
        c = [4.0, 1.0]
        d = [-3.0, 2.0] #3
        e = [-1.0, -2.0]
        f = [1.0, -2.0] #5

        vertices = [a, b, c, d,e,f]
        triangles = [[0,1,3], [1,0,2], [0,4,5], [0,5,2]] #abd bac aef afc


        point_coords = [[-2.0, 2.0],
                        [-1.0, 1.0],
                        [0.0, 2.0],
                        [1.0, 1.0],
                        [2.0, 1.0],
                        [0.0, 0.0],
                        [1.0, 0.0],
                        [0.0, -1.0],
                        [-0.2, -0.5],
                        [-0.9, -1.5],
                        [0.5, -1.9],
                        [3.0, 1.0]]

        interp = Interpolate(vertices, triangles)
        f = array([linear_function(vertices),2*linear_function(vertices) ])
        f = transpose(f)
        z = interp.interpolate(f, point_coords,
                               start_blocking_len=20)
        answer = [linear_function(point_coords),
                  2*linear_function(point_coords) ]
        answer = transpose(answer)
        #print "z",z
        #print "answer",answer
        assert allclose(z, answer)
        assert allclose(interp._A_can_be_reused, True)

        z = interp.interpolate(f)
        assert allclose(z, answer)
        
        # This causes blocking to occur. 
        z = interp.interpolate(f, start_blocking_len=10)
        assert allclose(z, answer)
        assert allclose(interp._A_can_be_reused, False)

        #A is recalculated
        z = interp.interpolate(f)
        assert allclose(z, answer)
        assert allclose(interp._A_can_be_reused, True)
        
        interp = Interpolate(vertices, triangles)
        #Must raise an exception, no points specified
        try:
            z = interp.interpolate(f)
        except:
            pass
        

#-------------------------------------------------------------
if __name__ == "__main__":
    suite = unittest.makeSuite(Test_Interpolate,'test')
    runner = unittest.TextTestRunner(verbosity=1)
    runner.run(suite)