#
#  This file is part of Sequana software
#
#  Copyright (c) 2016-2021 - Sequana Development Team
#
#  Distributed under the terms of the 3-clause BSD license.
#  The full license is in the LICENSE file, distributed with this software.
#
#  website: https://github.com/sequana/sequana
#  documentation: http://sequana.readthedocs.io
#
##############################################################################
"""Analysis of VCF file generated by freebayes. This method need the version
0.10 of pysam.
"""
from pysam import VariantFile
from sequana.lazy import pandas as pd

from sequana.vcftools import strand_ratio, compute_strand_balance, compute_frequency

import colorlog

logger = colorlog.getLogger(__name__)


__all__ = ["Variant"]


class Variant(object):
    """Variant class to stock variant reader and dictionary that resume most
    important informations.
    """

    def __init__(self, record):
        """.. rubric:: constructor

        :param RecordVariant record: variant record
        :param dict resume: most important informations of variant
        """
        self._record = record
        self._resume = self._bcf_line_to_dict(record)

    def __str__(self):
        return str(self.record)

    @property
    def record(self):
        return self._record

    @property
    def resume(self):
        return self._resume

    def _bcf_line_to_dict(self, bcf_line):
        """Convert a BCF line as a dictionnary with the most important
        information to detect real variants.
        """
        # Calcul all important information
        alt_freq = compute_frequency(bcf_line)
        strand_bal = compute_strand_balance(bcf_line)
        line_dict = {
            "chr": bcf_line.chrom,
            "position": str(bcf_line.pos),
            "depth": bcf_line.info["DP"],
            "reference": bcf_line.ref,
            "alternative": "; ".join(str(x) for x in bcf_line.alts),
            "freebayes_score": bcf_line.qual,
            "strand_balance": "; ".join("{0:.2f}".format(x) for x in strand_bal),
            "frequency": "; ".join("{0:.2f}".format(x) for x in alt_freq),
        }
        try:
            # If bcf is annotated by snpEff
            annotation = bcf_line.info["EFF"][0].split("|")
            effect_type, effect_lvl = annotation[0].split("(")
            try:
                prot_effect, cds_effect = annotation[3].split("/")
            except ValueError:
                cds_effect = annotation[3]
                prot_effect = ""
            ann_dict = {
                "CDS_position": cds_effect[2:],
                "effect_type": effect_type,
                "codon_change": annotation[2],
                "gene_name": annotation[5],
                "mutation_type": annotation[1],
                "prot_effect": prot_effect[2:],
                "prot_size": annotation[4],
                "effect_impact": effect_lvl,
            }
            line_dict = dict(line_dict, **ann_dict)
        except KeyError:
            pass
        return line_dict


