import optparse
import sys
from cyvcf2 import VCF
import numpy as np

import random

def cyvcf2msms(vcf_filename, chrom, sample_size, start, length, sample_list=None):

    # set new sample list
    vcf = VCF(vcf_filename)

    if sample_list == None:
        samples = vcf.samples
    else:
        samples = sample_list

    sample_set = set()

    # retrieve a sample from the total sample list
    if sample_size < len(samples):
        for i in range(sample_size):
            chosen = random.choice(samples)
            while chosen in sample_set:
                chosen = random.choice(samples)
            sample_set.add(chosen)
    else:
        print("Sample size bigger than total number of samples")
        return -1

    sample_list = list(sample_set)
    vcf.set_samples(sample_list)

    reformatted_hap_list = []
    reformatted_seg_site_idx_list = []

    sample_size = len(sample_list)
    pos = start
    for record in vcf(str(chrom) + ":" + str(start) + "-" + str(start+length)):
        hap = []
        for genotype in record.genotypes:
            hap.append(int(genotype[0]))
            # hap.append(int(genotype[1]))
        if sum(hap) == 0 or sum(hap) == sample_size:
            pass
        else:
            pos = record.POS
            if sample_size - sum(hap) < sum(hap):
                for idx, snp in enumerate(hap):
                    hap[idx] = 1 - snp
            reformatted_hap_list.append(hap)
            reformatted_seg_site_idx_list.append((pos-start))

    reformatted_hap_list = np.transpose(np.array(reformatted_hap_list))
    print(reformatted_hap_list.shape)
    reformatted_hap_list = reformatted_hap_list.tolist()

    assert(len(reformatted_hap_list) == sample_size)
    assert(len(reformatted_hap_list[0]) == len(reformatted_seg_site_idx_list))

    # write msms file
    msms_filename = '/scratch/saralab/first/real0.msms'
    msms_file = open(msms_filename, 'w')
    msms_file.write('fasta2msms.py Nov2014 -n ' + str(len(reformatted_hap_list)) + '\n')
    msms_file.write(str(chrom) + ', block: ' + str(start) + '-' + str(start+length))
    msms_file.write('\n\n//\nsegsites: ' + str(len(reformatted_seg_site_idx_list)) + '\npositions:')
    msms_file.write(' '.join(['%.5f'%(float(s)/length) for s in reformatted_seg_site_idx_list]) + '\n')
    for hap in reformatted_hap_list:
        hap = ''.join([str(x) for x in hap])
        msms_file.write(hap + '\n')
    
    msms_file.close()

def main():
    cyvcf2msms(vcf_filename='/scratch/hlee6/vcf/ALL.chr21.vcf.gz', chrom=21, sample_size=20, start=39988628, length=10e4, sample_list=None)

if __name__ == '__main__':
    main()
