import pyslim, msprime
import numpy as np
import random
import time
import re
import ast
import sys
import os
#NODE_IS_SAMPLE = np.array(msprime.NODE_IS_SAMPLE, dtype='uint32')

#Time how long it takes to read in file
#start = time.time()

#seed=sys.argv[1][37:50]
found = re.search(r'TSR_files/.+_Gen',sys.argv[1])
seed = found[0][10:23]
seed_round = round(seed/10000000)
print("seed = ",seed)

start = time.time()
T2 = pyslim.load(sys.argv[1],legacy_metadata=True)
end = time.time()
print("Time it took to load in file:", end - start)

print(f"The tree sequence has {T2.num_trees} trees on a genome of length {T2.sequence_length},"
      f" {T2.num_individuals} individuals, {T2.num_samples} 'sample' genomes,"
      f" and {T2.num_mutations} mutations.")

for t in np.unique(T2.individual_times):
  print(f"There are {np.sum(T2.individual_times == t)} individuals from time {t}.")

#Is recapitation necessary? How many trees have more than one root?
sum([t.num_roots == 1 for t in T2.trees()])

#Open the migration matrix (generated by SLiM code)
with open('results/SLiM_output/Mig_mat/' + str(seed) + '_mig_mat.txt') as infile:
    data = infile.read()

#Save migration matrix as 'mig_mat'
mig_mat = []
for i in re.findall(r"(\[.*?\])", data, flags=re.S):
    mig_mat.append(ast.literal_eval(i))

#Time how long it takes to recapitate the TSR file
start = time.time()
recapT2 = T2.recapitate(recombination_rate = 1e-06, migration_matrix=mig_mat, Ne=100,random_seed=seed_round) #Ne is multiplied by the number of populations in the file
end = time.time()
print("Time it took to to recapitate:", end - start)

#Add mutations to the recapitated tree
mutatedT2 = pyslim.SlimTreeSequence(msprime.mutate(recapT2, rate=1e-7, random_seed=seed_round, keep=True))

#Save recapitated and mutated file 
#mutatedT2.dump("TTT_SP_100_TSR_1581160262936_recapped_mutated.tree")

mutatedT2.dump(sys.argv[1][0:50]+"_recapped_mutated.tree")

print(f"The tree sequence now has {mutatedT2.num_trees} trees,"
      f" and {mutatedT2.num_mutations} mutations.")

#Look at the times at which individuals were recorded (in terms of time from end of simulation, i.e. end of simulation is time "0"
set(mutatedT2.individual_times)

#Save individuals at the two time points as a vector
T1ind = mutatedT2.individuals_alive_at(300)
T2ind = mutatedT2.individuals_alive_at(0)


indivlist_T1 = []
indivnames_T1 = []


with open("results/SLiM_output/VCF_files/T1_"+str(seed)+"_sim_individuals.txt", "w") as indfile:
    indfile.writelines("\t".join(["vcf_label", "tskit_id", "slim_id"] + ["birth_time_ago", "age"]) + "\n")
    for i in T1ind:
        indivlist_T1.append(i)
        ind = T2.individual(i)
        vcf_label = f"{ind.metadata.pedigree_id}"
        indivnames_T1.append(vcf_label)
        data = [vcf_label, str(ind.id), str(ind.metadata.pedigree_id), str(ind.time),str(ind.metadata.age)]
        indfile.writelines("\t".join(data) + "\n")


with open("results/SLiM_output/VCF_files/"+str(seed)+"_subset_T1.txt", "w") as indfile:
    indfile.writelines("vcf_label" + "\n")
    Start = np.arange(0,10001,100)
    Stop = np.arange(100,10001,100)
    for i in range(0, (len(Start)-1)):
      samp = random.sample(indivnames_T1[Start[i]:Stop[i]],10)
      indfile.writelines("\n".join(samp) + "\n")


with open("results/SLiM_output/VCF_files/T1_"+str(seed)+".vcf", "w") as vcffile:
  mutatedT2.write_vcf(vcffile, individuals=indivlist_T1, individual_names=indivnames_T1)


indivlist_T2 = []
indivnames_T2 = []


with open("results/SLiM_output/VCF_files/T2_"+str(seed)+"_sim_individuals.txt", "w") as indfile:
    indfile.writelines("\t".join(["vcf_label", "tskit_id", "slim_id"] + ["birth_time_ago", "age"]) + "\n")
    for i in T2ind:
        indivlist_T2.append(i)
        ind = T2.individual(i)
        vcf_label = f"{ind.metadata.pedigree_id}"
        indivnames_T2.append(vcf_label)
        data = [vcf_label, str(ind.id), str(ind.metadata.pedigree_id), str(ind.time),str(ind.metadata.age)]
        indfile.writelines("\t".join(data) + "\n")
     

with open("results/SLiM_output/VCF_files/"+str(seed)+"_subset_T2.txt", "w") as indfile:
    indfile.writelines("vcf_label" + "\n")
    Start = np.arange(0,10001,100)
    Stop = np.arange(100,10001,100)
    for i in range(0, (len(Start)-1)):
      samp = random.sample(indivnames_T2[Start[i]:Stop[i]],10)
      indfile.writelines("\n".join(samp) + "\n")


#Write out the genomes of the individuals at each time point as VCF files
with open("results/SLiM_output/VCF_files/T2_"+str(seed)+".vcf", "w") as vcffile:
  mutatedT2.write_vcf(vcffile, individuals=indivlist_T2, individual_names=indivnames_T2)


os.system("vcftools --vcf "+'results/SLiM_output/VCF_files/T1_'+str(seed)+'.vcf'+" --maf 0.01 --recode --recode-INFO-all --out "'results/SLiM_output/VCF_files/T1_'+str(seed))

os.system("vcftools --vcf "+'results/SLiM_output/VCF_files/T2_'+str(seed)+'.vcf'+" --maf 0.01 --recode --recode-INFO-all --out "'results/SLiM_output/VCF_files/T2_'+str(seed))

os.system("vcftools --vcf "+'results/SLiM_output/VCF_files/T1_'+str(seed)+'.vcf'+" --keep "+'results/SLiM_output/VCF_files/'+str(seed)+'_subset_T1.txt'" --recode --recode-INFO-all --out "'results/SLiM_output/VCF_files/T1_'+str(seed)+'_unfiltered_subset')

os.system("vcftools --vcf "+'results/SLiM_output/VCF_files/T2_'+str(seed)+'.vcf'+" --keep "+'results/SLiM_output/VCF_files/'+str(seed)+'_subset_T2.txt'" --recode --recode-INFO-all --out "'results/SLiM_output/VCF_files/T2_'+str(seed)+'_unfiltered_subset')

os.system("vcftools --vcf "+'results/SLiM_output/VCF_files/T1_'+str(seed)+'.recode.vcf'+" --keep "+'results/SLiM_output/VCF_files/'+str(seed)+'_subset_T1.txt'" --recode --recode-INFO-all --out "'results/SLiM_output/VCF_files/T1_'+str(seed)+'_filtered_subset')

os.system("vcftools --vcf "+'results/SLiM_output/VCF_files/T2_'+str(seed)+'.recode.vcf'+" --keep "+'results/SLiM_output/VCF_files/'+str(seed)+'_subset_T2.txt'" --recode --recode-INFO-all --out "'results/SLiM_output/VCF_files/T2_'+str(seed)+'_filtered_subset')
