"""Analysis of VCF file generated by freebayes.
"""
import sys
import vcf


class VCF(vcf.Reader):
    """VCF class (Variant Calling Format)

    This class is a wrapping of vcf.Reader class from the pyVCF package

    ::

        from sequana import sequana_data, VCF
        vcf_filename = sequana_data("test.vcf", "testing")

        # Read the data
        v = VCF(vcf_filename)

        # Filter the data
        filter_dict = {"QUAL": 10000, "FREQ": 0.85, 
            "INFO": {"DP": ">10", "AO": ">200", "SRP": "<100"}}
        v.filter_vcf(filter_dict)

    """
    def __init__(self, filename, **kwargs):
        """
        Filter vcf file with a dictionnary.

        :param str filename: a vcf file.
        :param kwargs: any arguments accepted by vcf.Reader
        """
        try:
            filin = open(filename, "r")
            vcf.Reader.__init__(self, fsock=filin, **kwargs)
        except IOError as e:
            print("I/O error({0}): {1}".format(e.errno, e.strerror))

    def _compute_freq(self, vcf_line):
        alt_freq = [float(count)/vcf_line.INFO["DP"] for count in \
                vcf_line.INFO["AO"]]
        return alt_freq

    def _filter_info_field(self, info_value, threshold):
        if(threshold.startswith("<")):
            if(threshold.startswith("<=")):
                if(info_value <= int(threshold[2:])):
                    return False
            elif(info_value < int(threshold[1:])):
                return False
        else:
            if(threshold.startswith(">=")):
                if(info_value >= int(threshold[2:])):
                    return False
            elif(info_value > int(threshold[1:])):
                return False
        return True

    def _filter_line(self, vcf_line, filter_dict):
        # dictionary must have QUAL/FREQ/INFO keys

        if(vcf_line.QUAL < filter_dict["QUAL"]):
            return False
        alt_freq = self._compute_freq(vcf_line)

        if(alt_freq[0] < filter_dict["FREQ"]):
            return False

        for key, value in filter_dict["INFO"].items():
            try:
                if(type(vcf_line.INFO[key]) != list):
                    if(self._filter_info_field(vcf_line.INFO[key], value)):
                        return False
                else:
                    if(self._filter_info_field(vcf_line.INFO[key][0], value)):
                        return False
            except KeyError:
                print("The information key doesn't exist in VCF file.\n")
        return True

    def filter_vcf(self, filter_dict, output):
        """Read the VCF file and write the filter vcf file.
        """
        with open(output, "w") as fp:
            vcf_writer = vcf.Writer(fp, self)
            for variant in self:
                if(self._filter_line(variant, filter_dict)):
                    vcf_writer.write_record(variant)


@@ -45,34 +45,219 @@
 es.%0A
+%0A
 %22%22%22%0A
-import sys%0Aimport vcf%0A
+%0A# Import -----------------------------------------------------------------------%0A%0Aimport sys%0Aimport vcf%0Aimport pandas as pd%0A%0A# Class ------------------------------------------------------------------------
 %0A%0Acl
@@ -2065,32 +2065,38 @@
     return False
+, None
 %0A        alt_fre
@@ -2200,16 +2200,22 @@
 rn False
+, None
 %0A%0A      
@@ -2438,32 +2438,38 @@
     return False
+, None
 %0A               
@@ -2582,32 +2582,38 @@
     return False
+, None
 %0A            exc
@@ -2657,54 +2657,398 @@
 The 
-information key doesn't exist in VCF file.%5Cn%22)
+key %7B0%7D does not exist in VCF file.%5Cn%22.format(key))%0A        line_dict = %7B%22chr%22: vcf_line.CHROM, %22position%22: vcf_line.POS,%0A                %22depth%22: vcf_line.INFO%5B%22DP%22%5D, %22reference%22: vcf_line.REF,%0A                %22alternative%22: %22; %22.join(str(x) for x in vcf_line.ALT), %0A                %22freebayes%22: vcf_line.QUAL,%0A                %22frequency%22: %22; %22.join(%22%7B0:.2f%7D%22.format(x) for x in alt_freq)%7D
 %0A   
@@ -3055,32 +3055,43 @@
      return True
+, line_dict
 %0A%0A    def filter
@@ -3185,21 +3185,185 @@
 file
-.%0A        %22%22%22
+ and return a data%0A        frame.%0A%0A        %22%22%22%0A        df = pd.DataFrame(columns=%5B%22chr%22, %22position%22, %22depth%22, %22reference%22, %0A            %22alternative%22, %22frequency%22, %22freebayes%22%5D)
 %0A   
@@ -3484,35 +3484,55 @@
                 
-if(
+keep_line, line_dict = 
 self._filter_lin
@@ -3554,18 +3554,111 @@
 er_dict)
+%0A                if keep_line:%0A                    df = df.append(line_dict, ignore_index=True
 )
-:
 %0A       
@@ -3703,9 +3703,26 @@
 ariant)%0A
+        return df
 %0A

Add report vcf