#!/home/asiimwe/anaconda3/bin/python
#Running STAR with no variants- no wasp output
#Script creates sample directories and copies executable BaseCode files (/BaseCode_STAR_Runs.sh) and/edits to respective sample and thread directories
import sys
import csv
import os
import string
import subprocess
from itertools import chain
import pandas as pd
from os import mkdir
import shutil
import glob
import fileinput
import stat

samples_path = "/scratch/asiimwe/STAR-WASP_FASTQs_VCFs/FASTQ/"

base_dir = "/home/asiimwe/projects/run_env/alpha_star_wasp_benchmarking/STAR"
os.chdir(base_dir)
star_dir = "STAR_Runs"
shutil.rmtree(star_dir)
mkdir(star_dir)
os.chdir(star_dir)
base_code = base_dir + "/BaseCode_STAR_Runs.sh"

sub_directories = ["8threads", "16threads", "32threads"]

for path, dirs, files in os.walk(samples_path):
	for dir in dirs:
		#print(dir) #sample dirs
		for sub_directory in sub_directories:
			os.makedirs(os.path.join(dir, sub_directory))
			full_path = os.path.join(base_dir, star_dir, dir, sub_directory)
			#print(full_path) #/home/asiimwe/projects/run_env/alpha_star_wasp_benchmarking/STAR/STAR_Runs/HG00512/8threads

			for i in os.walk(full_path):
				if i[0].split("/")[-1] == "8threads":
					#print(i[0])
					thread_dir = i[0]
					line = thread_dir + "/8threads_resource_log.txt"
					#print(line)
					subprocess.call(["touch", line]) #include sample_id
					shutil.copy(base_code, thread_dir)
					run_code = thread_dir + "/BaseCode_STAR_Runs.sh"
					st = os.stat(run_code)
					os.chmod(run_code, st.st_mode | stat.S_IEXEC)
					#y=i[0].split("_")
					#print("this is y ", y)
					#print(i[0].split("_")[5])
					
					if i[0].split("_")[5] == "Runs/NA12878":
						#path.split("_")[:8] 
						#print(i[0])
						#print()
						xthreads = 8
						sample_id = i[0].split("/")[8]
						sample_snp_dir = "NA12878" #we want to maintain the same snp dir for all NA12878 samples to avoid having multiple copies of the same as 1 vcf files applies to all NA12878 associated/derived samples
						#print(sample_id)
						#print(xthreads)
						base_code_path = i[0] + "/BaseCode_STAR_Runs.sh"
							
						#Replacing sections of base code to represent sample specific paths and threads being run
						with fileinput.FileInput(base_code_path, inplace=True) as file: #alternative argument to create a file backup: backup='.bak'
							for line in file:
								print(line.replace("xthreads", "8threads"), end='')
						with fileinput.FileInput(base_code_path, inplace=True) as file:
							for line in file:
								print(line.replace("runThreadN x", "runThreadN 8"), end='')
						with fileinput.FileInput(base_code_path, inplace=True) as file:
							for line in file:
								print(line.replace("sample_id", sample_id), end='')
						with fileinput.FileInput(base_code_path, inplace=True) as file:
							for line in file:
								print(line.replace("ssd", sample_snp_dir), end='')

					else:
						#print(i[0].split("/")[8]) #extracting other sample IDs
						sample_id = i[0].split("/")[8]
						#print(sample_id)
						base_code_path = i[0] + "/BaseCode_STAR_Runs.sh"

						with fileinput.FileInput(base_code_path, inplace=True) as file:
							for line in file:
								print(line.replace("xthreads", "8threads"), end='')
						with fileinput.FileInput(base_code_path, inplace=True) as file:
							for line in file:
								print(line.replace("runThreadN x", "runThreadN 8"), end='')
						with fileinput.FileInput(base_code_path, inplace=True) as file:
							for line in file:
								print(line.replace("sample_id", sample_id), end='')
						with fileinput.FileInput(base_code_path, inplace=True) as file:
							for line in file:
								print(line.replace("ssd", sample_id), end='')

				elif i[0].split("/")[-1] == "16threads":
					#print(i[0])
					thread_dir = i[0]
					line = thread_dir + "/16threads_resource_log.txt"
					#print(line)
					subprocess.call(["touch", line])
					shutil.copy(base_code, thread_dir)
					run_code = thread_dir + "/BaseCode_STAR_Runs.sh"
					st = os.stat(run_code)
					os.chmod(run_code, st.st_mode | stat.S_IEXEC)

	
					if i[0].split("_")[5] == "Runs/NA12878":
						#path.split("_")[:8] 6th item
						#print(i[0])
						#print()
						xthreads = 16
						sample_id = i[0].split("/")[8]
						sample_snp_dir = "NA12878" #we want to maintain the same snp dir for all NA12878 samples to avoid having multiple copies of the same as 1 vcf files applies to all NA12878 associated/derived samples
						#print(sample_id)
						#print(xthreads)
						base_code_path = i[0] + "/BaseCode_STAR_Runs.sh"
						
						#Replacing sections of base code to represent sample specific paths and threads being run
						with fileinput.FileInput(base_code_path, inplace=True) as file:
							for line in file:
								print(line.replace("xthreads", "16threads"), end='')
						with fileinput.FileInput(base_code_path, inplace=True) as file:
							for line in file:
								print(line.replace("runThreadN x", "runThreadN 16"), end='')
						with fileinput.FileInput(base_code_path, inplace=True) as file:
							for line in file:
								print(line.replace("sample_id", sample_id), end='')
						with fileinput.FileInput(base_code_path, inplace=True) as file:
							for line in file:
								print(line.replace("ssd", sample_snp_dir), end='')

					else:
						sample_id = i[0].split("/")[8]
						base_code_path = i[0] + "/BaseCode_STAR_Runs.sh"

						with fileinput.FileInput(base_code_path, inplace=True) as file:
							for line in file:
								print(line.replace("xthreads", "16threads"), end='')
						with fileinput.FileInput(base_code_path, inplace=True) as file:
							for line in file:
								print(line.replace("runThreadN x", "runThreadN 16"), end='')
						with fileinput.FileInput(base_code_path, inplace=True) as file:
							for line in file:
								print(line.replace("sample_id", sample_id), end='')
						with fileinput.FileInput(base_code_path, inplace=True) as file:
							for line in file:
								print(line.replace("ssd", sample_id), end='')


				elif i[0].split("/")[-1] == "32threads":
					#print(i[0])
					thread_dir = i[0]
					line = thread_dir + "/32threads_resource_log.txt"
					#print(line)
					subprocess.call(["touch", line])
					shutil.copy(base_code, thread_dir)
					run_code = thread_dir + "/BaseCode_STAR_Runs.sh"
					st = os.stat(run_code)
					os.chmod(run_code, st.st_mode | stat.S_IEXEC)


					if i[0].split("_")[5] == "Runs/NA12878":
						#path.split("_")[:8] 6th item
						#print(i[0])
						#print()
						xthreads = 32
						sample_id = i[0].split("/")[8]
						sample_snp_dir = "NA12878" #we want to maintain the same snp dir for all NA12878 samples to avoid having multiple copies of the same as 1 vcf files applies to all NA12878 associated/derived samples
						#print(sample_id)
						#print(xthreads)
						base_code_path = i[0] + "/BaseCode_STAR_Runs.sh"
							
						#Replacing sections of base code to represent sample specific paths and threads being run
						with fileinput.FileInput(base_code_path, inplace=True) as file:
							for line in file:
								print(line.replace("xthreads", "32threads"), end='')
						with fileinput.FileInput(base_code_path, inplace=True) as file:
							for line in file:
								print(line.replace("runThreadN x", "runThreadN 32"), end='')
						with fileinput.FileInput(base_code_path, inplace=True) as file:
							for line in file:
								print(line.replace("sample_id", sample_id), end='')
						with fileinput.FileInput(base_code_path, inplace=True) as file:
							for line in file:
								print(line.replace("ssd", sample_snp_dir), end='')

					else:
						sample_id = i[0].split("/")[8]
						base_code_path = i[0] + "/BaseCode_STAR_Runs.sh"
	
						with fileinput.FileInput(base_code_path, inplace=True) as file:
							for line in file:
								print(line.replace("xthreads", "32threads"), end='')
						with fileinput.FileInput(base_code_path, inplace=True) as file:
							for line in file:
								print(line.replace("runThreadN x", "runThreadN 32"), end='')
						with fileinput.FileInput(base_code_path, inplace=True) as file:
							for line in file:
								print(line.replace("sample_id", sample_id), end='')
						with fileinput.FileInput(base_code_path, inplace=True) as file:
							for line in file:
								print(line.replace("ssd", sample_id), end='')