import os
import numpy as np
import logging
import unittest
import doctest
from pathlib import Path
from pysnptools.util.pairs import _Pairs as UtilPairs

from pysnptools.snpreader import SnpReader

class _Pairs(SnpReader):
    '''
    Experimental.
    '''

    #!!!see fastlmm\association\epistasis.py for code that allows ranges of snps to be specified when making pairs
    def __init__(self, snpreader0,snpreader1=None, do_standardize=True,sid_materialize_limit=1000*1000,_include_single_times_single=False): 
        #!!! could add option to change snp separator and another to encode chrom, etc in the snp name
        super(_Pairs, self).__init__()
        self._ran_once = False
        self.snpreader0 = snpreader0
        self.snpreader1 = snpreader1
        self.do_standardize = do_standardize
        self.sid_materialize_limit = sid_materialize_limit
        self._include_single_times_single=_include_single_times_single

    def __repr__(self):
        part2 = '' if self.snpreader1 is None else ',{0}'.format(self.snpreader1)
        return "{0}({1}{2})".format(self.__class__.__name__,self.snpreader0,part2)

    @property
    def row(self):
        """*same as* :attr:`iid`
        """
        if not hasattr(self,"_row"):
            self._row = self.snpreader0.row
            assert self.snpreader1 is None or np.array_equal(self._row,self.snpreader1.row), "Expect snpreaders to have the same iids in the same order"
        return self._row

    @property
    def col(self):
        """*same as* :attr:`sid`
        """
        if not hasattr(self,"_col"):
            assert self.col_count < self.sid_materialize_limit, '{:,} is too many sids to materialize'.format(self.col_count)
            if self.snpreader1 is None:
                col_0 = self.snpreader0.col
                col_list = []
                self.index0_list = [] #!!!should be _index0_list, etc
                self.index1_list = []
                for index0 in xrange(self.snpreader0.col_count):
                    #logging.info("index0={0} of {1}".format(index0,self.snpreader0.col_count))#!!!
                    start1 = index0 if self._include_single_times_single else index0+1
                    for index1 in xrange(start1,self.snpreader0.col_count):
                        col_list.append('{0},{1}'.format(col_0[index0],col_0[index1]))
                        self.index0_list.append(index0)
                        self.index1_list.append(index1)
                self._col = np.array(col_list)
                self.index0_list = np.array(self.index0_list)
                self.index1_list = np.array(self.index1_list)
            else:
                col_0 = self.snpreader0.col
                col_1 = self.snpreader1.col
                assert len(set(col_0)&set(col_1))==0,"Pairs currently requires two snpreaders to have no sids in common"
                col_list = []
                self.index0_list = []
                self.index1_list = []
                for index0 in xrange(self.snpreader0.col_count):
                    #logging.info("index0={0} of {1}".format(index0,self.snpreader0.col_count))#!!!
                    for index1 in xrange(self.snpreader1.col_count):
                        col_list.append('{0},{1}'.format(col_0[index0],col_1[index1]))
                        self.index0_list.append(index0)
                        self.index1_list.append(index1)
                self._col = np.array(col_list)
                self.index0_list = np.array(self.index0_list)
                self.index1_list = np.array(self.index1_list)

        assert self.col_count == len(self._col), "real assert"
        return self._col

    @property
    def col_count(self):
        n0 = self.snpreader0.col_count
        if self.snpreader1 is None:
            if self._include_single_times_single:
                return (n0*n0+n0)//2
            else:
                return (n0*n0-n0)//2
        else:
            return n0*self.snpreader1.col_count

    @property
    def col_property(self):
        """*same as* :attr:`pos`
        """
        if not hasattr(self,"_col_property"):
            self._col_property = np.zeros([self.sid_count,3],dtype=np.int64)
        return self._col_property

    def copyinputs(self, copier):
        # doesn't need to self._run_once() because only uses original inputs !!!is this true?
        self.snpreader0.copyinputs(copier)
        if self.snpreader1 is not None:
            self.snpreader1.copyinputs(copier)

    def _run_once(self):
        if self._ran_once:
            return

        self._ran_once = True
        self.col

    def _read(self, iid_index_or_none, sid_index_or_none, order, dtype, force_python_only, view_ok, num_threads):
        self._run_once()
        dtype = np.dtype(dtype)

        iid_count_in = self.iid_count
        sid_count_in = self.sid_count

        if iid_index_or_none is not None:
            iid_count_out = len(iid_index_or_none)
            iid_index = iid_index_or_none
        else:
            iid_count_out = iid_count_in
            iid_index = range(iid_count_in)

        if sid_index_or_none is not None:
            sid_count_out = len(sid_index_or_none)
            sid_index = sid_index_or_none
        else:
            sid_count_out = sid_count_in
            sid_index = range(sid_count_in)

        
        sid_index_inner_0 = self.index0_list[sid_index] #Find the sid_index of the left snps of interest
        sid_index_inner_1 = self.index1_list[sid_index] #Find the sid_index of the right snps of interest
        if self.snpreader1 is None:
            sid_index_inner_01 = np.unique(np.r_[sid_index_inner_0,sid_index_inner_1]) #Index of every snp of interest
            inner_01 = self.snpreader0[iid_index_or_none,sid_index_inner_01].read(order=order, dtype=dtype, force_python_only=force_python_only, view_ok=True,num_threads=num_threads) #read every val of interest
            val_inner_01 = inner_01.standardize(num_threads=num_threads).val if self.do_standardize else inner_01.val

            sid_index_inner_01_reverse = {v:i for i,v in enumerate(sid_index_inner_01)} #Dictionary of snp_index to position in sid_index_inner_01
            sid_index_inner_0_in_val = np.array([sid_index_inner_01_reverse[i] for i in sid_index_inner_0])  #Replace snp_index0 with column # in val_inner_01
            sid_index_inner_1_in_val = np.array([sid_index_inner_01_reverse[i] for i in sid_index_inner_1])  #Replace snp_index1 with column # in val_inner_01
            val_inner_0 = val_inner_01[:,sid_index_inner_0_in_val] #Extract the vals for the left snps of interest
            val_inner_1 = val_inner_01[:,sid_index_inner_1_in_val]#Extract the vals for the right snps of interest
        else:
            inner_0 = self.snpreader0[iid_index_or_none,sid_index_inner_0].read(order=order, dtype=dtype, force_python_only=force_python_only, view_ok=True,num_threads=num_threads) #read every val of interest
            inner_1 = self.snpreader1[iid_index_or_none,sid_index_inner_1].read(order=order, dtype=dtype, force_python_only=force_python_only, view_ok=True,num_threads=num_threads) #read every val of interest
            val_inner_0 = inner_0.standardize(num_threads=num_threads).val if self.do_standardize else inner_0.val
            val_inner_1 = inner_1.standardize(num_threads=num_threads).val if self.do_standardize else inner_1.val
        val = val_inner_0*val_inner_1 #Element multiplication creates the vals for the pairs
        return val



