from insightspy.SessionCore import RequestCore
from insightspy.utils import _list_to_int_list, _download
import pandas as pd
import os


class SamplesMixin(RequestCore):
    def samples(self, tag_suffix="_tag_dupl", sample_ids=None):
        """List samples

        Lists samples accessible withthin the current portal session. Sample access is
        limited by user id and the current session project if one is set.

        Args:
            tag_suffix (str):  the suffix that is appended to the column name if a tag
                shares the same name as one of the columns in the core database used to
                describe the samples
            sample_ids (int, [int]):  one or more sample ids to filter on (optional).
                This will also override any filtering from the project scope.

        Examples:
            >>> # p is a logged in portal session
            >>> p.samples()
            >>> # To only retrieve metadata about specific samples
            >>> p.samples(sample_ids = [32, 35, 42])

        Returns:
            DataFrame: table of sample metadata
        """
        if sample_ids is not None:
            ignore_project = True
            if not isinstance(sample_ids, list):
                sample_ids = [sample_ids]
            sample_ids = _list_to_int_list(sample_ids)
        else:
            ignore_project = False
        # List keys to export to user
        export_keys = [
            "sample_id",
            "description",
            "sample_type",
            "reference_genome",
            "revision_id",
            "tags",
        ]
        out = [
            {key: v[key] for key in export_keys}
            for v in self._post(
                "sample/retrieve",
                {
                    "tagsAsDict": True,
                    "sample_ids": sample_ids,
                    "ignoreProject": ignore_project,
                },
                expect_data=True,
            )["response"]["data"]
        ]
        out = pd.DataFrame.from_dict(out)
        if "tags" in out:
            out = out.join(
                pd.DataFrame(out.pop("tags").values.tolist()), rsuffix=tag_suffix
            )
        return out

    def gene_tpm(self, sample_ids):
        """Get gene TPMs

        Get transcripts per million (TPMs) values for all pre-quantified genes in a
            sample

        Args:
            sample_ids (int, [int]):  one or more sample ids

        Examples:
            >>> # p is a logged in portal session
            >>> p.gene_tpm([33, 48])

        Returns:
            DataFrame:
                table of gene TPMs with sample ids as column names and gene names as row
                names
        """
        if not isinstance(sample_ids, list):
            sample_ids = [sample_ids]
        sample_ids = _list_to_int_list(sample_ids)
        counts = self._post(
            "sample/TPM",
            {
                "sample_ids": sample_ids,
                "minimal": True,
                "raw_counts": False,
                "pseudocounts": 0,
            },
            expect_data=True,
        )["response"]["data"]
        out = pd.DataFrame.from_dict(counts[0], orient="index").sort_index()
        out.columns = [str(x["sample_id"]) for x in counts[1]]
        return out

    def gene_count(self, sample_ids):
        """Get gene TPMs

        Get read counts for all pre-quantified genes in a sample

        Args:
            sample_ids (int, list):  one or more sample ids

        Examples:
            >>> # p is a logged in portal session
            >>> p.gene_count([33, 48])

        Returns:
            DataFrame:
                table of read counts with sample ids as column names and gene names
                as row names
        """
        if not isinstance(sample_ids, list):
            sample_ids = [sample_ids]
        sample_ids = _list_to_int_list(sample_ids)
        counts = self._post(
            "sample/TPM",
            {
                "sample_ids": sample_ids,
                "minimal": True,
                "raw_counts": True,
                "pseudocounts": 0,
            },
            expect_data=True,
        )["response"]["data"]
        out = pd.DataFrame.from_dict(counts[0], orient="index").sort_index()
        out.columns = [str(x["sample_id"]) for x in counts[1]]
        return out

    def locus_density(self, sample_ids, locus):
        """Get locus density profile

        Get binned coverage counts for a locus

        Args:
            sample_ids (int, [int]):  one or more sample ids
            locus(str): locus in the form "chr:start-stop"

        Examples:
            >>> # p is a logged in portal session
            >>> p.locus_density(33, "chr1:1000000-1100000")

        Returns:
            dict(DataFrame):
                dictionary of DataFrames containing RPK normalized mean signal
                per bin with one value per strand. One entry in the dictionary
                per sample being queried.
        """
        if not isinstance(sample_ids, list):
            sample_ids = [sample_ids]
        sample_ids = _list_to_int_list(sample_ids)
        density = self._post(
            "density/locus",
            {"sample_ids": sample_ids, "locus": locus},
            expect_data=True,
        )["response"]["data"]
        out = {}
        # Iterate over all returned values (2 per sample, one for each strand) and
        # place them in a dictionary of DataFrames where the key is the sample id
        for l in density:
            id = str(l["ID"])
            x = pd.DataFrame.from_dict(l["values"])
            x["strand"] = l["strand"]
            if id not in out:
                out[id] = x
            else:
                out[id] = pd.concat([out[id], x])
        for k in out.keys():
            out[k].columns = ["bin_center", "value", "strand"]
        return out

    def download_fastq(self, sample_ids, target_directory="."):
        """Download fastq

        Download the fastq file(s) associated with the sample id(s)

        Args:
            sample_ids (int, [int]):  one or more sample ids
            target_directory(str): path to directory where files should be downloaded
                to. Will be created if it does not exist

        Examples:
            >>> # p is a logged in portal session
            >>> p.download_fastq(33, "./my_project/fastq")

        """
        if not isinstance(sample_ids, list):
            sample_ids = [sample_ids]
        sample_ids = _list_to_int_list(sample_ids)
        os.makedirs(target_directory, exist_ok=True)
        types = ["fastq", "fastq2"]
        urls = {}
        for i in sample_ids:
            for t in types:
                filesize = self._post(f"fileSize/sample/{t}/{i}")["response"]["data"]
                if filesize > 0:
                    read_suffix = "R1" if t == "fastq" else "R2"
                    filepath = os.path.join(
                        target_directory, f"{i}_{read_suffix}.fastq.gz"
                    )
                    urls[filepath] = self._post(f"download/sample/{t}/{i}")["response"][
                        "data"
                    ]
        _download(urls)

    def download_bam(self, sample_ids, target_directory="."):
        """Download bam

        Download the bam file(s) associated with the sample id(s)

        Args:
            sample_ids (int, [int]):  one or more sample ids
            target_directory(str): path to directory where files should be downloaded
                to. Will be created if it does not exist

        Examples:
            >>> # p is a logged in portal session
            >>> p.download_bam(33, "./my_project/bam")

        """
        if not isinstance(sample_ids, list):
            sample_ids = [sample_ids]
        sample_ids = _list_to_int_list(sample_ids)
        os.makedirs(target_directory, exist_ok=True)
        urls = {}
        for i in sample_ids:
            filesize = self._post(f"fileSize/sample/bam/{i}")["response"]["data"]
            if filesize > 0:
                filepath = os.path.join(target_directory, f"{i}.bam")
                urls[filepath] = self._post(f"download/sample/bam/{i}")["response"][
                    "data"
                ]
        _download(urls)

    def download_bigwig(self, sample_ids, target_directory="."):
        """Download bam

        Download the bigwig file(s) associated with the sample id(s)

        Args:
            sample_ids (int, [int]):  one or more sample ids
            target_directory(str): path to directory where files should be downloaded
                to. Will be created if it does not exist

        Examples:
            >>> # p is a logged in portal session
            >>> p.download_bigwig(33, "./my_project/bigwig")

        """
        if not isinstance(sample_ids, list):
            sample_ids = [sample_ids]
        sample_ids = _list_to_int_list(sample_ids)
        os.makedirs(target_directory, exist_ok=True)
        # Dictionary of the file types to use for both stranded and unstranded data
        suffix_dict = {
            "stranded": {
                "bigwigPos": "pos.normalized.bw",
                "bigwigNeg": "neg.normalized.bw",
            },
            "unstranded": {"bigwig": "normalized.bw"},
        }
        # set of the stranded datatypes
        stranded_datatypes = set(["proseq", "PRO-seq", "rnaseq", "RNA-seq"])
        # get sample info
        sample_info = self.samples(sample_ids=sample_ids).set_index("sample_id")
        urls = {}
        for i in sample_ids:
            filesize = self._post(f"fileSize/sample/bigwig/{i}")["response"]["data"]
            is_stranded = sample_info.loc[i].sample_type in stranded_datatypes
            str_key = "stranded" if is_stranded else "unstranded"
            for t in suffix_dict[str_key].keys():
                if filesize > 0:
                    suffix = suffix_dict[str_key][t]
                    filepath = os.path.join(target_directory, f"{i}.{suffix}")
                    urls[filepath] = self._post(f"download/sample/{t}/{i}")["response"][
                        "data"
                    ]
        _download(urls)
