#
# MIT License
#
# Copyright (c) 2019-2022 Tskit Developers
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
"""
Convert tree sequences to VCF.
"""
import numpy as np

from . import provenance


def legacy_position_transform(positions):
    """
    Transforms positions in the tree sequence into VCF coordinates under
    the pre 0.2.0 legacy rule.
    """
    last_pos = 0
    transformed = []
    for pos in positions:
        pos = int(round(pos))
        if pos <= last_pos:
            pos = last_pos + 1
        transformed.append(pos)
        last_pos = pos
    return transformed


class VcfWriter:
    """
    Writes a VCF representation of the genotypes tree sequence to a
    file-like object.
    """

    def __init__(
        self,
        tree_sequence,
        *,
        ploidy,
        contig_id,
        individuals,
        individual_names,
        position_transform,
        site_mask,
        sample_mask,
        isolated_as_missing,
    ):
        self.tree_sequence = tree_sequence
        self.contig_id = contig_id
        self.isolated_as_missing = isolated_as_missing

        if individuals is None:
            individuals = np.arange(tree_sequence.num_individuals, dtype=int)
        self.individuals = individuals

        self.__make_sample_mapping(ploidy)

        if individual_names is None:
            individual_names = [f"tsk_{j}" for j in range(self.num_individuals)]
        self.individual_names = individual_names
        if len(self.individual_names) != self.num_individuals:
            raise ValueError(
                "individual_names must have length equal to the number of individuals"
            )

        # Transform coordinates for VCF
        if position_transform is None:
            position_transform = np.round
        elif position_transform == "legacy":
            position_transform = legacy_position_transform
        self.transformed_positions = np.array(
            position_transform(tree_sequence.tables.sites.position), dtype=int
        )
        if self.transformed_positions.shape != (tree_sequence.num_sites,):
            raise ValueError(
                "Position transform must return an array of the same length"
            )
        self.contig_length = max(
            1, int(position_transform([tree_sequence.sequence_length])[0])
        )
        if len(self.transformed_positions) > 0:
            # Arguably this should be last_pos + 1, but if we hit this
            # condition the coordinate systems are all muddled up anyway
            # so it's simpler to stay with this rule that was inherited
            # from the legacy VCF output code.
            self.contig_length = max(self.transformed_positions[-1], self.contig_length)

        if site_mask is None:
            site_mask = np.zeros(tree_sequence.num_sites, dtype=bool)
        self.site_mask = np.array(site_mask, dtype=bool)
        if self.site_mask.shape != (tree_sequence.num_sites,):
            raise ValueError("Site mask must be 1D a boolean array of length num_sites")

        self.sample_mask = sample_mask
        if sample_mask is not None:
            if not callable(sample_mask):
                sample_mask = np.array(sample_mask, dtype=bool)
                self.sample_mask = lambda _: sample_mask

    def __make_sample_mapping(self, ploidy):
        """
        Compute the sample IDs for each VCF individual and the template for
        writing out genotypes.
        """
        self.samples = None
        self.individual_ploidies = []
        if len(self.individuals) > 0:
            if ploidy is not None:
                raise ValueError("Cannot specify ploidy when individuals present")
            self.samples = []
            for i in self.individuals:
                if i < 0 or i >= self.tree_sequence.num_individuals:
                    raise ValueError("Invalid individual IDs provided.")
                ind = self.tree_sequence.individual(i)
                self.samples.extend(ind.nodes)
                self.individual_ploidies.append(len(ind.nodes))
            if len(self.samples) == 0:
                raise ValueError("The individuals do not map to any sampled nodes.")
        else:
            if ploidy is None:
                ploidy = 1
            if ploidy < 1:
                raise ValueError("Ploidy must be >= 1")
            if self.tree_sequence.num_samples % ploidy != 0:
                raise ValueError("Sample size must be divisible by ploidy")
            self.individual_ploidies = [
                ploidy for _ in range(self.tree_sequence.sample_size // ploidy)
            ]
        self.num_individuals = len(self.individual_ploidies)

    def __write_header(self, output):
        print("##fileformat=VCFv4.2", file=output)
        print(f"##source=tskit {provenance.__version__}", file=output)
        print('##FILTER=<ID=PASS,Description="All filters passed">', file=output)
        print(
            f"##contig=<ID={self.contig_id},length={self.contig_length}>", file=output
        )
        print(
            '##FORMAT=<ID=GT,Number=1,Type=String,Description="Genotype">', file=output
        )
        vcf_samples = "\t".join(self.individual_names)
        print(
            "#CHROM",
            "POS",
            "ID",
            "REF",
            "ALT",
            "QUAL",
            "FILTER",
            "INFO",
            "FORMAT",
            vcf_samples,
            sep="\t",
            file=output,
        )

    def write(self, output):
        self.__write_header(output)

        # Build the array for hold the text genotype VCF data and the indexes into
        # this array for when we're updating it.
        gt_array = []
        indexes = []
        for ploidy in self.individual_ploidies:
            for _ in range(ploidy):
                indexes.append(len(gt_array))
                # First element here is a placeholder that we'll write the actual
                # genotypes into when for each variant.
                gt_array.extend([0, ord("|")])
            gt_array[-1] = ord("\t")
        gt_array[-1] = ord("\n")
        gt_array = np.array(gt_array, dtype=np.int8)
        # TODO Unclear here whether using int64 or int32 will be faster for this index
        # array. Test it out.
        indexes = np.array(indexes, dtype=int)

        for variant in self.tree_sequence.variants(
            samples=self.samples, isolated_as_missing=self.isolated_as_missing
        ):
            site_id = variant.site.id
            # We check the mask before we do any checks so we can use this as a
            # way of skipping problematic sites.
            if self.site_mask[site_id]:
                continue

            if variant.num_alleles > 9:
                raise ValueError(
                    "More than 9 alleles not currently supported. Please open an issue "
                    "on GitHub if this limitation affects you."
                )
            pos = self.transformed_positions[variant.index]
            ref = variant.alleles[0]
            alt = "."
            if variant.num_alleles > 1:
                alt = ",".join(variant.alleles[1 : variant.num_alleles])
            print(
                self.contig_id,
                pos,
                site_id,
                ref,
                alt,
                ".",
                "PASS",
                ".",
                "GT",
                sep="\t",
                end="\t",
                file=output,
            )
            genotypes = variant.genotypes
            gt_array[indexes] = genotypes + ord("0")
            if self.sample_mask is not None:
                genotypes = genotypes.copy()
                sample_mask = np.array(self.sample_mask(variant), dtype=bool)
                if sample_mask.shape != genotypes.shape:
                    raise ValueError(
                        "Sample mask must be a numpy array of size num_samples"
                    )
                genotypes[sample_mask] = -1
            if self.sample_mask is not None or variant.has_missing_data:
                missing = genotypes == -1
                gt_array[indexes[missing]] = ord(".")
            g_bytes = memoryview(gt_array).tobytes()
            g_str = g_bytes.decode()
            print(g_str, end="", file=output)