def split_on_sids(snpreader,part_count):
    sid_count = snpreader.sid_count
    start = 0
    for part_index in xrange(1,part_count+1):
        end = part_index*sid_count//part_count
        yield snpreader[:,start:end]
        start=end


class _Pairs2(SnpReader):
    
    def __init__(self, snpreader0,snpreader1=None, do_standardize=True,sid_materialize_limit=1000*1000,_include_single_times_single=False): #!!! could add option to change snp separator and another to encode chrom, etc in the snp name
        super(_Pairs2, self).__init__()
        self._ran_once = False
        self.snpreader0 = snpreader0
        self.snpreader1 = snpreader1
        self.do_standardize = do_standardize
        self.sid_materialize_limit = sid_materialize_limit
        self._include_single_times_single=_include_single_times_single
        self._utilpairs = UtilPairs(snpreader0.sid,snpreader1.sid if snpreader1 is not None else snpreader0.sid,include_singles=_include_single_times_single,duplicates_ok=False)

    def __repr__(self):
        part2 = '' if self.snpreader1 is None else ',{0}'.format(self.snpreader1)
        return "{0}({1}{2})".format(self.__class__.__name__,self.snpreader0,part2)

    @property
    def row(self):
        """*same as* :attr:`iid`
        """
        if not hasattr(self,"_row"):
            self._row = self.snpreader0.row
            assert self.snpreader1 is None or np.array_equal(self._row,self.snpreader1.row), "Expect snpreaders to have the same iids in the same order"
        return self._row

    @property
    def col(self):
        """*same as* :attr:`sid`
        """
        if not hasattr(self,"_col"):
            assert self.col_count < self.sid_materialize_limit, '{:,} is too many sids to materialize'.format(self.col_count)
            snpreader1 = self.snpreader1 if self.snpreader1 is not None else self.snpreader0
            #!!!self.index0_list = self.snpreader0.sid_to_index((sid0 for sid0,sid1 in self._utilpairs[:])) #!!! can we do without these?
            #!!!self.index1_list = snpreader1.sid_to_index(sid1 for sid0,sid1 in self._utilpairs[:])#!!!can we do without these?
            self._col = np.array(list(','.join(pair) for pair in self._utilpairs[:]))
        return self._col

    @property
    def col_count(self):
        return len(self._utilpairs)

    @property
    def col_property(self):
        """*same as* :attr:`pos`
        """
        if not hasattr(self,"_col_property"):
            self._col_property = np.zeros([self.sid_count,3],dtype=np.int64)
        return self._col_property

    def copyinputs(self, copier):
        # doesn't need to self.run_once() because only uses original inputs !!!is this true?
        self.snpreader0.copyinputs(copier)
        if self.snpreader1 is not None:
            self.snpreader1.copyinputs(copier)

    def run_once(self):
        if self._ran_once:
            return

        self._ran_once = True
        self.col

    def _read(self, iid_index_or_none, sid_index_or_none, order, dtype, force_python_only, view_ok, num_threads):
        self.run_once()

        iid_count_in = self.iid_count
        sid_count_in = self.sid_count

        if iid_index_or_none is not None:
            iid_count_out = len(iid_index_or_none)
            iid_index_out = iid_index_or_none
        else:
            iid_count_out = iid_count_in
            iid_index_out = range(iid_count_in)

        if sid_index_or_none is not None:
            sid_count_out = len(sid_index_or_none)
            sid_index_out = sid_index_or_none
        else:
            sid_count_out = sid_count_in
            sid_index_out = splice(None)#!!! test this

        pair_array = np.array(list(self._utilpairs[sid_index_out])) #!!!make more efficient with npfromiter?
        snpreader1 = self.snpreader1 if self.snpreader1 is not None else self.snpreader0
        sid_index_inner_0 = self.snpreader0.sid_to_index(pair_array[:,0])
        sid_index_inner_1 = snpreader1.sid_to_index(pair_array[:,1])


        if self.snpreader1 is None:
            sid_index_inner_01 = np.unique(np.r_[sid_index_inner_0,sid_index_inner_1]) #Index of every snp of interest
            inner_01 = self.snpreader0[iid_index_or_none,sid_index_inner_01].read(order=order, dtype=dtype, force_python_only=force_python_only, view_ok=True, num_threads=num_threads) #read every val of interest
            val_inner_01 = inner_01.standardize().val if self.do_standardize else inner_01.val

            sid_index_inner_01_reverse = {v:i for i,v in enumerate(sid_index_inner_01)} #Dictionary of snp_index to position in sid_index_inner_01
            sid_index_inner_0_in_val = np.array([sid_index_inner_01_reverse[i] for i in sid_index_inner_0])  #Replace snp_index0 with column # in val_inner_01
            sid_index_inner_1_in_val = np.array([sid_index_inner_01_reverse[i] for i in sid_index_inner_1])  #Replace snp_index1 with column # in val_inner_01
            val_inner_0 = val_inner_01[:,sid_index_inner_0_in_val] #Extract the vals for the left snps of interest
            val_inner_1 = val_inner_01[:,sid_index_inner_1_in_val]#Extract the vals for the right snps of interest
        else:
            inner_0 = self.snpreader0[iid_index_or_none,sid_index_inner_0].read(order=order, dtype=dtype, force_python_only=force_python_only, view_ok=True, num_threads=num_threads) #read every val of interest
            inner_1 = self.snpreader1[iid_index_or_none,sid_index_inner_1].read(order=order, dtype=dtype, force_python_only=force_python_only, view_ok=True, num_threads=num_threads) #read every val of interest
            val_inner_0 = inner_0.standardize().val if self.do_standardize else inner_0.val
            val_inner_1 = inner_1.standardize().val if self.do_standardize else inner_1.val
        val = val_inner_0*val_inner_1 #Element multiplication creates the vals for the pairs
        return val