class BCF_freebayes(VariantFile):
    """BCF_freebayes class (Binary Variant Calling Format)

    This class is a wrapping of VariantFile class from the pysam package. It
    is dedicated for VCF file generated by freebayes and compressed by
    bcftools. BCF file is faster to parse than VCF. A data frame with all
    variants is produced which can be write as csv file. It can filter variants
    with a dictionnary of filter parameter. Filter variants are wrotte in a new
    VCF file.

    Example:

    ::

        from sequana import BCF_freebayes

        # Read the data
        b = BCF_freebayes(bcf_filename)

        # Filter the data
        filter_dict = {'freebayes_score': 200,
                       'frequency': 0.8,
                       'min_depth': 10,
                       'forward_depth':3,
                       'reverse_depth':3,
                       'strand_ratio': 0.2}
        b.filter_bcf(filter_dict, "output.vcf")
    """

    def __init__(self, input_filename, **kwargs):
        """
        :param str filename: a bcf file.
        :param kwargs: any arguments accepted by VariantFile.
        """
        try:
            super().__init__(input_filename, **kwargs)
        except OSError:
            logger.error("OSError: {0} doesn't exist.".format(input_filename))
            raise OSError
        # initiate filters dictionary
        self._filters = {
            "freebayes_score": 0,
            "frequency": 0,
            "min_depth": 0,
            "forward_depth": 0,
            "reverse_depth": 0,
            "strand_ratio": 0,
        }

    @property
    def filters(self):
        """Get or set the filters parameters to select variants of interest.
        Setter take a dictionnary as parameter to update the attribute
        :attr:`BCF_freebayes.filters`. Delete will reset different variable to
        0.

        ::

            bcf = BCF_freebayes("input.bcf")
            bcf.filter = {"freebayes_score": 200,
                          "frequency": 0.8,
                          "min_depth": 10,
                          "forward_depth":3,
                          "reverse_depth":3,
                          "strand_ratio": 0.2}
        """
        return self._filters

    @filters.setter
    def filters(self, d):
        self._filters.update(d)

    @filters.deleter
    def filters(self):
        self._filters = {
            "freebayes_score": 0,
            "frequency": 0,
            "min_depth": 0,
            "forward_depth": 0,
            "reverse_depth": 0,
            "strand_ratio": 0,
        }

    def filter_bcf(self, filter_dict=None):
        """Filter variants in the BCF file and write them in a BCF file.

        :param str output_filename: BCF output filename.
        :param dict filter_dict: dictionary of filters. It updates the
            attribute :attr:`BCF_freebayes.filters`
        Return BCF_freebayes object with new BCF file.
        """
        if filter_dict:
            self.filters = filter_dict
        variants = [Variant(v) for v in self if self._filter_line(v)]
        # Rewind the iterator
        self.reset()
        return Filtered_freebayes(variants, self)

    def _filter_line(self, bcf_line):
        """Filter variant with parameter set in :attr:`BCF_freebayes.filters`.

        :param VariantRecord bcf_line

        Return line if all filters are passed.
        """
        # SRF="Number of reference observations on the forward strand
        # SRR="Number of reference observations on the reverse strand
        # SAF="Number of alternate observations on the forward strand
        # SAR=Number of alternate observations on the reverse strand
        if bcf_line.qual < self.filters["freebayes_score"]:
            return False

        if bcf_line.info["DP"] <= self.filters["min_depth"]:
            return False

        forward_depth = bcf_line.info["SRF"] + sum(bcf_line.info["SAF"])
        if forward_depth <= self.filters["forward_depth"]:
            return False

        reverse_depth = bcf_line.info["SRR"] + sum(bcf_line.info["SAR"])
        if reverse_depth <= self.filters["reverse_depth"]:
            return False

        alt_freq = compute_frequency(bcf_line)
        if alt_freq[0] < self.filters["frequency"]:
            return False

        strand_bal = compute_strand_balance(bcf_line)
        if strand_bal[0] < self.filters["strand_ratio"]:
            return False

        return True


class Filtered_freebayes(object):
    """Variants filtered with BCF_freebayes."""

    _col_index = [
        "chr",
        "position",
        "reference",
        "alternative",
        "depth",
        "frequency",
        "strand_balance",
        "freebayes_score",
        "effect_type",
        "mutation_type",
        "effect_impact",
        "gene_name",
        "CDS_position",
        "codon_change",
        "prot_effect",
        "prot_size",
    ]

    def __init__(self, variants, bcf):
        """.. rubric:: constructor

        :param list variants: list of variants record.
        :param BCF_freebayes bcf: class parent.
        """
        self._variants = variants
        self._bcf = bcf
        self._df = self._bcf_to_df()

    @property
    def variants(self):
        """Get the variant list."""
        return self._variants

    @property
    def df(self):
        return self._df

    def _bcf_to_df(self):
        """Create a data frame with the most important information contained
        in the bcf file.
        """
        dict_list = [v.resume for v in self.variants]
        df = pd.DataFrame.from_records(dict_list)
        try:
            df = df[Filtered_freebayes._col_index]
        except (ValueError, KeyError):
            df = df[Filtered_freebayes._col_index[: len(df.columns)]]
        return df

    def to_csv(self, output_filename):
        """Write DataFrame in CSV format.

        :params str output_filename: output CSV filename.
        """
        with open(output_filename, "w") as fp:
            print("# sequana_variant_calling; {0}".format(self._bcf.filters), file=fp)
            if self.df.empty:
                print(",".join(Filtered_freebayes._col_index), file=fp)
            else:
                self.df.to_csv(fp, index=False)

    def to_vcf(self, output_filename):
        """Write BCF file in VCF format.

        :params str output_filename: output VCF filename.
        """
        with open(output_filename, "w") as fp:
            print(self._bcf.header, end="", file=fp)
            for variant in self.variants:
                print(variant, end="", file=fp)
