# -*- coding: utf-8 -*-
#
#  This file is part of Sequana software
#
#  Copyright (c) 2016 - Sequana Development Team
#
#  File author(s):
#      Thomas Cokelaer <thomas.cokelaer@pasteur.fr>
#      Dimitri Desvillechabrol <dimitri.desvillechabrol@pasteur.fr>,
#          <d.desvillechabrol@gmail.com>
#
#  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."""
from sequana.lazy import vcf
from sequana.lazy import pandas as pd
from sequana import logger


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._vcf_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 _vcf_line_to_dict(self, vcf_line):
        """ Convert a VCF line as a dictionnary with the most important
        information to detect real variants.
        """
        # Calcul all important information
        alt_freq = compute_frequency(vcf_line)
        strand_bal = compute_strand_balance(vcf_line)
        line_dict = {"chr": vcf_line.CHROM, "position": str(vcf_line.POS),
                     "depth": vcf_line.INFO["DP"], "reference": vcf_line.REF,
                     "alternative": "; ".join(str(x) for x in vcf_line.ALT),
                     "freebayes_score": vcf_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 = vcf_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 VCF_freebayes(vcf.Reader):
    """VCF class (Variant Calling Format)

    This class is a wrapping of vcf.Reader class from the pyVCF package. It
    is dedicated for VCF file generated by freebayes. 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.

    ::

        from sequana import sequana_data
        from sequana.freebayes_vcf_filter import VCF_freebayes
        vcf_filename = sequana_data("JB409847.vcf")

        # Read the data
        v = VCF_freebayes(vcf_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}
        filter_v = v.filter_vcf(filter_dict)
        filter_v.to_vcf('output.filter.vcf')
    """
    def __init__(self, filename, **kwargs):
        """.. rubric:: constructor 

        :param str filename: a vcf file.
        :param kwargs: any arguments accepted by vcf.Reader
        """
        try:
            self.filename = filename
            filin = open(filename, "r")
            vcf.Reader.__init__(self, fsock=filin, **kwargs)
            self._get_start_index()
        except FileNotFoundError as e:
            logger.error("FileNotFoundError({0}): {1}".format(e.errno,
                                                              e.strerror))
            raise FileNotFoundError
        # initiate filters dictionary
        self._filters_params = {'freebayes_score': 0,
                                'frequency': 0,
                                'min_depth': 0,
                                'forward_depth':0,
                                'reverse_depth':0,
                                'strand_ratio': 0}

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

        ::

            v = VCF_freebayes("input.bcf")
            v.filter_params = {"freebayes_score": 200,
                               "frequency": 0.8,
                               "min_depth": 10,
                               "forward_depth":3,
                               "reverse_depth":3,
                               "strand_ratio": 0.2}
        """
        return self._filters_params

    @filters_params.setter
    def filters_params(self, d):
        self._filters_params.update(d)

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

    def _get_start_index(self):
        self._reader.seek(0)
        for line in iter(self._reader.readline, ''):
            if line.startswith("#"):
                self._start_index = self._reader.tell()
            else:
                self.rewind()
                break

    def rewind(self):
        """ Rewind the reader
        """
        self._reader.seek(self._start_index)
        self.reader = (line.strip() for line in self._reader if line.strip())

    def filter_vcf(self, filter_dict=None):
        """ Filter variants in the VCF file.

        :param dict filter_dict: dictionary of filters. It updates the
            attribute :attr:`VCF_freebayes.filters`
        Return Filtered_freebayes object.
        """
        if filter_dict:
            self.filters_params = filter_dict
        variants = [Variant(v) for v in self if self._filter_line(v)]
        self.rewind()
        return Filtered_freebayes(variants, self)

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

        :param vcf.model._Record vcf_line

        Return line if all filters are passed.
        """

        if vcf_line.QUAL < self.filters_params["freebayes_score"]:
            return False

        if vcf_line.INFO["DP"] <= self.filters_params["min_depth"]:
            return False

        forward_depth = vcf_line.INFO["SRF"] + sum(vcf_line.INFO["SAF"])
        if forward_depth <= self.filters_params["forward_depth"]:
            return False

        reverse_depth = vcf_line.INFO["SRR"] + sum(vcf_line.INFO["SAR"])
        if reverse_depth <= self.filters_params["reverse_depth"]:
            return False

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

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

        return True


class Filtered_freebayes(object):
    """ Variants filtered with VCF_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, fb_vcf):
        """.. rubric:: constructor

        :param list variants: list of variants record.
        :param BCF_freebayes bcf: class parent.
        """
        self._variants = variants
        self._vcf = fb_vcf
        self._df = self._vcf_to_df()

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

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

    def _vcf_to_df(self):
        """ Create a data frame with the most important information contained
        in the VCF 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._vcf.filters_params), 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:
            vcf_writer = vcf.Writer(fp, self._vcf)
            for variant in self.variants:
                vcf_writer.write_record(variant.record)


def strand_ratio(number1, number2):
    """ Compute ratio between two number. Return result between [0:0.5].
    """
    try:
        division = float(number1) / (number1 + number2)
        if division > 0.5:
            division = 1 - division
    except ZeroDivisionError:
        return 0
    return division

def compute_frequency(vcf_line):
    """ Compute frequency of alternate allele.
        alt_freq = Count Alternate Allele / Depth

    :param vcf.model._Record vcf_line: variant record
    """
    alt_freq = [float(count)/vcf_line.INFO["DP"]
                for count in vcf_line.INFO["AO"]]
    return alt_freq

def compute_strand_balance(vcf_line):
    """ Compute strand balance of alternate allele include in [0,0.5].
    strand_bal = Alt Forward / (Alt Forward + Alt Reverse)

    :param vcf.model._Record vcf_line: variant record
    """
    strand_bal = [
        strand_ratio(vcf_line.INFO["SAF"][i], vcf_line.INFO["SAR"][i])
        for i in range(len(vcf_line.INFO["SAF"]))]
    return strand_bal
