"""

Generate individual files and chr22.json

Running this code need following command:

    mkdir individual

    mkdir promoters

"""

import sys

import os

import vcf

import pandas as pd

import numpy as np

import csv

import time

import json

def subfinder(mylist, pattern):

    for i in range(len(mylist)):

        if mylist[i] == pattern[0] and mylist[i:i + len(pattern)] == pattern:

            return i, i + len(pattern)

def promoter_var_idx(p_idx, promoter, var_idx):

    pmt_var = promoter[p_idx]  # which promoter

    p_start, p_end = subfinder(var_idx, pmt_var)

    return p_start, p_end

# read used ALS variants in each .vcf file

with open('var_in_file.json') as json_data:

    var_file_dic = json.load(json_data)

# read promoter table

with open('promoter1.csv', 'rb') as f:

    reader = csv.reader(f)

    promoter = list(reader)

# read used variant posotions

with open('all_ALS_var.txt','r') as f:

    var_idx = [i.replace('\n','') for i in f]

# mapping dictionary

var_num_dict = {"0/0":'0',

                "0/1":'1',

                "1/0":'1',

                "1/1":'2',

                "./.":'-1'}

# IDS:

labels_file = 'labes.csv'

labels_df = pd.read_csv(labels_file,index_col=0)

ids = labels_df.FID.tolist()

# create individual files

print 'Create individual files'

for ind in ids:

    id_file_name = 'individual/'+str(ind)+'.txt'

    id_file = open(id_file_name, 'w')

    id_file.write('')

## path to vcf files

#files = os.listdir('./chr22')

#files.sort()

#files = files[1:]

files = ['xaa.vcf']

print 'Start writing....'

num_vcf = len(files)

num_vcf_batch = len(files)*0.05

vcf_i = 0

for f_str in files:

    if vcf_i % num_vcf_batch == 0:

        print vcf_i / num_vcf*100., '%....'

    file_name = '' + f_str

    with open(file_name, 'r') as f:

        lines = [l.replace('\n', '').split("\t") for l in f if not l.startswith('##')]

    info = lines[2:]

    # create dataframe

    v_df = pd.DataFrame.from_records(info, columns=lines[0])

    use_less_l = ["ID", "REF", "ALT", "QUAL", "FILTER", "INFO", "FORMAT"]

    v_df.drop(use_less_l, axis=1, inplace=True)

    # 0/0 ==> 0 i.e.

    samples = lines[0][9:]

    replace_dict = {}

    for i in samples:

        replace_dict[i] = var_num_dict

    v_df_r = v_df.replace(to_replace=replace_dict)

    var_df = v_df_r[v_df_r['POS'].isin(var_file_dic[f_str])]

    # write to file

    for ind in ids:

        # the id in vcf file is 'LP6008192-DNA_E10_LP6008192-DNA_E10' for 'LP6008192-DNA_E10'

        new_ind = str(ind)+'_'+str(ind)

        if new_ind in var_df.columns.tolist()[2:]:

            info = var_df[new_ind].tolist()

            id_file_name = 'individual/' + str(ind) + '.txt'

            id_file = open(id_file_name, 'a')

            for item in info:

                id_file.write("%s\n" % item)

            id_file.close()

        else:

            print ind

    vcf_i += 1.

print "100.0%... Done, generated all used ALS variants for each individual"

num_pro = 0

print "Generating No.",num_pro," promoter in chr",22

p_start,p_end = promoter_var_idx(0,promoter,var_idx=var_idx)

promoter_ind = {}

for ind in ids:

    indiv_file = 'individual/'+str(ind)+'.txt'

    with open(indiv_file,'r') as f:

        ind_v = [i.replace('\n','') for i in f]

    promo = ind_v[p_start:p_end]

    promoter_ind[ind] = promo

with open('promoters/chr22.json', 'w') as fp:

    json.dump(promoter_ind, fp)