from typing import OrderedDict
import gzip


def af_from_vcf(vcf_gz, snps_only=False):
    """Parse AF from INFO fields of LoFreq vcf's"""
    afs = OrderedDict()
    with gzip.open(vcf_gz) as fh:
        for line in fh:
            if line.decode().startswith("#"):
                continue
            sq, pos, _, ref, alt, _, _, info = line.decode().rstrip().split("\t")[:8]
            if snps_only:
                if len(ref) > 1 or len(alt) > 1:
                    continue
            pos = int(pos)-1
            af = [float(x[4:]) for x in info.split(";") if x.startswith("AF=")][0]
            if sq not in afs:
                afs[sq] = OrderedDict()
            # keep max af if multi-allelic
            af2 = afs[sq].get(pos)
            if af2 and af2 > af:
                af = af2
            afs[sq][pos] = af
    return afs