import os
import math
import random
import time
curdata = CURDATA()
print '**************'
for items in curdata:
	print items
MSG("A tool to sparse Bruker parameter sets by Aleks Gutmanas and Murthy Karra")
MSG("Before running this program, make sure that all parameters have been set as in a regular experiment.\
\nMake sure that the total experiment time (Bruker command expt) is approxomately 3 times the desired time.\
\nFor example if you want to run the experiment for 2 days, the time shown by Bruker should be ~6 days")
cwd=os.getcwd()


td1 = int(GETPAR("TD",1))
td2 = int(GETPAR("TD",2))
fastCplx=td2/2
slowCplx=td1/2

nucSlow=GETPAR("NUC1",1)
nucFast=GETPAR("NUC1",2)

if (nucFast=="15N"):
	T2spFast=50.0
elif (nucFast=="13C"):
	T2spFast=20.0
else:
	T2spFast=35.0

if (nucSlow=="15N"):
	T2spSlow=50.0
elif (nucSlow=="13C"):
	T2spSlow=20.0
else:
	T2spSlow=35.0

now=time.localtime()
year=int(now[0])%100
month=int(now[1])
day=int(now[2])
tday=day+100*month+10000*year
if (year<10):
	sday="0"+str(tday)
else:
	sday=str(tday)

datadir=curdata[3]+"/data/"+curdata[4]+"/nmr/"+curdata[0]+"/"+curdata[1]

try:
	ft=open(cwd+"/vdstd",'r')
	vddir=ft.readline()
	vddir=vddir.strip()
	success=1
	ft.close()
except:
	vddir="/opt/topspin/exp/stan/nmr/lists/vd"
	success=0

input = INPUT_DIALOG("Path to vdlist folder","Please modify if needed",["VDSTD"],[vddir])


if (vddir!=input[0].strip() or success==0):
	vddir=input[0].strip()
	message="Do you wish to write the new path into the standard file?\n"+vddir
	result = CONFIRM("Standard vdlist folder",message)
	if (result==1):
		ft=open(cwd+"/vdstd",'w')
		ft.write(vddir+"\n")
		ft.close()

vdname="v"+sday+"_"+curdata[1]
sname="schedule"+curdata[1]


input = INPUT_DIALOG("Parameters","Please enter the following",\
["Level of Sparsing (>30), %", "T2sp for Slow (F1, "+nucSlow+") dimension, ms",\
"T2sp for Fast (F2, "+nucFast+") dimension, ms","Random seed"],\
["30",str(T2spSlow),str(T2spFast),"709"])

sparse=0.01*float(input[0])
TspSlow=float(input[1])
TspFast=float(input[2])
rseed=int(input[3])

if (sparse<0.3) :
	MSG("Sparse level too low. Must be above 30. No VDLIST for you...")
	EXIT()

total=int(slowCplx*fastCplx*sparse)
if (total/(slowCplx*fastCplx) < sparse):
	total=total+1
	sparse=float(total)/(slowCplx*fastCplx)
	print "Sparse level recalculated to "+str(sparse*100)

print "Total number of hypercomplex points that will be recorded "+str(total)
swhSlow = float(GETPAR("SWH",1))
print swhSlow
swhFast = float(GETPAR("SWH",2))
print swhFast

P=0.0
tol=0.4
nminSlow=sparse*tol*slowCplx
if (nminSlow < 3):
	nminSlow=3

nminFast=sparse*tol*fastCplx
if (nminFast < 3):
	nminFast=3

for k in range(slowCplx):
	for j in range(fastCplx):
		P=P+math.exp(-float(k)*1000.0/(swhSlow*TspSlow)-float(j)*1000.0/(swhFast*TspFast))

P=P/(1.0*slowCplx*fastCplx)
if (P<sparse):
	MSG("Consider increasing one or both of the T2 values as the distribution will be truncated. Proceeding anyway...")

Pf=sparse/P
Pff=1.0
ntot=2*total
random.seed(rseed)
checkSlow=[]
checkFast=[]
sel=[]

for k in range(slowCplx):
	checkSlow.append(0)

for j in range(fastCplx):
	checkFast.append(0)

for k in range(slowCplx):
	sel.append([])
	for j in range(fastCplx):
		sel[k].append(0)

ncyc=0
while (ntot > total) or (ncyc<10):
	Pff=Pff*0.98
	ntot=0
	ncyc=ncyc+1
	for k in range(slowCplx):
		checkSlow[k]=0
	for j in range(fastCplx):
		checkFast[j]=0
	for k in range(slowCplx):
		for j in range(fastCplx):
			temp=Pf*Pff*math.exp(-k*1000.0/(swhSlow*TspSlow)-j*1000.0/(swhFast*TspFast))
			if (k==0 or j==0 or (k==slowCplx-1 and j==fastCplx-1) or (random.random() < temp)):
				sel[k][j]=1
				ntot=ntot+1
				checkSlow[k]=checkSlow[k]+1
				checkFast[j]=checkFast[j]+1
			else:
				sel[k][j]=0
	for k in range(slowCplx):
		while (checkSlow[k]<nminFast):
			j=random.randint(0,fastCplx-1)
			if (sel[k][j]==0):
				sel[k][j]=1
				ntot=ntot+1
				checkSlow[k]=checkSlow[k]+1
				checkFast[j]=checkFast[j]+1
	for j in range(fastCplx):
		while (checkFast[j]<nminSlow):
			k=random.randint(0,slowCplx-1)
			if (sel[k][j]==0):
				sel[k][j]=1
				ntot=ntot+1
				checkSlow[k]=checkSlow[k]+1
				checkFast[j]=checkFast[j]+1
	if (ntot>total):
		print "Too many points... Reducing probabilities and repeating the selection"

while (ntot < total):
	k=random.randint(0,slowCplx-1)
	j=random.randint(0,fastCplx-1)
	if (sel[k][j]==0):
		sel[k][j]=1
		ntot=ntot+1

fv=open(datadir+"/"+vdname,'w')
ft=open(vddir+"/"+vdname,'w')
fs=open(datadir+"/"+sname,'w')

for k in range(slowCplx):
	for j in range(fastCplx):
		if (sel[k][j]==1):
			fv.write("3u\n")
			ft.write("3u\n")
			fs.write("\t"+str(k)+"\t"+str(j)+"\n")
		else:
			fv.write("1u\n")
			ft.write("1u\n")
	for j in range(fastCplx):
		if (sel[k][j]==1):
			fv.write("3u\n")
			ft.write("3u\n")
		else:
			fv.write("1u\n")
			ft.write("1u\n")
			
fv.close()
fs.close()
ft.close()

XCMD("vdlist "+vdname)

MSG("The vdlist for sparse acquisition is created.\
\nFile "+vdname+" is copied to folder "+vddir+" and parameter vdlist is set accordingly.\
\nTotal number of FIDs to be recorded is 4*"+str(ntot)+" ("+str(4*ntot)+") out of 4*"+str(fastCplx)+"*"+str(slowCplx)+" ("+str(td1*td2)+")\
\nThe above file should have a total of "+str(2*fastCplx*slowCplx)+" records. Please check...\
\n The total experimental time will be approximately "+str(0.1*int(1000*sparse))+"% of the one reported by Bruker.")