import unittest
from clean_data import *
from Assignment6_functions import *
from cryptorandom.cryptorandom import SHA256
import pandas as pd 
import numpy as np

class utest(unittest.TestCase):
    """ 

    Unit test class that test assignment6's functions

    """
	
    def test_len_filenames(self):
        """
        test if the get_filenames function return the correct number of file names.

        """
        self.assertEqual(len(get_filenames()), 30)
        self.assertEqual(len(get_filenames(1980, 1981)), 2)

    def test_filenames(self):
        """
        test if the get_filenames function return the correct file names (1980 to 2009)
        """
        filenames = ['ucr_offenses_known_monthly_2004.dta',
                     'ucr_offenses_known_monthly_2005.dta',
                     'ucr_offenses_known_monthly_2007.dta',
                     'ucr_offenses_known_monthly_2006.dta',
                     'ucr_offenses_known_monthly_2002.dta',
                     'ucr_offenses_known_monthly_2003.dta',
                     'ucr_offenses_known_monthly_2001.dta',
                     'ucr_offenses_known_monthly_2000.dta',
                     'ucr_offenses_known_monthly_1988.dta',
                     'ucr_offenses_known_monthly_1989.dta',
                     'ucr_offenses_known_monthly_1999.dta',
                     'ucr_offenses_known_monthly_1998.dta',
                     'ucr_offenses_known_monthly_1981.dta',
                     'ucr_offenses_known_monthly_1995.dta',
                     'ucr_offenses_known_monthly_1994.dta',
                     'ucr_offenses_known_monthly_1980.dta',
                     'ucr_offenses_known_monthly_1996.dta',
                     'ucr_offenses_known_monthly_1982.dta',
                     'ucr_offenses_known_monthly_1983.dta',
                     'ucr_offenses_known_monthly_1997.dta',
                     'ucr_offenses_known_monthly_1993.dta',
                     'ucr_offenses_known_monthly_1987.dta',
                     'ucr_offenses_known_monthly_1986.dta',
                     'ucr_offenses_known_monthly_1992.dta',
                     'ucr_offenses_known_monthly_1984.dta',
                     'ucr_offenses_known_monthly_1990.dta',
                     'ucr_offenses_known_monthly_1991.dta',
                     'ucr_offenses_known_monthly_1985.dta',
                     'ucr_offenses_known_monthly_2008.dta',
                     'ucr_offenses_known_monthly_2009.dta']
        self.assertEqual(get_filenames(), filenames)
        self.assertEqual(get_filenames(2004, 2007), filenames[0:4])

    def test_all_alameda_data(self):
        """
        test if the get_clean_data function calculate the sum of crime correctly
        """
        df = get_clean_data(start_year = 1980, end_year = 1981, )
        self.assertEqual(df.shape, (24, 4))
        self.assertEqual(df.loc[0, 'crime_sum'], 9980)
        self.assertEqual(df.loc[12, 'crime_sum'], 10375)

    def test_west_alameda(self):
        """
        test if, for west alameda, the get_clean_data function calculate the sum of crime correctly
        """
        df = get_clean_data(start_year = 1980, end_year = 1981, selected_area = "west")
        self.assertEqual(df.shape, (24, 4))
        self.assertEqual(df.loc[0, 'crime_sum'], 9479)
        self.assertEqual(df.loc[11, 'crime_sum'], 9838)

    def test_east_alameda(self):
        """
        test if, for east alameda, the get_clean_data function calculate the sum of crime correctly
        """
        df = get_clean_data(start_year = 1980, end_year = 1981, selected_area = "east")
        self.assertEqual(df.shape, (24, 4))
        self.assertEqual(df.loc[0, 'crime_sum'], 501)
        self.assertEqual(df.loc[11, 'crime_sum'], 471)


    def test_get_temp_features(self):
        """
            This function is to test get_temp_features function
            First, read a max temparature in Alameda csv file to create a pandas data frame.
            Then test the elements of data frame generated by get_temp_feature
            """
        data1 = pd.read_csv('../group_assignment3/TMAX_data.csv')
        
        self.assertEqual(get_temp_features(data1).iloc[0,2],  34.0)
        self.assertEqual(get_temp_features(data1).iloc[1,3],  41.0)
        self.assertEqual(get_temp_features(data1).iloc[4,6],  1.0)

    def test_get_prcp_features_whole(self):
        """
            This function is to test get_prcp_features function
            First, read a prcp in Alameda csv file to create a pandas data frame which have five columns.
            Then test the elements of data frame generated by get_prcp_feature
        """
        data2 = pd.read_csv("../group_assignment3/PRCP_data.csv")
        self.assertEqual(get_prcp_features_whole(data2).iloc[0,0],29.0)
        self.assertEqual(get_prcp_features_whole(data2).iloc[2,1],22.0)

    def test_get_prcp_features(self):
        """
            This function is to test get_prcp_features function
            First, read a prcp in Alameda csv file to create a pandas data frame which have 3 columns.
            Then test the elements of data frame generated by get_prcp_feature
        """
        data2 = pd.read_csv("../group_assignment3/PRCP_data.csv")
        self.assertEqual(get_prcp_features(data2).iloc[0,0],29.0)
        self.assertEqual(get_prcp_features(data2).iloc[2,1],22.0)
        self.assertEqual(get_prcp_features(data2).iloc[2,2],4.0)

    def test_get_dummy_features(self):
        """
            This function is to test get_dummy_features function
            First, read a  max temparature in Alameda csv file to create a pandas data frame.
            Then test elements in both phi and theta dummy varibles
            """
        
        data3 = pd.read_csv('../group_assignment3/TMAX_data.csv')
        
        self.assertEqual(get_dummy_features(data3)[0].iloc[6,0],  1.0)
        self.assertEqual(get_dummy_features(data3)[0].iloc[10,10],  0)
        self.assertEqual(get_dummy_features(data3)[0].iloc[358,29],  1)
        self.assertEqual(get_dummy_features(data3)[1].iloc[0,1],  1.0)
        self.assertEqual(get_dummy_features(data3)[1].iloc[9,9],  0.0)
        self.assertEqual(get_dummy_features(data3)[1].iloc[355,8],  1.0)

    def test_get_permute(self):
        """
            This function is to test get_permute function. For each k, the length of return value should
            be less and equal than k.
        
        """
        r = SHA256(seed=123456)
        self.assertEqual(len(get_permute(5,r=r)), 5)
        self.assertEqual(len(get_permute(360,r=r)), 360)
        self.assertEqual(len(get_permute(77,r=r)), 77)
        r = SHA256(seed=123456)
        self.assertEqual(get_permute(256,r=r)[:5], [True, True, False, False, False])
    
    def test_rms(self):
        """
            This fucntion test rms function. First to generate three panda data frames
            Then test the rms of these three data frames.
            """
        data1 = [1,2,3,4,5]
        df1 = pd.DataFrame(data1)
        data2 = [8,4,3,7,5]
        df2 = pd.DataFrame(data2)
        data3= [9,10,8,7,9]
        df3 = pd.DataFrame(data3)
        
        self.assertEqual(rms(df1,df2)[0], 3.521363372331802)
        self.assertEqual(rms(df1,df3)[0], 5.966573556070519)
        self.assertEqual(rms(df2,df3)[0], 3.9496835316262997)
            
if __name__ == '__main__':
    unittest.main()