#!!! keep these?

def split_on_sids(snpreader,part_count):
    sid_count = snpreader.sid_count
    start = 0
    for part_index in range(1,part_count+1):
        end = part_index*sid_count//part_count
        yield snpreader[:,start:end]
        start=end

def epi_reml(pair_snps,pheno,covar=None,kernel_snps=None,output_dir='results',part_count=33,runner=None,override=False):
    from pysnptools.kernelreader import SnpKernel
    from pysnptools.standardizer import Unit
    import datetime
    from fastlmm.association import single_snp

    part_list = list(split_on_sids(pair_snps,part_count))
    part_pair_count = (part_count*part_count+part_count)/2
    part_pair_index = -1
    print("part_pair_count={0:,}".format(part_pair_count))
    K0 = SnpKernel(kernel_snps or pair_snps,standardizer=Unit()).read() #Precompute the similarity
    if not os.path.exists(output_dir): os.makedirs(output_dir)
    start_time = datetime.datetime.now()
    for i in range(part_count):
        part_i = part_list[i]        
        for j in range(i,part_count):
            part_pair_index+=1
            pairs = _Pairs2(part_i) if i==j else _Pairs2(part_i,part_list[j])
            print("Looking at pair {0},{1} which is {2} of {3}".format(i,j,part_pair_index,part_pair_count))
            output_file = '{0}/result.{1}.{2}.tsv'.format(output_dir,part_pair_index,part_pair_count)
            if override or not os.path.exists(output_file):
                result_df_ij = single_snp(pairs, K0=K0, pheno=pheno, covar=covar, leave_out_one_chrom=False, count_A1=True,runner=runner)
                result_df_ij.to_csv(output_file, sep="\t", index=False)
                print(result_df_ij[:1])
                time_so_far = datetime.datetime.now()-start_time
                total_time_estimate = time_so_far*part_pair_count/(part_pair_index+1)
                print(total_time_estimate)


    
