"""Script has as input a snp set it will find the locations of the snps and calculates the allele
frequency over the new samples of that location. This alles me to match the allel frequency of new samples with that
of my snp set and see if the results cluster correct."""

import vcf
from optparse import OptionParser

def parse_options():
    """Option parser"""
    parser = OptionParser()

    parser.add_option("-v", "--vcf", dest="vcf_file",
                      help="Raw vcf file", metavar="VCF", type="string")

    parser.add_option("-l", "--loci", dest="loci_file",
                      help="Selected Snps",
                      metavar="PAIR_FILe", type="string")

    parser.add_option("-o", "--out", dest="output",
                      help="Output name", metavar="OUT", type="string")

    opts, args = parser.parse_args()
    return opts, args

def calc_pool_freq(variant):
    """Calculates the frequency over the pooled samples, the frequency is estimated from the read depth using the
    reference allele"""
    all_pool = []

    for sample in range(0, len(variant.samples)):
        individual = variant.samples[sample]
        if individual.called:
            if individual.is_het:
                freq = individual["AD"][0] / sum(individual["AD"])
            else:
                base = individual.gt_bases.split("/")
                if base[0] == variant.REF:
                    freq = float(1)
                else:
                    freq = float(0)
        else:
            freq = "-"

        all_pool.append(freq)

    return all_pool


def main(opts):
    """Main function: input files are snp set generated by configure_dist_set.R and a vcf file the output is a
    csv file containing the calculated frequency's over the samples."""
    vcf_reader = vcf.Reader(filename=opts.vcf_file)

    write_file = []
    sample_names = vcf_reader.samples
    sample_names.insert(0, "-")


    with open(opts.loci_file, 'r') as f:
        first_line = f.readline()

    snp_index = first_line.split()
    print(snp_index)


    for snp in snp_index[1:]:

        # coordinate = snp.split(".") # for framboos
        coordinate = snp.rsplit(".", 1)
        try:
            info_snp = vcf_reader.fetch(coordinate[0][1:], int(coordinate[1][:-1]) - 1)

            variant = next(info_snp)
            frequencies = calc_pool_freq(variant)
            frequencies.insert(0, coordinate[0][1:] + ":" + str(coordinate[1][:-1]))
            write_file.append(frequencies)

        except ValueError:
            print("No information for loci: " + coordinate[0][1:] + ":" + str(coordinate[1][:-1]) + " in " + opts.vcf_file)
            print("Added empty entry")

            frequencies = [coordinate[0][1:] + ":" + str(coordinate[1][:-1])]
            for sample in vcf_reader.samples:
                frequencies.append("-")
            write_file.append(frequencies)

        # print(frequencies)

    write_file.insert(0, sample_names)
    with open(opts.output, 'w') as f:
        for row in zip(*write_file):
            f.write('\t'.join(list(map(str, row))) + '\n')

if __name__ == '__main__':
    opts, args = parse_options()
    main(opts)