# -*- coding: utf-8 -*-
"""
Created on Mon May  26 11:33:12 2014

@author: emmanuel
"""

import unittest

import numpy as np
import scipy as sp

from connectome.pes.weak_path_enumeration_sampling import weak_edge_pes


class TestWeakPES(unittest.TestCase):
    def test_weak_pes_triangle(self):
        triangle  = sp.array( [ [0 , 0 , 1],
                                [1 , 0 , 0],
                                [0 , 1 , 0]  ]  )
        expected_results = sp.array( [[0, 0, 0],
                                      [3, 0, 0],
                                      [0, 0, 3],] ,dtype=np.uint64)                                
        results = weak_edge_pes(triangle,
                                max_length_py=triangle.shape[0],
                                probabilities_py=sp.ones(triangle.shape[0]) * 2,
                                repeat=1)
        self.assertTrue( (expected_results == results).all() )

        
    def test_weak_pes_triangle_isolated(self):
        triangle  = sp.array( [ [0 , 0 , 1, 0],
                                [1 , 0 , 0, 0],
                                [0 , 1 , 0, 0],
                                [0 , 0 , 0, 0]  ]  )
        expected_results = sp.array( [[0, 0, 0],
                                      [3, 0, 0],
                                      [0, 0, 3],
                                      [0, 0, 0]] ,dtype=np.uint64)                                
        results = weak_edge_pes(triangle,
                                max_length_py=triangle.shape[0],
                                probabilities_py=sp.ones(triangle.shape[0]) * 2,
                                repeat=1)
        self.assertTrue( (expected_results == results).all() )
 
        
        
    def test_weak_pes_triangle_and_input(self):
        triangle  = sp.array( [ [0 , 0 , 1, 1],
                                [1 , 0 , 0, 0],
                                [0 , 1 , 0, 0],
                                [0 , 0 , 0, 0]  ]  )
        expected_results = sp.array( [[0, 0, 0],
                                      [4, 0, 0],
                                      [2, 0, 3],
                                      [2, 0, 0]] ,dtype=np.uint64)                                
        results = weak_edge_pes(triangle,
                                max_length_py=triangle.shape[0],
                                probabilities_py=sp.ones(triangle.shape[0]) * 2,
                                repeat=1)
        self.assertTrue( (expected_results == results).all() )
       
        
        
    def test_weak_pes_rectangle(self):
        triangle  = sp.array( [ [0 , 0 , 0, 1],
                                [1 , 0 , 0, 0],
                                [0 , 1 , 0, 0],
                                [0 , 0 , 1, 0]  ]  )
        expected_results = sp.array( [[0, 0, 0],
                                      [4, 0, 0],
                                      [4, 0, 0],
                                      [0, 0, 4]] ,dtype=np.uint64)                                
        results = weak_edge_pes(triangle,
                                max_length_py=triangle.shape[0],
                                probabilities_py=sp.ones(triangle.shape[0]) * 2,
                                repeat=1)
        self.assertTrue( (expected_results == results).all() )
        

        
    

if __name__ == "__main__":
    unittest.main()
