#!/usr/bin/python3
# -*- coding: Utf-8 -*-
import random 
import argparse 
import re
import os

parser = argparse.ArgumentParser(description="Generates a TE File.")
parser.add_argument("-TE",help="TE_Description generated by JDD.R")
parser.add_argument("-Assoc",help="Assoc file generated by JDD.R.")
parser.add_argument("-out",help="Name of the output.")
args = parser.parse_args()


def pick_nt(GC): 
	"""
	Pick a random nucleotide depending of GC count.
	"""
	
	RGC = random.random()
	if GC < RGC :
		if random.randint(0,1) == 0 :
			return ("A")
		else :
			return ("T")
	else :
		if random.randint(0,1) == 0 :
			return ("G")
		else :
			return ("C")
	
class Famille(object):
	"""
	Family object :
		- CopyNumber
		- Name
		- Length of the longest TE
		- Sequence of this TE
	"""
	def __init__(self,Copy,Name,Max):
		# While initiating a TE family : Creates the sequence of the longest TE in the family.
		self.Copy = Copy
		self.Name=Name
		self.Max=int(Max)
	
	
		seq=[]
		
		for nt in range(Max):
			nt = pick_nt(0.5)
			seq.append(nt)
		self.seq="".join(seq)

	def __str__(self):
		return (str(len(self.seq)))

class TE(object):
	"""
	TE object :
		- Length
		- Diversity
		- Family name
		- Copy Number of it's family
		- Assoc file from JDD.R
		- Longest TE of the family
	"""
	def __init__(self,TEL,Div,Fam,Assos,First):
		self.TEL=int(TEL)
		self.Div=float(Div)
		self.Fam=Fam
		self.Copy=0
		self.Assos=Assos
		self.First=First
		
		# Find the copy Number of it's family is the Assoc file :
		
		As = open(self.Assos,"r")
		for line in As :
			Fami = line.split("\t")
			
			if Fami[0].replace("\n","") == self.Fam :
				
				self.Copy = int(Fami[1])
		
	def __str__(self):
		"""
		Print function called : Length / Diversity / Name / Copy Number
		"""
		return(str(self.TEL) +'\t'+ str(self.Div) +'\t'+ str(self.Fam) +'\t'+ str(self.Copy))
	
	def cree_seq(self):
		"""
		Function to create the sequence of the TE.
		"""
		
		seq=self.First.seq # Retrieve the sequence of the longest TE in the family.
		
		position = random.randint(0,self.First.Max - self.TEL) # Retrieve a starting position for the TE we are studying now.
		
		Myseq = seq[position:position+self.TEL] # The sequence corresponds to the longest TE of the family from the position picked plus the length of the TE created at the moment.
		
		# Now we need to apply the divergence :
		Divseq=[]
		Pos=0

		for nt in Myseq:
			Pos+=1
			Div = self.Div/100
			Seed = random.random()
			if Seed > Div : # If the random number is above the diversity threshold, the nucleotide is picked randomly.
				act=nt			
			else :
				act=pick_nt(0.5)
			
			Divseq.append(act)
			if Pos % 80 == 0 :
				Divseq.append("\n")
		self.MyNewseq = "".join(Divseq) # Replace the not diverged sequence by the new one.
	
	def write_et(self,out,FBti):
		"""
		Write the TE in the TE file.
		"""
		if os.path.isfile(out) :
			typ = "a"
			
		else: 
			typ="w"
		fastaout = open(out,typ)
		
		fastaout.write(">FBti" + str(FBti)+ " type=transposable_element; loc=3R:0...0; name="+str(self.Fam)+"{}; dbxref=Created_by_TE_dataset; MD5=NoMD5; length="+str(self.TEL)+"; div="+str(self.Div)+"; nbcop="+str(self.Copy)+";\n")
		if len(self.MyNewseq.replace("\n",""))%80 != 0 :
			fastaout.write(self.MyNewseq+"\n")
		else :
			fastaout.write(self.MyNewseq)

def main():
	Fami = open(args.Assoc,"r")
	Fams = {}
	
	# Among the TE families :
	for line in Fami :
		Fam = line.split("\t")
		if Fam[0] != "Lite_Fam":
			# Add to the dictionary {FamilyName : Famille(Freq, Name, Max length)}
			Fams[Fam[0].replace("\n","")]=Famille(Fam[1],Fam[0].replace("\n",""),int(Fam[2].replace("\n","")))
	# We created N=50 Family stored in a dictionary.
	
	# Open the TEs File :
	ETs = open(args.TE,"r")
	FBti=0
	for ET in ETs :
		
		line = ET.split("\t")
		if line[0] != "Diversity" :
			# For each TE :
			FBti+=1
			Di = float(line[0])
			Si = int(line[1])
			Fa = line[2].replace("\n","")
			
			# Create a TE object
			Premier =TE(Si,Di,Fa,args.Assoc,Fams[Fa])
			# Create it's sequence
			Premier.cree_seq()
			# Write in the output
			Premier.write_et(args.out,FBti)

if __name__ == "__main__":
    main()
				
		
		
