#!/usr/bin/env python3
"""
Tool: vcf_variant_filter.py
Category: Variant Calling & VCF Analysis
Description: Stream-parses VCF v4.2 files (including gzip), filters by Depth (DP),
             Variant Quality (QUAL), and Allele Frequency (AF), computes Ti/Tv ratio,
             and exports filtered variants into a clean TSV table.
Usage:
    python vcf_variant_filter.py -i variants.vcf.gz -o filtered.tsv --min-dp 20 --min-qual 30
"""

import argparse
import gzip
import sys

def open_vcf(path):
    if path.endswith('.gz'):
        return gzip.open(path, 'rt')
    return open(path, 'rt')

def parse_info_field(info_str):
    """Parses key=value pairs from the VCF INFO column."""
    info_dict = {}
    for item in info_str.split(';'):
        if '=' in item:
            k, v = item.split('=', 1)
            info_dict[k] = v
        else:
            info_dict[item] = True
    return info_dict

def is_transition(ref, alt):
    """Returns True if single nucleotide variant is a Transition (A<->G or C<->T)."""
    transitions = {('A', 'G'), ('G', 'A'), ('C', 'T'), ('T', 'C')}
    return (ref, alt) in transitions

def filter_vcf(input_path, output_tsv, min_dp=10, min_qual=30.0, min_af=0.0):
    total_variants = 0
    passed_variants = 0
    transitions = 0
    transversions = 0
    
    out_fh = open(output_tsv, 'w') if output_tsv else sys.stdout
    out_fh.write("CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tDP\tAF\tTYPE\n")
    
    with open_vcf(input_path) as f:
        for line in f:
            if line.startswith('#'):
                continue
            cols = line.strip().split('\t')
            if len(cols) < 8:
                continue
                
            total_variants += 1
            chrom, pos, var_id, ref, alt, qual_str, filt, info_str = cols[:8]
            
            try:
                qual = float(qual_str) if qual_str != '.' else 0.0
            except ValueError:
                qual = 0.0
                
            info = parse_info_field(info_str)
            dp = int(info.get('DP', 0))
            
            af = 0.0
            if 'AF' in info:
                try:
                    af = float(info['AF'].split(',')[0])
                except ValueError:
                    af = 0.0
                    
            # Filter checks
            if qual < min_qual or dp < min_dp or af < min_af:
                continue
                
            passed_variants += 1
            var_type = "SNV" if (len(ref) == 1 and len(alt) == 1) else "INDEL"
            
            if var_type == "SNV":
                if is_transition(ref, alt):
                    transitions += 1
                else:
                    transversions += 1
                    
            out_fh.write(f"{chrom}\t{pos}\t{var_id}\t{ref}\t{alt}\t{qual:.1f}\t{filt}\t{dp}\t{af:.3f}\t{var_type}\n")
            
    if output_tsv and out_fh != sys.stdout:
        out_fh.close()
        
    ti_tv = transitions / transversions if transversions > 0 else 0.0
    return total_variants, passed_variants, ti_tv

def main():
    parser = argparse.ArgumentParser(description="VCF Variant Filter & Ti/Tv Calculator.")
    parser.add_argument("-i", "--input", help="Path to input VCF file (.vcf or .vcf.gz)")
    parser.add_argument("-o", "--output", help="Path to output TSV table")
    parser.add_argument("--min-dp", type=int, default=10, help="Minimum sequencing read depth (DP)")
    parser.add_argument("--min-qual", type=float, default=30.0, help="Minimum variant call quality (QUAL)")
    parser.add_argument("--min-af", type=float, default=0.0, help="Minimum allele frequency (AF)")
    args = parser.parse_args()

    if not args.input:
        demo = "/tmp/demo_variants.vcf"
        with open(demo, "w") as f:
            f.write("##fileformat=VCFv4.2\n")
            f.write("#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\n")
            f.write("chr1\t10020\trs001\tA\tG\t85.5\tPASS\tDP=45;AF=0.52\n")
            f.write("chr1\t10540\trs002\tC\tA\t20.0\tLowQual\tDP=5;AF=0.10\n")
            f.write("chr2\t22100\trs003\tT\tC\t99.0\tPASS\tDP=60;AF=0.98\n")
            f.write("chr2\t34000\trs004\tG\tT\t55.2\tPASS\tDP=35;AF=0.48\n")
        args.input = demo
        print(f"[!] Using demo VCF at {demo}\n")

    total, passed, ti_tv = filter_vcf(args.input, args.output, args.min_dp, args.min_qual, args.min_af)
    print(f"\n=== VCF Filtering Summary ===", file=sys.stderr)
    print(f"Total Variants : {total}", file=sys.stderr)
    print(f"Passed Variants: {passed} ({(passed/total*100):.1f}%)", file=sys.stderr)
    print(f"Ti/Tv Ratio    : {ti_tv:.3f}", file=sys.stderr)

if __name__ == "__main__":
    main()