#!/usr/bin/env python3

"""
FOREST Thomas (thomas.forest@college-de-france.fr)

Caution : At the moment for gzipped files only.

ARGS
--------

usage : vcf_to_sfs.py VCF.gz nb_indiv

"""

import gzip
import sys

# default folded SFS
folded = True
diploid = True
phased = False

# PARAM : Nb of indiv
n = int(sys.argv[2])

if diploid and not folded:
    n *= 2
# initiate SFS_values with a zeros dict
SFS_values = dict.fromkeys(range(n),0)

polycount = 0

with gzip.open(sys.argv[1], "rb") as inputgz:
    line = inputgz.readline()
    genotypes = []
    while line:
        # decode gzipped binary lines
        line = line.decode('utf-8').strip()
        # every snp line, not comment or header
        if not line.startswith("##") and not line.startswith("#"):
            FIELDS = line.split("\t")
            # REF is col 4 of VCF
            REF = FIELDS[3].split(",")
            # ALT is col 5 of VCF
            ALT = FIELDS[4].split(",")            
            FORMAT = line.split("\t")[8:9]
            SAMPLES = line.split("\t")[9:]
            snp_genotypes = []
            allele_counts = {}
            allele_counts_list = []
            # SKIP the SNP if :
            # 1 : missing
            # 2 : deletion among REF
            # 3 : deletion among ALT
            if "./.:." in line \
               or len(ALT[0]) > 1 \
               or len(REF[0]) > 1:
                line = inputgz.readline()
                continue
            for sample in SAMPLES:
                if not phased:
                    # for UNPHASED data
                    smpl_genotype = [int(a) for a in sample.split(':')[0].split('/') if a != '.']
                else:
                    # for PHASED
                    smpl_genotype = [int(a) for a in sample.split(':')[0].split('|') if a != '.']
                nb_alleles = set(smpl_genotype)
                snp_genotypes += smpl_genotype
            # if set(snp_genotypes) > 2:
            #     polyallelic = set(snp_genotypes)
            # else:
            #     polyallelic = False
            polyallelic = len(ALT)
            ##print(REF, ALT, snp_genotypes)
            # skip if all individuals have the same genotype
            if len(set(snp_genotypes)) == 1:
                line = inputgz.readline()
                continue
            for k in set(snp_genotypes):
                allele_counts[snp_genotypes.count(k)] = k
                allele_counts_list.append(snp_genotypes.count(k))
            if folded :
                #allele_counts_list = list(allele_counts.keys())
                ##print("ALC", allele_counts_list, "POLY", polyallelic, ALT)
                # for al in range(polyallelic-1):
                #     SFS_values[min(allele_counts_list)-1] += 1/len(ALT)
                #     allele_counts_list.remove(min(allele_counts_list))
                if len(ALT) == 1:
                    SFS_values[min(allele_counts_list)-1] += 1
                else:
                    for al in range(polyallelic-1):
                        SFS_values[min(allele_counts_list)-1] += 1/len(ALT)
                        allele_counts_list.remove(min(allele_counts_list))
                        polycount += 1
        line = inputgz.readline()
        print(SFS_values)
print(polycount)