#!/usr/bin/env python
import os, subprocess, sys, math
import numpy as np
import matplotlib.pyplot as plt

work_DIR = sys.argv[1]
genoFile = sys.argv[2]
nQTL = int(sys.argv[3])
heritability = float(sys.argv[4])
model = int(sys.argv[5])
os.chdir(work_DIR)

if model == 1:
	#################################################
	# MODEL 1: additive effects alone
	# y = Xb + e; e~N(0, Ve); Ve=Vg(1/1-h^2); Vg=sum(cor(Xq)bqbq')/4
	#################################################
	GEN = np.genfromtxt(genoFile, delimiter='\t', skip_header=1)
	nLOCI = GEN.shape[0]
	nIND = GEN.shape[1]
	QTL_locations = np.random.choice(range(0, nLOCI), replace=False, size=nQTL)
	QTL_locations.sort()
	# DEFINING THE DISTRIBUTIONS OF THE EFFECTS:
	mean_QTL = 100/(2*nQTL); var_QTL = 2	#normal QTL effects
	#QTL effects:
	QTL_effects = np.random.normal(mean_QTL, np.sqrt(var_QTL), size=nQTL) #for a mean QTL effect of ~5 and ~mean phenotypic value of 50
	QTL_OUT = np.column_stack((QTL_locations, QTL_effects)) #for writing out

	##########################################

	X=np.transpose(GEN)
	GFX=QTL_effects
	nFX=nQTL
	h2=heritability
	#partionning the variance taking into account the linkage among associated SNPs
	AssoX=X[:,QTL_locations]
	Rho=np.corrcoef(AssoX,rowvar=0)
	XtX=GFX.reshape(1,nFX)*GFX.reshape(nFX,1) #GFX * GFX' (nFXxnFX dimensions)
	Vg=np.sum(Rho*XtX)/4
	Ve=Vg*(1/h2-1)
	#Generating the phenotypes based on the variance components
	Xb=np.matmul(AssoX, GFX) #alpha
	e=np.random.normal(0,Ve**(0.5),nIND)
	Y_model1=Xb+e
	#OUTPUT
	np.savetxt("Simulated_Lolium_perenne_QTL.out", QTL_OUT, fmt='%s' ,delimiter="\t")
	np.savetxt("Simulated_Lolium_perenne_PHENOTYPES.data", Y_model1, fmt='%s' ,delimiter="\t")

elif model == 2:
	#################################################		&#### FIX ME!!!! Y_model2 is giving me the middle finger! ;-P
	# MODEL 2: additive genetic effects + transcript levels
	# y = Xg + Ba + Zt + e; e~N(0,Ve); Ve = (Vg+Vt)(1/1-h^2)
	#################################################
	transBase = sys.argv[6]
	transGeno = sys.argv[7]

	GEN = np.genfromtxt(genoFile, delimiter='\t', skip_header=1)
	T_base = np.genfromtxt(transBase, delimiter='\t', dtype=int)
	T_geno = np.genfromtxt(transGeno, delimiter='\t', dtype=int)
	nLOCI = GEN.shape[0]
	nIND = GEN.shape[1]
	nTRANS = len(T_base)
	QTL_locations = np.random.choice(range(0, nLOCI), replace=False, size=nQTL)
	QTL_locations.sort()
	# DEFINING THE DISTRIBUTIONS OF THE EFFECTS:
	mean_QTL = 100/(2*nQTL); var_QTL = 2	#normal QTL effects
	mean_bT = (mean_QTL /4); var_bT = 1		#normal base trancript level effects
	mean_gT = (mean_QTL /2); var_gT = 1		#normal genotype-specific trancript level effecs
	#QTL effects:
	QTL_effects = np.random.normal(mean_QTL, np.sqrt(var_QTL), size=nQTL) #for a mean QTL effect of ~5 and ~mean phenotypic value of 50
	QTL_OUT = np.column_stack((QTL_locations, QTL_effects)) #for writing out
	#Transcript Base-levels effects:
	nCausalTrans = int(np.ceil(np.random.normal(nQTL, 1, size=1))[0]) #number of transcripts that affect the phenotype
	locCausalTrans = np.random.choice(nTRANS, size=nCausalTrans, replace=False)
	#	(i.)base-level effects:
	T_base_effects = np.random.normal(mean_bT, np.sqrt(var_bT), size=nCausalTrans)
	#	(ii.)genotype-specific level effects:
	T_geno_effects = np.random.normal(mean_gT, np.sqrt(var_gT), size=nCausalTrans)

	##########################################

	X=np.transpose(GEN)
	GFX=QTL_effects
	nFX=nQTL
	h2=heritability

	T0 = T_base
	T1 = T_geno
	t0FX = np.zeros((nTRANS,1)); t0FX[locCausalTrans,0] = T_base_effects
	t1FX = np.zeros((nTRANS,1)); t1FX[locCausalTrans,0] = T_geno_effects

	#variance partitioning adding Vg with Vt--> variation due to genotype-specific transcripts abundance
	AssoZ=T1[:,locCausalTrans]
	Rho=np.corrcoef(AssoZ,rowvar=0)
	ZtZ=T_geno_effects.reshape(1,nCausalTrans)*T_geno_effects.reshape(nCausalTrans,1)
	Vt=np.sum(Rho*ZtZ)/4
	Vet=(Vg+Vt)*(1/h2-1)
	#generating the phenotypes using the new residual distribution et:Vet
	Xg = np.matmul(X[:,QTL_locations], GFX)
	Ba = np.sum(T0[locCausalTrans]*T_base_effects)
	Zt = np.matmul(T1, t1FX)
	et=np.random.normal(0,Vet**(0.5),nIND)
	Y_model2 = Xg + Ba + Zt[:,0] + et
	#OUTPUT
	np.savetxt("Simulated_Lolium_perenne_QTL.out", QTL_OUT, fmt='%s' ,delimiter="\t")
	np.savetxt("Simulated_Lolium_perenne_PHENOTYPES.data", Y_model1, fmt='%s' ,delimiter="\t")