import numpy as np
import logging
from pysnptools.snpreader import SnpReader
from pysnptools.snpreader import SnpData
import warnings
from pysnptools.pstreader import PstData
from pysnptools.pstreader import _OneShot


def zero_family(s):
    """Given an input id from the file, returns 0 as the family id and that input id as the individual id."""
    return "0", s


def zero_pos(s):
    """Given an input id from the file, return the input id as the sid and returns (0,0,0) as the chromosome, genetic distance, basepair distance."""
    return 0, s, 0, 0


def just_case(iid):
    """Use just the individual id (ignoring the family id) in the file."""
    return iid[1]


def just_sid(sid, pos):
    """Use just the sid (ignoring the chromosome, genetic distance, and basepair distance) in the file."""
    return sid


class Dense(_OneShot, SnpReader):
    """
    A :class:`.SnpReader` for reading \*.dense.txt (0,1,2 text files) from disk.

    See :class:`.SnpReader` for general examples of using SnpReaders.

    This format looks like::

        var	4006	9570    22899	37676	41236	41978	55159	66828...
        1-10004-rs12354060	22222222...
        1-707348-rs12184279	222222?2...
        1-724325-rs12564807	00000000...
        ...

    where rows are sids and columns are iids.

    **Constructor:**
        :Parameters: * **filename** (*string*) -- The Dense file to read.
                     * **extract_iid_function** (*string*) -- A function that breaks the row id from the file into a family id and an individual id. Defaults to
                       setting the family id to 0 and using the whole row id as the iid.
                     * **extract_sid_pos_function** (*string*) -- A function that breaks the column id from the file into (chromosome, sid, genetic distance, basepair distance).
                       Defaults to setting the :attr:`SnpReader.pos` information to 0 and using the whole column id as the sid.

        :Example:

        >>> from pysnptools.snpreader import Dense
        >>> from pysnptools.util import example_file # Download and return local file name
        >>> dense_file = example_file("pysnptools/examples/toydata100.dense.txt")
        >>> data_on_disk = Dense(dense_file)
        >>> print((data_on_disk.iid_count, data_on_disk.sid_count))
        (500, 100)

    **Methods beyond** :class:`.SnpReader`

    """

    def __init__(
        self,
        filename,
        extract_iid_function=zero_family,
        extract_sid_pos_function=zero_pos,
    ):
        """
        filename    : string of the name of the Dat file.
        """
        super(Dense, self).__init__()
        self.filename = str(filename)
        self.extract_iid_function = extract_iid_function
        self.extract_sid_pos_function = extract_sid_pos_function

    def _read_pstdata(self):
        bim_list = []
        val_list_list = []
        with open(self.filename, "r") as fp:
            header = fp.readline()
            iid_string_list = header.strip().split()[1:]
            iid = np.array(
                [
                    self.extract_iid_function(iid_string)
                    for iid_string in iid_string_list
                ],
                dtype="str",
            )
            val_list = []
            zerofloat = float("0"[0])
            missing_char = "?"[0]
            for line_index, line in enumerate(fp):
                if line_index % 1000 == 0:
                    logging.info(
                        "reading sid and iid info from line {0} of file '{1}'".format(
                            line_index, self.filename
                        )
                    )
                sid_string_rest = line.strip().split()
                sid_string = sid_string_rest[0]
                rest = [] if len(sid_string_rest) == 1 else sid_string_rest[1]
                assert len(rest) == len(iid)
                bim_list.append(self.extract_sid_pos_function(sid_string))
                val_list = np.array(
                    [
                        float(val) - zerofloat if val != missing_char else np.NaN
                        for val in rest
                    ]
                )
                val_list_list.append(val_list)

        col = np.array([bim[1] for bim in bim_list], dtype="str")
        col_property = np.array(
            [[bim[0], bim[2], bim[3]] for bim in bim_list], dtype=np.float64
        )

        val = np.zeros((len(iid), len(col)))
        for col_index in range(len(col)):
            val[:, col_index] = val_list_list[col_index]

        return PstData(iid, col, val, col_property=col_property, name=self.filename)

    @staticmethod
    def write(
        filename, snpdata, join_iid_function=just_case, join_sid_pos_function=just_sid
    ):
        """Writes a :class:`SnpData` to Dense format. The values must be 0,1,2 (or missing). Returns a :class:`.Dense` with the default extractors.

        :param filename: the name of the file to create
        :type filename: string
        :param snpdata: The in-memory data that should be written to disk.
        :type snpdata: :class:`SnpData`
        :param join_iid_function: function to turn a family id and individual id into a file id for columns.
                                  Defaults ignoring the family id and using the individual id as the column id.
        :type join_iid_function: a function
        :param join_sid_pos_function: function to turn a sid and pos data into a file id for rows.
                                      Defaults ignoring the :attr:`SnpReader.pos` information and using the sid id as the row id.
        :type join_sid_pos_function: a function
        :rtype: :class:`.Dense`

        >>> from pysnptools.snpreader import Dense, Bed
        >>> import pysnptools.util as pstutil
        >>> from pysnptools.util import example_file # Download and return local file name
        >>> bed_file = example_file("pysnptools/examples/toydata.5chrom.*","*.bed")
        >>> snpdata = Bed(bed_file,count_A1=False)[:,:10].read()  # Read first 10 snps from Bed format
        >>> pstutil.create_directory_if_necessary("tempdir/toydata10.dense.txt")
        >>> Dense.write("tempdir/toydata10.dense.txt",snpdata)        # Write data in Dense format
        Dense('tempdir/toydata10.dense.txt')
        """

        if isinstance(filename, SnpData) and isinstance(
            snpdata, "str"
        ):  # For backwards compatibility, reverse inputs if necessary
            warnings.warn(
                "write statement should have filename before data to write",
                DeprecationWarning,
            )
            filename, snpdata = snpdata, filename

        snpsarray = snpdata.val
        with open(filename, "w") as filepointer:
            filepointer.write("var\t")
            filepointer.write(
                "\t".join((join_iid_function(iid_pair) for iid_pair in snpdata.iid))
                + "\n"
            )

            for sid_index, sid in enumerate(snpdata.sid):
                pos = snpdata.pos[sid_index]
                if sid_index % 1000 == 0:
                    logging.info(
                        "Writing snp # {0} to file '{1}'".format(sid_index, filename)
                    )
                filepointer.write("%s\t" % join_sid_pos_function(sid, pos))
                row = snpsarray[:, sid_index]
                filepointer.write(
                    "".join((str(int(i)) if i == i else "?" for i in row)) + "\n"
                )
        logging.info("Done writing " + filename)
        return Dense(filename)


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

    import doctest

    doctest.testmod(optionflags=doctest.ELLIPSIS)
