import numpy as np
from itertools import *
import pandas as pd
import logging
from pysnptools.kernelreader import KernelReader
from pysnptools.pstreader import PstData
from pysnptools.pstreader import PstReader

class Identity(KernelReader):
    '''
    A :class:`.KernelReader` for that represents an identity matrix. No memory for the values is allocated until :meth:`Identity.read` is called.

    See :class:`.KernelReader` for general examples of using KernelReaders.

    **Constructor:**
        :Parameters: * **iid** (an array of strings) -- The :attr:`KernelReader.iid` information

        :Example:

        >>> from pysnptools.kernelreader import Identity
        >>> identity = Identity(iid=[['fam0','iid0'],['fam0','iid1']])
        >>> print(identity.iid_count)
        2
        >>> print(identity.read().val) # '...' is possible space character
        [[...1.  0.]
         [...0.  1.]]

        >>> identity = Identity(iid=[['fam0','iid0'],['fam0','iid1'],['fam0','iid2']],test=[['fam0','iid1'],['fam0','iid3']])
        >>> print((identity.iid0_count, identity.iid1_count))
        (3, 2)
        >>> print(identity.read().val) # '...' is possible space character
        [[...0.  0.]
         [...1.  0.]
         [...0.  0.]]

    '''
    def __init__(self, iid, test=None): #!!! add docs and test for test
        super(Identity, self).__init__()

        if test is None:
            test = iid

        if test is iid:
            iid = PstData._fixup_input(iid,empty_creator=lambda ignore:np.empty([0,2],dtype='str'),dtype='str')
            test = iid
        else:
            iid = PstData._fixup_input(iid,empty_creator=lambda ignore:np.empty([0,2],dtype='str'),dtype='str')
            test = PstData._fixup_input(test,empty_creator=lambda ignore:np.empty([0,2],dtype='str'),dtype='str')

        if len(iid)>0:
            self._row0 = iid
        else:
            self._row0 = self._empty

        if len(test)>0:
            self._row1 = test
        else:
            self._row1 = self._empty

    _empty = np.empty([0,2],dtype='str')

    def __repr__(self): 
        return "{0}({1}x{2})".format(self.__class__.__name__, self.row_count, self.col_count)


    @property
    def row(self):
        return self._row0

    @property
    def col(self):
        return self._row1

    def _read(self, row_index_or_none, col_index_or_none, order, dtype, force_python_only, view_ok, num_threads):
        dtype = np.dtype(dtype)
        if row_index_or_none is None and col_index_or_none is None and self._row0 is self._row1: #read all of a square ID
            val = np.identity(self.row_count,dtype=dtype)
            if (order=='F'and not val.flags["F_CONTIGUOUS"]) or (order=='C'and not val.flags["C_CONTIGUOUS"]):
                val = val.T
            return val
        else: #Non-square
            #!!! This is also less efficient than it could be because it create a big identity matrix and then slices it.

            #In about O(col_count + row_count) fill in zeros
            big = np.zeros([self.row_count,self.col_count],dtype=dtype)
            common = set([PstReader._makekey(x) for x in self.row]) & set([PstReader._makekey(x) for x in self.col])
            big[self.row_to_index(common),self.col_to_index(common)] = 1.0
            val, shares_memory = self._apply_sparray_or_slice_to_val(big, row_index_or_none, col_index_or_none, order, dtype, force_python_only, num_threads)
            return val

    def __getitem__(self, iid_indexer_and_snp_indexer):
        if isinstance(iid_indexer_and_snp_indexer,tuple):
            row_index_or_none, col_index_or_none = iid_indexer_and_snp_indexer
        else:
            row_index_or_none = iid_indexer_and_snp_indexer
            col_index_or_none = row_index_or_none

        #If iid0 will be iid1 then make 'is' equal
        iid=self.iid0[row_index_or_none]
        test=self.iid1[col_index_or_none]
        if len(iid) == len(test) and np.array_equal(iid,test):
            result = Identity(iid)
        else: 
            result = Identity(iid=iid,test=test)
        return result


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

    import doctest
    doctest.testmod(optionflags=doctest.ELLIPSIS|doctest.NORMALIZE_WHITESPACE)
