import numpy as np
import scipy as sp
import logging
import doctest

from pysnptools.pysnptools.snpreader.bed import *
from pysnptools.pysnptools.snpreader.hdf5 import Hdf5
from pysnptools.pysnptools.snpreader.dat import Dat
from pysnptools.pysnptools.snpreader.ped import Ped
from pysnptools.pysnptools.standardizer.unit import Unit
from pysnptools.pysnptools.standardizer.beta import Beta


import unittest
import os.path
import time



   
class TestLoader(unittest.TestCase):     

    def xtest_aaa_hdf5_speed(self): #!!ck

        #currentFolder + "/examples/toydata"
        #currentFolder + "/examples/delme.hdf5"
        bedFileName = r"d:\data\carlk\cachebio\genetics\synthetic\wtccclikeH\snps" #!!ck local paths
        hdf5Pattern = r"d:\data\carlk\cachebio\genetics\synthetic\wtccclikeH\del.{0}.hdf5"#!!ck

        tt0 = time.time()
        snpreader_bed = Bed(bedFileName)

        S0 = snpreader_bed.sid_count
        snp_index_list0 = range(min(S0,15000)) 

        hdf5FileName = hdf5Pattern.format(len(snp_index_list0))

        #!ck        snpDataBed = snpreader_bed.read(SnpIndexList(snp_index_list0))
        tt1 = time.time()
        logging.info("Read bed %.2f seconds" % (tt1 - tt0))
        #!ck        Hdf5.write(snpDataBed, hdf5FileName)
        tt2 = time.time()
        logging.info("write Hdf5 bed %.2f seconds" % (tt2 - tt1))

        snpreader_hdf5 = Hdf5(hdf5FileName)
        assert(snpreader_hdf5.iid_count == snpreader_bed.iid_count)
        S = snpreader_hdf5.sid_count
        N_original = snpreader_hdf5.iid_count
        iid_index_list = sorted(range(N_original - 1,0,-2))

        snp_index_list = sorted(range(S - 1,0,-2))#!ck
        #snp_index_list = range(S/2)

        snpreader_hdf5 = snpreader_hdf5[iid_index_list,:]
        snpDataHdf5 = snpreader_hdf5[:,snp_index_list].read()
        tt3 = time.time()
        logging.info("read Hdf5 with reversed indexes bed %.2f seconds" % (tt3 - tt2))

        snpDataHdf5C = snpreader_hdf5[:,snp_index_list].read(order = "C")
        tt4 = time.time()
        logging.info("read Hdf5 C with reversed indexes bed %.2f seconds" % (tt4 - tt3))

        print "the end"
    

    @classmethod
    def setUpClass(self):
        self.currentFolder = os.path.dirname(os.path.realpath(__file__))
        #TODO: get data set with NANs!
        snpreader = Bed(self.currentFolder + "/examples/toydata")
        self.pheno_fn = self.currentFolder + "/examples/toydata.phe"
        self.snps = snpreader.read(order='F',force_python_only=True).val

    def test_c_reader_bed(self):
        snpreader = Bed(self.currentFolder + "/examples/toydata")
        self.c_reader(snpreader)

    def test_c_reader_hdf5(self):
        snpreader = Hdf5(self.currentFolder + "/examples/toydata.snpmajor.hdf5")
        self.c_reader(snpreader)

    def test_c_reader_dat(self):
        snpreader = Dat(self.currentFolder + "/examples/toydata.dat")
        self.c_reader(snpreader)

    def c_reader(self,snpreader):
        """
        make sure c-reader yields same result
        """
        snp_c = snpreader.read(order='F',force_python_only=False).val
        
        self.assertEqual(np.float64, snp_c.dtype)
        self.assertTrue(np.allclose(self.snps, snp_c, rtol=1e-05, atol=1e-05))

    def test_standardize_bed(self):
        snpreader = Bed(self.currentFolder + "/examples/toydata")
        self.standardize(snpreader)

    def test_standardize_hdf5(self):
        snpreader = Hdf5(self.currentFolder + "/examples/toydata.iidmajor.hdf5")
        self.standardize(snpreader)

    def test_standardize_dat(self):
        snpreader = Dat(self.currentFolder + "/examples/toydata.dat")
        self.standardize(snpreader)

    def test_standardize_ped(self):
        snpreader = Ped(self.currentFolder + "/examples/toydata")
        self.standardize(snpreader)


    def standardize(self,snpreader):
        """
        make sure blocked standardize yields same result as regular standardize
        """

        for dtype in [sp.float64,sp.float32]:

            snps = snpreader.read(order='F',force_python_only=True,dtype=dtype).val
            self.assertEqual(dtype, snps.dtype)

            snp_s1 = Unit().standardize(snps.copy(), force_python_only=True)
            snp_s2 = Unit().standardize(snps.copy(), blocksize=100, force_python_only=True)
            snps_F = np.array(snps, dtype=dtype, order="F")
            snp_s3 = Unit().standardize(snps_F)
            snps_C = np.array(snps, dtype=dtype, order="C")
            snp_s4 = Unit().standardize(snps_C)

            snp_beta1 = Beta(1, 25).standardize(snps.copy(), force_python_only=True)
            snps_F = np.array(snps, dtype=dtype, order="F")
            snp_beta2 = Beta(1, 25).standardize(snps_F)
            snps_C = np.array(snps, dtype=dtype, order="C")
            snp_beta3 = Beta(1, 25).standardize(snps_C)

            self.assertEqual(snp_s1.shape[0], snp_s2.shape[0])
            self.assertEqual(snp_s1.shape[1], snp_s2.shape[1])

            self.assertEqual(snp_s1.shape[0], snp_s3.shape[0])
            self.assertEqual(snp_s1.shape[1], snp_s3.shape[1])
        
            self.assertEqual(snp_s1.shape[0], snp_s4.shape[0])
            self.assertEqual(snp_s1.shape[1], snp_s4.shape[1])

            self.assertTrue(np.allclose(snp_s1, snp_s2, rtol=1e-05, atol=1e-05))
            self.assertTrue(np.allclose(snp_s1, snp_s3, rtol=1e-05, atol=1e-05))
            self.assertTrue(np.allclose(snp_s1, snp_s4, rtol=1e-05, atol=1e-05))

            self.assertEqual(snp_beta1.shape[0], snp_beta2.shape[0])
            self.assertEqual(snp_beta1.shape[1], snp_beta2.shape[1])
            self.assertEqual(snp_beta1.shape[0], snp_beta3.shape[0])
            self.assertEqual(snp_beta1.shape[1], snp_beta3.shape[1])
        
            self.assertTrue(np.allclose(snp_beta1, snp_beta2, rtol=1e-05, atol=1e-05))
            self.assertTrue(np.allclose(snp_beta1, snp_beta3, rtol=1e-05, atol=1e-05))

    def test_load_and_standardize_bed(self):
        snpreader2 = Bed(self.currentFolder + "/examples/toydata")
        self.load_and_standardize(snpreader2, snpreader2)

    def test_subset_view(self):
        snpreader2 = Bed(self.currentFolder + "/examples/toydata")[:,:]
        result = snpreader2.read(view_ok=True)
        self.assertFalse(snpreader2 is result)
        result2 = result[:,:].read()
        self.assertFalse(SP.may_share_memory(result2.val,result.val))
        result3 = result[:,:].read(view_ok=True)
        self.assertTrue(SP.may_share_memory(result3.val,result.val))
        result4 = result3.read()
        self.assertFalse(SP.may_share_memory(result4.val,result3.val))
        result5 = result4.read(view_ok=True)
        self.assertTrue(SP.may_share_memory(result4.val,result5.val))


    def test_load_and_standardize_hdf5(self):
        snpreader2 = Hdf5(self.currentFolder + "/examples/toydata.snpmajor.hdf5")
        snpreader3 = Hdf5(self.currentFolder + "/examples/toydata.iidmajor.hdf5")
        self.load_and_standardize(snpreader2, snpreader3)
        snpreaderref = Bed(self.currentFolder + "/examples/toydata")
        self.load_and_standardize(snpreader2, snpreaderref)



    def test_load_and_standardize_dat(self):
        snpreader2 = Dat(self.currentFolder + "/examples/toydata.dat")
        self.load_and_standardize(snpreader2, snpreader2)
        #snpreaderref = Bed(self.currentFolder + "/examples/toydata")
        #self.load_and_standardize(snpreader2, snpreaderref)

    def test_load_and_standardize_ped(self):

        #!!!cmk04072014 Ped columns can be ambiguous
        ###Creating Ped data ...
        #currentFolder = os.path.dirname(os.path.realpath(__file__))
        #snpData = Bed(currentFolder + "/examples/toydata").read()
        ##Ped.write(snpData, currentFolder + "/examples/toydata.ped")
        #fromPed = Ped(currentFolder + "/examples/toydata").read()
        #self.assertTrue(np.allclose(snpData.val, fromPed.val, rtol=1e-05, atol=1e-05))


        snpreader2 = Ped(self.currentFolder + "/examples/toydata")
        self.load_and_standardize(snpreader2, snpreader2)
        #snpreaderref = Bed(self.currentFolder + "/examples/toydata")
        #self.load_and_standardize(snpreader2, snpreaderref)

    def load_and_standardize(self, snpreader2, snpreader3):
        """
        test c-version of load and standardize
        """

        S = snpreader2.sid_count
        N_original = snpreader2.iid_count

        iid_index_list = range(N_original - 1,0,-2)
        snpreader3 = snpreader3[iid_index_list,:]

        for dtype in [sp.float64,sp.float32]:

            G2 = snpreader2.read(order='F',force_python_only=True).val
            G2 = Unit().standardize(G2, blocksize=10000, force_python_only=True)

            SNPs_floatF = snpreader2.read(order="F", dtype=dtype, force_python_only=False).val
            GF = Unit().standardize(SNPs_floatF)

            SNPs_floatC = snpreader2.read(order="C", dtype=dtype, force_python_only=False).val
            GC = Unit().standardize(SNPs_floatC)

            self.assertTrue(np.allclose(GF, G2, rtol=1e-05, atol=1e-05))
            self.assertTrue(np.allclose(GF, GC, rtol=1e-05, atol=1e-05))

            #testing selecting a subset of snps and iids
            snp_index_list = range(S - 1,0,-2)

            G2x = snpreader2.read(order='F',force_python_only=True).val
            G2x = G2x[iid_index_list,:][:,snp_index_list]
            G2x = Unit().standardize(G2x, blocksize=10000, force_python_only=True)


            SNPs_floatFx = snpreader3[:,snp_index_list].read(order="F", dtype=dtype, force_python_only=False).val
            GFx = Unit().standardize(SNPs_floatFx)
            self.assertTrue(np.allclose(GFx, G2x, rtol=1e-05, atol=1e-05))

            SNPs_floatCx = snpreader3[:,snp_index_list].read(order="C", dtype=dtype, force_python_only=False).val
            GCx = Unit().standardize(SNPs_floatCx)
            self.assertTrue(np.allclose(GFx, G2x, rtol=1e-05, atol=1e-05))