class TestPairs(unittest.TestCase):

    def test_run1(self):
        from pysnptools.snpreader import Bed
        root = Path(r'D:\OneDrive\programs\epireml') #!!! make this work without this
        runner = None
        #runner = LocalMultiProc(multiprocessing.cpu_count(),just_one_process=False)

        bed_original = Bed(str(root / 'syndata.bed'),count_A1=False)        
        pheno= str(root/'pheno.txt')
        covar= str(root/'cov.txt') #set to None, if none
        output_dir = r'm:/deldir/results.run1'
        part_count = 2

        epi_reml(bed_original[:,-20:],pheno,covar=covar,output_dir=output_dir,part_count=part_count,runner=runner,override=True)        
        #!!! check answer?

    def test_run2(self):
        from pysnptools.snpreader import Bed
        root = Path(r'D:\OneDrive\programs\epireml') #!!! make this work without this
        runner = None
        #runner = LocalMultiProc(multiprocessing.cpu_count(),just_one_process=False)
        bed_original = Bed(str(root / 'syndata.bed'),count_A1=False) #Read only the first 10 SNPs
        pheno= str(root / 'pheno.txt')
        covar= str(root / 'cov.txt')
        output_dir = r'm:/deldir/results.run2'
        part_count = 1

        epi_reml(bed_original[:,:20],pheno,kernel_snps=bed_original,covar=covar,output_dir=output_dir,part_count=part_count,runner=runner,override=True)

        #!!! check answer? against M:\deldir\refresults.run1 and 2

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


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

    TestPairs().test_run1()
    TestPairs().test_run2()

    suites = getTestSuite()
    r = unittest.TextTestRunner(failfast=True)
    r.run(suites)

    result = doctest.testmod()
    assert result.failed == 0, "failed doc test: " + __file__


    print("done")


    from pysnptools.snpreader import Pheno, Bed
    import pysnptools.util as pstutil

    if False:
        from pysnptools.snpreader import Bed
        from pysnptools.util import example_file # Download and return local file name
        bed_file = example_file('doc/ipynb/all.*','*.bed')
        bed = Bed(bed_file,count_A1=False)
        print(f'{bed.sid_count:,}')
        pairs = _Pairs(bed)
        print(f'{pairs.sid_count:,}')
        subset = pairs[:10,10000:10010]
        print(subset.read().val.shape)
        print(subset.sid)

    if False:

        data_file = 'd:\OneDrive\programs\epiCornell\syndata.bed'
        if False:
            from pysnptools.snpreader import SnpData
            import numpy as np
            bed1 = Bed("../../tests/datasets/synth/all")
            print(bed1.iid_count, bed1.sid_count, bed1.iid_count * bed1.sid_count)
            #goal 1500 individuals x 27000 SNP
            snpdata1 = bed1.read()
            iid = bed1.iid
            sid = ['sid{0}'.format(i) for i in xrange(27000)]
            val = np.tile(snpdata1.val,(3,6))[:,:27000].copy()
            #snpdata = Pheno('pysnptools/examples/toydata.phe').read()         # Read data from Pheno format
            snpdata2 = SnpData(iid, sid, val)
            print(snpdata2.iid_count, snpdata2.sid_count, snpdata2.iid_count * snpdata2.sid_count)
            Bed.write(snpdata2,data_file,count_A1=False)

        synbed = Bed(data_file)
        print(synbed.iid_count, synbed.sid_count, synbed.iid_count * synbed.sid_count)

        part_count = 1000
        part_list = list(split_on_sids(synbed,part_count))

        pairs00 = _Pairs(part_list[0])
        from fastlmm.association import single_snp
        pheno_fn = r"d:\OneDrive\programs\epiCornell\pheno.txt"
        cov_fn = r"d:\OneDrive\programs\epiCornell\cov.txt"
        results_df = single_snp(pairs00, K0=synbed, pheno=pheno_fn, covar=cov_fn, leave_out_one_chrom=False, count_A1=True)

        if False:
            for i,synbed_part_i in enumerate(synbed_part_list):
                for j,synbed_part_j in enumerate(synbed_part_list):
                    if j<i:
                        continue #not break
                    print("Looking at pair {0},{1}".format(i,j))
                    pairs = _Pairs(synbed_part_i) if i==j else _Pairs(synbed_part_i,synbed_part_j)
                    #print(pairs.iid)
                    print('{:,}'.format(pairs.sid_count))
                    #print(pairs.sid)
                    #print(pairs.pos)
                    #print(pairs.row_property)
                    snpdata = pairs.read()#
                    #print(snpdata.val)

        import datetime
        from pysnptools.kernelreader import SnpKernel
        from pysnptools.standardizer import Unit
        from pysnptools.util.mapreduce1.runner import LocalMultiProc
        from pysnptools.util.mapreduce1 import map_reduce
        #runner=None
        runner = LocalMultiProc(1,just_one_process=False)

        part_pair_count = (part_count*part_count+part_count)//2
        part_pair_index = -1
        print("part_pair_count={0:,}".format(part_pair_count))

        K0 = SnpKernel(synbed,standardizer=Unit()).read() #Precompute the similarity

        start_time = datetime.datetime.now()
        for i,part_i in enumerate(part_list):
            def mapper1(j):
                #from fastlmm.association import single_snp
                #from pysnptools.snpreader import Pairs
                #print('Z')
                #part_j = part_list[j]
                #print('A')
                print("Looking at pair {0},{1} which is {2} of {3}".format(i,j,part_pair_index+j+1,part_pair_count))
                #pairs = Pairs(part_i) if i==j else Pairs(part_i,part_j)
                #result_df_ij = single_snp(pairs, K0=K0, pheno=pheno_fn, covar=cov_fn, leave_out_one_chrom=False, count_A1=True)
                #print(result_df_ij[:1])
                #return result_df_ij

            result_df_i = map_reduce(range(i,part_count),
                                     mapper=mapper1,
                                     reducer=lambda result_j_list:pd.concat(result_j_list),
                                     runner=runner,
                                     name='js')
            part_pair_index+=(part_count-i)
            time_so_far = datetime.datetime.now()-start_time
            total_time_estimate = time_so_far*part_pair_count/(part_pair_index+1)
            print(total_time_estimate)

