import numpy as np
import logging
import doctest
import unittest
import os.path
import time

from pysnptools.kernelreader import *
from pysnptools.snpreader import Bed
from pysnptools.util import create_directory_if_necessary
from pysnptools.pstreader import PstReader
from pysnptools.snpreader import SnpData
import pysnptools.standardizer as stdizer

class _fortesting_JustCheckExists(object): #Implements ICopier

    def __init__(self,doPrintOutputNames=False):
        self.doPrintOutputNames = doPrintOutputNames
    
    def input(self,item):
        if isinstance(item, str):
            if not os.path.exists(item): raise Exception("Missing input file '{0}'".format(item))
        elif hasattr(item,"copyinputs"):
            item.copyinputs(self)
        # else -- do nothing

    def output(self,item):
        if isinstance(item, str):
            if not os.path.exists(item): raise Exception("Missing output file '{0}'".format(item))
            if self.doPrintOutputNames:
                print(item)
        elif hasattr(item,"copyoutputs"):
            item.copyoutputs(self)
        # else -- do nothing




class TestKernelReader(unittest.TestCase):
    @classmethod
    def setUpClass(self):
        self.currentFolder = os.path.dirname(os.path.realpath(__file__))

    def test_merge_std(self):
        #unit vs beta
        for std in [stdizer.Beta(2,10),stdizer.Unit()]:
            np.random.seed(0)
            sid_count = 20
            snpreader = SnpData(iid=[["0","0"],["1","1"],["2","2"]],sid=[str(i) for i in range(sid_count)],val=np.array(np.random.randint(3,size=[3,sid_count]),dtype=np.float64,order='F'))
            kerneldata0,trained0,diag0 = SnpKernel(snpreader,std,block_size=1)._read_with_standardizing(to_kerneldata=True,return_trained=True)
            kerneldata1,trained1,diag1 = SnpKernel(snpreader,std,block_size=None)._read_with_standardizing(to_kerneldata=True,return_trained=True)
            np.testing.assert_array_almost_equal(kerneldata0.val,kerneldata1.val, decimal=10)
            np.testing.assert_array_almost_equal(trained0.stats,trained1.stats, decimal=10)
            assert abs(diag0.factor-diag1.factor) < 1e-7
    
    def test_cpp_std(self):

        #Order C vs F
        for order in ['C','F']:
            #32 vs 64
            for dtype in [np.float64,np.float32]:
                #unit vs beta
                for std in [stdizer.Unit(),stdizer.Beta(2,10)]:
                        np.random.seed(0)
                        sid_count = 20
                        snpreader0 = SnpData(iid=[["0","0"],["1","1"],["2","2"]],sid=[str(i) for i in range(sid_count)],val=np.array(np.random.randint(3,size=[3,sid_count]),dtype=dtype,order=order))
                        snpreader1 = SnpData(iid=[["3","3"],["4","4"]],sid=[str(i) for i in range(sid_count)],val=np.array(np.random.randint(3,size=[2,sid_count]),dtype=dtype,order=order))

                        #has SNC
                        for has_SNC_in_train in [False, True]:
                            if has_SNC_in_train:
                                snpreader0.val[:,1] = 0

                            #missing data
                            for has_missing_data in [False, True]:
                                if has_missing_data:
                                    snpreader0.val[0,2]=np.nan
                                    snpreader1.val[0,2]=np.nan

                                #gather stats vs not
                                cppa, stdcppa = snpreader0.read(order=order,dtype=dtype).standardize(std,return_trained=True,force_python_only=False)
                                pya, stdpya = snpreader0.read(order=order,dtype=dtype).standardize(std,return_trained=True,force_python_only=True)
                                np.testing.assert_array_almost_equal(cppa.val, pya.val, decimal=10 if dtype==np.float64 else 5)

                                np.testing.assert_array_almost_equal(stdcppa.stats,stdpya.stats, decimal=10 if dtype==np.float64 else 5)
                                assert (np.inf in stdcppa.stats[:,1]) == has_SNC_in_train
                                assert (np.inf in stdpya.stats[:,1]) == has_SNC_in_train

                                if has_SNC_in_train:
                                    assert np.array_equal(cppa.val[:,1],np.zeros([cppa.val.shape[0]]))
                                    assert np.array_equal(pya.val[:,1],np.zeros([pya.val.shape[0]]))

                                if has_missing_data:
                                    assert 0 == cppa.val[0,2]
                                    assert 0 == pya.val[0,2]
                                        
                                #uses stats
                                cppb = snpreader1.read(order=order,dtype=dtype).standardize(stdcppa,force_python_only=False)
                                pyb = snpreader1.read(order=order,dtype=dtype).standardize(stdpya,force_python_only=True)
                                np.testing.assert_array_almost_equal(cppb.val, pyb.val, decimal=10 if dtype==np.float64 else 5)
                                np.testing.assert_array_almost_equal(stdcppa.stats,stdpya.stats, decimal=10 if dtype==np.float64 else 5) #Make sure we haven't messed up the train stats

                                if has_SNC_in_train:
                                    assert np.array_equal(cppb.val[:,1],np.zeros([cppb.val.shape[0]]))
                                    assert np.array_equal(pyb.val[:,1],np.zeros([pyb.val.shape[0]]))

                                if has_missing_data:
                                    assert cppb.val[0,2]==0
                                    assert pyb.val[0,2]==0
        logging.info("done with 'test_cpp_std'")


    def test_intersection(self):

        from pysnptools.standardizer import Unit
        from pysnptools.kernelreader import SnpKernel
        from pysnptools.snpreader import Pheno
        from pysnptools.kernelreader._subset import _KernelSubset
        from pysnptools.snpreader._subset import _SnpSubset
        from pysnptools.util import intersect_apply

        snps_all = Bed(self.currentFolder + "/../examples/toydata.5chrom.bed",count_A1=False)
        k = SnpKernel(snps_all,stdizer.Identity())

        pheno = Pheno(self.currentFolder + "/../examples/toydata.phe")
        pheno = pheno[1:,:] # To test intersection we remove a iid from pheno

        k1,pheno = intersect_apply([k,pheno]) #SnpKernel is special because it standardizes AFTER intersecting.
        assert isinstance(k1.snpreader,_SnpSubset) and not isinstance(k1,_KernelSubset)

        #What happens with fancy selection?
        k2 = k[::2]
        assert isinstance(k2,SnpKernel)

        logging.info("Done with test_intersection")

    def test_respect_read_inputs(self):
        from pysnptools.kernelreader import KernelHdf5,Identity,KernelNpz,SnpKernel
        from pysnptools.standardizer import Unit
        from pysnptools.standardizer import Identity as StdIdentity
        from pysnptools.snpreader import Bed

        previous_wd = os.getcwd()
        os.chdir(os.path.dirname(os.path.realpath(__file__)))

        iidref = KernelNpz('../examples/toydata.kernel.npz').iid

        for kernelreader in [
                           SnpKernel(Bed('../examples/toydata.5chrom.bed',count_A1=True),StdIdentity())[::2,::2],
                           Bed('../examples/toydata.5chrom.bed',count_A1=True)[::2,::2].read_kernel(StdIdentity()),
                           KernelHdf5('../examples/toydata.kernel.hdf5'),
                           Identity(iidref,test=[('0','x'),('0','y')]),
                           Identity(iidref),
                           KernelNpz('../examples/toydata.kernel.npz'),
                           KernelNpz('../examples/toydata.kernel.npz').read(),
                           KernelNpz('../examples/toydata.kernel.npz')[::2,::2],
                           Bed('../examples/toydata.5chrom.bed',count_A1=True).read_kernel(Unit()),
                           SnpKernel(Bed('../examples/toydata.5chrom.bed',count_A1=True),Unit())
                          ]:
            logging.info(str(kernelreader))
            for order in ['F','C','A']:
                for dtype in [np.float32,np.float64]:
                    for force_python_only in [True,False]:
                        for view_ok in [True,False]:
                            val = kernelreader.read(order=order,dtype=dtype,force_python_only=force_python_only,view_ok=view_ok).val
                            has_right_order = order=="A" or (order=="C" and val.flags["C_CONTIGUOUS"]) or (order=="F" and val.flags["F_CONTIGUOUS"])
                            if hasattr(kernelreader,'val') and not view_ok:
                                assert kernelreader.val is not val
                            if (hasattr(kernelreader,'val') and view_ok and kernelreader.val is not val and
                                (order == 'A' or (order == 'F' and kernelreader.val.flags['F_CONTIGUOUS']) or (order == 'C' and kernelreader.val.flags['C_CONTIGUOUS'])) and
                                (dtype is None or  kernelreader.val.dtype == dtype)):
                                logging.info("{0} could have read a view, but didn't".format(distreader))
                            assert val.dtype == dtype and has_right_order
        os.chdir(previous_wd)


    def test_respect_inputs(self):
        np.random.seed(0)
        for dtype_start,decimal_start in [(np.float32,5),(np.float64,10)]:
            for order_start in ['F','C','A']:
                for sid_count in [20,2]:
                    snpdataX = SnpData(iid=[["0","0"],["1","1"],["2","2"]],sid=[str(i) for i in range(sid_count)],val=np.array(np.random.randint(3,size=[3,sid_count]),dtype=dtype_start,order=order_start))
                    for stdx in [stdizer.Beta(1,25),stdizer.Identity(),stdizer.Unit()]:
                        for snpreader0 in [snpdataX,snpdataX[:,1:]]:
                            snpreader1 = snpreader0[1:,:]

                            refdata0, trained_standardizer = snpreader0.read().standardize(stdx,return_trained=True,force_python_only=True)
                            refval0 = refdata0.val.dot(refdata0.val.T)
                            refdata1 = snpreader1.read().standardize(trained_standardizer,force_python_only=True)#LATER why aren't these used?
                            refval1 = refdata0.val.dot(refdata1.val.T)#LATER why aren't these used?
                            for dtype_goal,decimal_goal in [(np.float32,5),(np.float64,10)]:
                                for order_goal in ['F','C','A']:
                                    k = snpreader0.read_kernel(standardizer=stdx,block_size=1,order=order_goal,dtype=dtype_goal)
                                    PstReader._array_properties_are_ok(k.val,order_goal,dtype_goal)
                                    np.testing.assert_array_almost_equal(refval0,k.val, decimal=min(decimal_start,decimal_goal))

    def test_fail(self):
        did_fail = True
        try:
            kd = KernelData(iid=[["0"],["1"],["2"]],val=[[1,2,3],[4,5,6],[7,8,9]]) #Wrong iid shape
            did_fail = False
        except:
            pass
        assert did_fail, "The constructor should fail because the iid is the wrong shape"

    def test_kernel2(self):
        logging.info("in kernel2")
        kd = KernelData(iid=[["0","0"],["1","1"],["2","2"]],val=[[1,2,3],[4,5,6],[7,8,9]])
        assert np.array_equal(kd.iid_to_index([["1","1"],["2","2"]]),np.array([1,2]))
        assert np.array_equal(kd.iid0_to_index([["1","1"],["2","2"]]),np.array([1,2]))
        kd = kd.standardize()
        assert np.abs(np.diag(kd.val).sum()-3)<1e-7
        assert kd.iid1_count == 3

    def test_snp_kernel2(self):
        logging.info("in test_snp_kernel2")
        snpreader = Bed(self.currentFolder + "/../examples/toydata.5chrom.bed",count_A1=False)
        snpkernel = SnpKernel(snpreader,standardizer=stdizer.Beta(1,25))
        s  = str(snpkernel)
        _fortesting_JustCheckExists().input(snpkernel)
        
    def test_npz(self):
        logging.info("in test_npz")
        snpreader = Bed(self.currentFolder + "/../examples/toydata.5chrom.bed",count_A1=False)
        kerneldata1 = snpreader.read_kernel(standardizer=stdizer.Unit())
        s = str(kerneldata1)
        output = "tempdir/kernelreader/toydata.kernel.npz"
        create_directory_if_necessary(output)
        KernelNpz.write(output,kerneldata1)
        kernelreader2 = KernelNpz(output)
        kerneldata2 = kernelreader2.read()
        np.testing.assert_array_almost_equal(kerneldata1.val, kerneldata2.val, decimal=10)
        logging.info("done with test")

    def test_subset(self):
        logging.info("in test_subset")
        snpreader = Bed(self.currentFolder + "/../examples/toydata.5chrom.bed",count_A1=False)
        snpkernel = SnpKernel(snpreader,stdizer.Unit())
        krsub = snpkernel[::2,::2]
        kerneldata1 = krsub.read()
        expected = snpreader.read_kernel(stdizer.Unit())[::2].read()
        np.testing.assert_array_almost_equal(kerneldata1.val, expected.val, decimal=10)

        krsub2 = snpkernel[::2]
        kerneldata2 = krsub2.read()
        np.testing.assert_array_almost_equal(kerneldata2.val, expected.val, decimal=10)
        logging.info("done with test")

    def test_identity(self):
        logging.info("in test_identity")
        snpreader = Bed(self.currentFolder + "/../examples/toydata.5chrom.bed",count_A1=False)
        assert snpreader.iid is snpreader.row
        kid = Identity(snpreader.row)
        assert np.array_equal(kid.row,kid.iid) and np.array_equal(kid.iid,kid.iid0) and np.array_equal(kid.iid0,kid.iid1) and np.array_equal(kid.iid1, kid.col)
        np.testing.assert_array_almost_equal(kid.read().val, np.identity(snpreader.iid_count))

        #subset
        sub1 = kid[1:5]
        np.testing.assert_array_almost_equal(sub1.read().val, np.identity(4))

        #zero size
        sub2 = kid[0:0]
        np.testing.assert_array_almost_equal(sub2.read().val, np.identity(0))

    def test_identity_sub(self):
        logging.info("in test_identity_sub")
        snpreader = Bed(self.currentFolder + "/../examples/toydata.5chrom.bed",count_A1=False)
        assert snpreader.iid is snpreader.row
        kid = Identity(snpreader.row)
        sub3 = kid[::2,1:5]
        expected = np.identity(kid.iid_count)[::2,:][:,1:5]
        np.testing.assert_array_almost_equal(sub3.read().val,expected)

        logging.info("done with test")

    def test_underscore_read1(self):
        logging.info("in test_underscore_read1")
        snpreader = Bed(self.currentFolder + "/../examples/toydata.5chrom.bed",count_A1=False)
        assert snpreader.iid is snpreader.row
        kid = Identity(snpreader.row)
        sub3 = kid[::2,::2]
        expected = np.identity(kid.iid_count)[::2,:][:,::2]
        np.testing.assert_array_almost_equal(sub3.read().val,expected)

        logging.info("done with test")

    def test_underscore_read2(self):
        logging.info("in test_underscore_read2")
        snpreader = Bed(self.currentFolder + "/../examples/toydata.5chrom.bed",count_A1=False)
        assert snpreader.iid is snpreader.row
        kid = Identity(snpreader.row)
        sub3 = kid[::2,::2]
        expected = np.identity(kid.iid_count)[::2,:][:,::2]
        np.testing.assert_array_almost_equal(sub3.read().val,expected)

        logging.info("done with test")