class NaNCNCTestCases(unittest.TestCase):
    def __init__(self, iid_index_list, snp_index_list, standardizer, snpreader, dtype, order, force_python_only, reference_snps, reference_dtype):
        self.iid_index_list = iid_index_list
        self.snp_index_list = snp_index_list
        self.standardizer = standardizer
        self.snpreader = snpreader
        self.dtype = dtype
        self.order = order
        self.force_python_only = force_python_only
        self.reference_snps = reference_snps
        self.reference_dtype = reference_dtype

    _testMethodName = "runTest"
    _testMethodDoc = None

    @staticmethod
    def factory_iterator():
        snp_reader_factory_bed = lambda : Bed("examples/toydata")
        snp_reader_factory_snpmajor_hdf5 = lambda : Hdf5("examples/toydata.snpmajor.hdf5")
        snp_reader_factory_iidmajor_hdf5 = lambda : Hdf5("examples/toydata.iidmajor.hdf5",blocksize=6000)
        snp_reader_factory_dat = lambda : Dat("examples/toydata.dat")

        previous_wd = os.getcwd()
        os.chdir(os.path.dirname(os.path.realpath(__file__)))
        snpreader0 = snp_reader_factory_bed()
        S_original = snpreader0.sid_count
        N_original = snpreader0.iid_count

        snps_to_read_count = min(S_original, 100)

        for iid_index_list in [range(N_original), range(N_original/2), range(N_original - 1,0,-2)]:
            for snp_index_list in [range(snps_to_read_count), range(snps_to_read_count/2), range(snps_to_read_count - 1,0,-2)]:
                for standardizer in [Unit(),Beta(1,25)]:
                    reference_snps, reference_dtype = NaNCNCTestCases(iid_index_list, snp_index_list, standardizer, snp_reader_factory_bed(), sp.float64, "C", "False", None, None).read_and_standardize()
                    for snpreader_factory in [snp_reader_factory_bed, snp_reader_factory_snpmajor_hdf5, snp_reader_factory_iidmajor_hdf5, snp_reader_factory_dat]:
                        for dtype in [sp.float64,sp.float32]:
                            for order in ["C", "F"]:
                                for force_python_only in [False, True]:
                                    snpreader = snpreader_factory()
                                    test_case = NaNCNCTestCases(iid_index_list, snp_index_list, standardizer, snpreader, dtype, order, force_python_only, reference_snps, reference_dtype)
                                    yield test_case
        os.chdir(previous_wd)

    def __str__(self):
        iid_index_list = self.iid_index_list
        snp_index_list = self.snp_index_list
        standardizer = self.standardizer
        snpreader = self.snpreader
        dtype = self.dtype
        order = self.order
        force_python_only = self.force_python_only
        return "{0}(iid_index_list=[{1}], snp_index_list=[{2}], standardizer={3}, snpreader={4}, dtype={5}, order='{6}', force_python_only=={7})".format(
            self.__class__.__name__,
            ",".join([str(i) for i in iid_index_list]) if len(iid_index_list) < 10 else ",".join([str(i) for i in iid_index_list[0:10]])+",...",
            ",".join([str(i) for i in snp_index_list]) if len(snp_index_list) < 10 else ",".join([str(i) for i in snp_index_list[0:10]])+",...",
            standardizer, snpreader, dtype, order, force_python_only)

    def read_and_standardize(self):
        previous_wd = os.getcwd()
        os.chdir(os.path.dirname(os.path.realpath(__file__)))
        iid_index_list = self.iid_index_list
        snp_index_list = self.snp_index_list
        standardizer = self.standardizer
        snpreader = self.snpreader
        dtype = self.dtype
        order = self.order
        force_python_only = self.force_python_only
        
        snps = snpreader[iid_index_list,snp_index_list].read(order=order, dtype=dtype, force_python_only=force_python_only).val
        snps[0,0] = np.nan # add a NaN
        snps[:,1] = 2 # make a SNC
        snps = standardizer.standardize(snps,force_python_only=force_python_only)
        os.chdir(previous_wd)
        return snps, dtype

    def doCleanups(self):
        pass
        #return super(NaNCNCTestCases, self).doCleanups()

    def runTest(self, result = None):
        snps, dtype = self.read_and_standardize()
        self.assertTrue(snps[0,0] == 0)
        self.assertTrue(np.all(snps[:,1] == 0))
        if self.reference_snps is not None:
            self.assertTrue(np.allclose(self.reference_snps, snps, rtol=1e-04 if dtype == sp.float32 or self.reference_dtype == sp.float32 else 1e-12))

