import numpy as np
import subprocess, sys, os.path
from itertools import *
import pandas as pd
import logging
from pysnptools.kernelreader import KernelReader
from pysnptools.pstreader import PstData
from pysnptools.kernelstandardizer import Identity as KS_Identity
from pysnptools.kernelstandardizer import DiagKtoN
import pysnptools.util as pstutil


class KernelData(KernelReader, PstData):
    """A :class:`.KernelReader` for holding kernel values in-memory, along with related iid information.
    It is usually created by calling the :meth:`.SnpReader.read_kernel` method on a :class:`.SnpReader`, for example, :class:`.Bed`.
    It can also be created by calling the :meth:`.KernelReader.read` method on :class:`.SnpKernel`, for example, :class:`.KernelNpz`.
    It can also be constructed directly.

    See :class:`.KernelReader` for details and examples.

    **Constructor:**
        :Parameters: * **iid** (an array of string pairs) -- The :attr:`KernelReader.iid` information.
                     * **iid0** (an array of string pairs) -- The :attr:`KernelReader.iid0` information.
                     * **iid1** (an array of strings pairs) -- The :attr:`KernelReader.iid1` information.
                     * **val** (a 2-D array of floats) -- The SNP values
                     * **name** (optional, string) -- Information to be display about the origin of this data
                     * **xp** (optional, a Python module or string) -- The array module that controls **val**,
                               for example, 'numpy' or 'cupy'. Defaults to numpy. Can also be set with the
                               'ARRAY_MODULE' environment variable.

        If *iid* is provided, don't provide *iid0* and *iid1*. Likewise, if *iid0* and *iid1* are provided, don't provide *iid*.

        :Example:

        >>> from pysnptools.kernelreader import KernelData
        >>> kerneldata = KernelData(iid=[['fam0','iid0'],['fam0','iid1']], val=[[1.,.5],[.5,1.]])
        >>> print((kerneldata.val[0,1], kerneldata.iid_count))
        (0.5, 2)

    **Equality:**

        Two KernelData objects are equal if their three arrays (:attr:`.iid0`, :attr:`.iid1`, and :attr:`KernelData.val`) are 'array_equal'.
        (Their 'string' does not need to be the same).
        If either :attr:`.KernelData.val` contains NaN, the objects will not be equal. However, :meth:`.KernelData.allclose` can be used to treat NaN as
        regular values.


        :Example:

        >>> import numpy as np
        >>> from pysnptools.kernelreader import KernelData
        >>> kerneldata1 = KernelData(iid=[['fam0','iid0'],['fam0','iid1']], val=[[1.,.5],[.5,1.]])
        >>> kerneldata2 = KernelData(iid=[['fam0','iid0'],['fam0','iid1']], val=[[1.,.5],[.5,1.]])
        >>> print(kerneldata1 == kerneldata2) #True, because all the arrays have the same values.
        True
        >>> print(kerneldata1.val is kerneldata2.val) #False, because the two arrays have different memory.
        False
        >>> kerneldata3 = KernelData(iid=[['a','0'],['b','0']], val=[[1.,.5],[.5,1.]])
        >>> kerneldata4 = KernelData(iid=[['fam0','iid0'],['fam0','iid1']], val=[[1.,.5],[.5,1.]])
        >>> print(kerneldata3 == kerneldata4) #False, because the iids are different.
        False
        >>> kerneldata5 = KernelData(iid=[['fam0','iid0'],['fam0','iid1']], val=[[1.,.5],[.5,np.nan]])
        >>> kerneldata6 = KernelData(iid=[['fam0','iid0'],['fam0','iid1']], val=[[1.,.5],[.5,np.nan]])
        >>> print(kerneldata5 == kerneldata6) #False, because the val's contain NaN
        False
        >>> print(kerneldata5.allclose(kerneldata6)) #True, if we consider the NaN as regular values, all the arrays have the same values.
        True


    **Methods beyond** :class:`.KernelReader`
    """

    def __init__(
        self,
        iid=None,
        iid0=None,
        iid1=None,
        val=None,
        name=None,
        parent_string=None,
        xp=None,
    ):  # !!!autodoc doesn't generate good doc for this constructor
        # We don't have a 'super(KernelData, self).__init__()' here because KernelData takes full responsibility for initializing both its superclasses
        xp = pstutil.array_module(xp)
        self._val = None

        # !!why does SnpData __init__ have a copy_inputs, but KernelData doesn't?
        assert (iid is None) != (
            iid0 is None and iid1 is None
        ), "Either 'iid' or both 'iid0' 'iid1' must be provided."
        assert (
            name is None or parent_string is None
        ), "Can't set both 'name' and the deprecated 'parent_string'"
        if parent_string is not None:
            warnings.warn(
                "'parent_string' is deprecated. Use 'name'", DeprecationWarning
            )

        if iid is not None:
            self._row = PstData._fixup_input(
                iid,
                empty_creator=lambda ignore: np.empty([0, 2], dtype="str"),
                dtype="str",
            )
            self._col = self._row
        else:
            self._row = PstData._fixup_input(
                iid0,
                empty_creator=lambda ignore: np.empty([0, 2], dtype="str"),
                dtype="str",
            )
            self._col = PstData._fixup_input(
                iid1,
                empty_creator=lambda ignore: np.empty([0, 2], dtype="str"),
                dtype="str",
            )
        self._row_property = PstData._fixup_input(
            None,
            count=len(self._row),
            empty_creator=lambda count: np.empty([count, 0], dtype="str"),
            dtype="str",
        )
        self._col_property = PstData._fixup_input(
            None,
            count=len(self._col),
            empty_creator=lambda count: np.empty([count, 0], dtype="str"),
            dtype="str",
        )
        self._val = PstData._fixup_input_val(
            val,
            row_count=len(self._row),
            col_count=len(self._col),
            empty_creator=lambda row_count, col_count: xp.empty(
                [row_count, col_count], dtype=xp.float64
            ),
            xp=xp,
        )
        self._assert_iid0_iid1(check_val=True)
        self._name = name or parent_string or ""
        self._std_string_list = []
        self._xp = xp

    @property
    def val(self):
        """The 2D NumPy array of floats that represents the values of the kernel. You can get or set this property.

        >>> from pysnptools.kernelreader import KernelData
        >>> kerneldata = KernelData(iid=[['fam0','iid0'],['fam0','iid1']], val=[[1.,.5],[.5,1.]])
        >>> print((kerneldata.val[0,1], kerneldata.iid_count))
        (0.5, 2)
        """
        return self._val

    @val.setter
    def val(self, new_value):
        self._val = PstData._fixup_input_val(
            new_value,
            row_count=len(self._row),
            col_count=len(self._col),
            empty_creator=lambda row_count, col_count: np.empty(
                [row_count, col_count], dtype=np.float64
            ),
            xp=self._xp,
        )
        self._assert_iid0_iid1(check_val=True)

    def allclose(self, value, equal_nan=True):
        """
        :param value: Other object with which to compare.
        :type value: :class:`KernelData`
        :param equal_nan: (Default: True) Tells if NaN in :attr:`KernelData.val` should be treated as regular values when testing equality.
        :type equal_nan: bool

        >>> import numpy as np
        >>> kerneldata5 = KernelData(iid=[['fam0','iid0'],['fam0','iid1']], val=[[1.,.5],[.5,np.nan]])
        >>> kerneldata6 = KernelData(iid=[['fam0','iid0'],['fam0','iid1']], val=[[1.,.5],[.5,np.nan]])
        >>> print(kerneldata5.allclose(kerneldata6)) #True, if we consider the NaN as regular values, all the arrays have the same values.
        True
        >>> print(kerneldata5.allclose(kerneldata6,equal_nan=False)) #False, if we consider the NaN as special values, all the arrays are not equal.
        False

        """
        return PstData.allclose(self, value, equal_nan=equal_nan)

    # !! SnpData.standardize() changes the str to help show that the data has been standardized. Should this to that too?
    def standardize(
        self,
        standardizer=DiagKtoN(),
        return_trained=False,
        force_python_only=False,
        num_threads=None,
    ):
        """Does in-place standardization of the in-memory
        kernel data. The method multiples the values with a scalar factor such that the diagonal sums to iid_count. Although it works in place, for convenience
        it also returns the KernelData.

        :rtype: :class:`.KernelData` (standardizes in place, but for convenience, returns 'self')

        >>> from pysnptools.kernelreader import KernelNpz
        >>> from pysnptools.util import example_file # Download and return local file name
        >>> import numpy as np
        >>> kernel_file = example_file('pysnptools/examples/toydata.kernel.npz')
        >>> kernel_on_disk = KernelNpz(kernel_file)
        >>> kerneldata1 = kernel_on_disk.read() # read all kernel values into memory
        >>> print(np.diag(kerneldata1.val).sum())
        5000000.0
        >>> kerneldata1.standardize() # standardize changes the values in kerneldata1.val
        KernelData(KernelNpz(...pysnptools/examples/toydata.kernel.npz'))
        >>> print(np.diag(kerneldata1.val).sum())
        500.0
        >>> kerneldata2 = kernel_on_disk.read().standardize() # Read and standardize in one expression with only one ndarray allocated.
        >>> print(np.diag(kerneldata2.val).sum())
        500.0
        """
        return standardizer.standardize(
            self,
            return_trained=return_trained,
            force_python_only=force_python_only,
            num_threads=num_threads,
        )


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

    import doctest

    doctest.testmod(optionflags=doctest.ELLIPSIS)
    # There is also a unit test case in 'pysnptools\test.py' that calls this doc test