# 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 pickleable so that they can be run on a cluster.
class TestKrDocStrings(unittest.TestCase):
    def test_kernelreader(self):
        import pysnptools.kernelreader.kernelreader
        old_dir = os.getcwd()
        os.chdir(os.path.dirname(os.path.realpath(__file__)))
        result = doctest.testmod(pysnptools.kernelreader.kernelreader,optionflags=doctest.ELLIPSIS)
        os.chdir(old_dir)
        assert result.failed == 0, "failed doc test: " + __file__

    def test_snpkernel(self):
        import pysnptools.kernelreader.snpkernel
        old_dir = os.getcwd()
        os.chdir(os.path.dirname(os.path.realpath(__file__)))
        result = doctest.testmod(pysnptools.kernelreader.snpkernel,optionflags=doctest.ELLIPSIS)
        os.chdir(old_dir)
        assert result.failed == 0, "failed doc test: " + __file__

    def test_kerneldata(self):
        import pysnptools.kernelreader.kerneldata
        old_dir = os.getcwd()
        os.chdir(os.path.dirname(os.path.realpath(__file__)))
        result = doctest.testmod(pysnptools.kernelreader.kerneldata,optionflags=doctest.ELLIPSIS)
        os.chdir(old_dir)
        assert result.failed == 0, "failed doc test: " + __file__

    def test_identity(self):
        import pysnptools.kernelreader.identity
        old_dir = os.getcwd()
        os.chdir(os.path.dirname(os.path.realpath(__file__)))
        result = doctest.testmod(pysnptools.kernelreader.identity,optionflags=doctest.ELLIPSIS|doctest.NORMALIZE_WHITESPACE)
        os.chdir(old_dir)
        assert result.failed == 0, "failed doc test: " + __file__

    def test_kernelhdf5(self):
        import pysnptools.kernelreader.kernelhdf5
        old_dir = os.getcwd()
        os.chdir(os.path.dirname(os.path.realpath(__file__)))
        result = doctest.testmod(pysnptools.kernelreader.kernelhdf5,optionflags=doctest.ELLIPSIS)
        os.chdir(old_dir)
        assert result.failed == 0, "failed doc test: " + __file__

    def test_kernelnpz(self):
        import pysnptools.kernelreader.kernelnpz
        old_dir = os.getcwd()
        os.chdir(os.path.dirname(os.path.realpath(__file__)))
        result = doctest.testmod(pysnptools.kernelreader.kernelnpz,optionflags=doctest.ELLIPSIS)
        os.chdir(old_dir)
        assert result.failed == 0, "failed doc test: " + __file__



def getTestSuite():
    """
    set up composite test suite
    """
    
    test_suite = unittest.TestSuite([])
    test_suite.addTests(unittest.TestLoader().loadTestsFromTestCase(TestKernelReader))
    test_suite.addTests(unittest.TestLoader().loadTestsFromTestCase(TestKrDocStrings))
    return test_suite

if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO)

    suites = getTestSuite()
    r = unittest.TextTestRunner(failfast=False)
    ret = r.run(suites)
    assert ret.wasSuccessful()