# We do it this way instead of using doctest.DocTestSuite because doctest.DocTestSuite requires modules to be pickled, which python doesn't allow.
# We need tests to be picklable so that they can be run on a cluster.
class TestDocStrings(unittest.TestCase):
    def test_snpreader(self):
        import pysnptools.pysnptools.snpreader.snpreader
        os.chdir(snpreader_path)
        doctest.testmod(pysnptools.pysnptools.snpreader.snpreader)

    def test_bed(self):
        import pysnptools.pysnptools.snpreader.bed
        os.chdir(snpreader_path)
        doctest.testmod(pysnptools.pysnptools.snpreader.bed)

    def test_snpdata(self):
        import pysnptools.pysnptools.snpreader.snpdata
        os.chdir(snpreader_path)
        doctest.testmod(pysnptools.pysnptools.snpreader.snpdata)


    def test_util(self):
        import pysnptools.pysnptools.util.util
        os.chdir(snpreader_path)
        doctest.testmod(pysnptools.pysnptools.util.util)

#!!!cmk04072014 Need to add more of these?



def getTestSuite():
    """
    set up composite test suite
    """
    
    test_suite = unittest.TestSuite([])
    test_suite.addTests(unittest.TestLoader().loadTestsFromTestCase(TestLoader))
    test_suite.addTests(NaNCNCTestCases.factory_iterator())
    test_suite.addTests(unittest.TestLoader().loadTestsFromTestCase(TestDocStrings))

    return test_suite

snpreader_path = os.path.join(os.path.dirname(os.path.realpath(__file__)),"snpreader")


if __name__ == '__main__':

    # TestFeatureSelection().test_aaa_hdf5_speed()
    r = unittest.TextTestRunner(failfast=False)
    r.run(getTestSuite())

