import matplotlib.pyplot as plt
import numpy as np
from qutip import *
from lintrap_sim_hamiltonian import *
# set magnetic field range. 
voltagelist_to_currentdriver = np.arange(9.0,-0.25,-0.25) 
voltagelist_to_currentdriver[-1] = 0.02 # instead of zero-Bfield we add a small field to avoid some coherent phenomena that are occur when all transitions are degenerated
volt_to_B = 1.604 # Gauss/V
bfieldlist = voltagelist_to_currentdriver * volt_to_B # in Gauss
# set trap depth
voltagelist_to_trappowerstab = np.array( [1.3, 3.3] )
volt_to_trap = 1/1.47 # mk/V
trapdepth_list = voltagelist_to_trappowerstab * volt_to_trap
#set time steps for sim etc
psi0 = basis(12,0)
tsteps = 10000
tfinal = 1000e-6
tlist = np.linspace(0.0, tfinal, tsteps)
tstep=tlist[1]-tlist[0]

# norminal power
A_sigma_m = 0.0515 * gamma1   #
max_ph = (A_sigma_m**2/4)/(gamma1**2/4+A_sigma_m**2/2)*gamma1*tfinal
print('max photons: '+ str(max_ph))

# set freq range
freqlist = np.linspace(0, 39, 80)* 1e6 * 2 * np.pi

emission_spec = getTxMatrix(psi0, tlist, A_sigma_m, trapdepth_list[0], bfieldlist, freqlist)
spec = emission_spec / max_ph

k=0
h=0
for V in voltagelist_to_currentdriver:
	outfile = open(format(V, '05.2f') + 'V_tmp_sim_' + format(trapdepth_list[0], '0.2f') + 'mK_trapdepth_fine.dat','w')
	for f in freqlist:
		for item in [f, spec[k,h]]:
			outfile.write("%s " % item)	
		outfile.write("\n")
		k=k+1
	h=h+1
	k=0
	outfile.close()	
  
####### second trap depth
# set freq range
freqlist = np.linspace(44, 83, 80)* 1e6 * 2 * np.pi

emission_spec = getTxMatrix(psi0, tlist, A_sigma_m, trapdepth_list[1], bfieldlist, freqlist)
spec = emission_spec / max_ph

k=0
h=0
for V in voltagelist_to_currentdriver:
	outfile = open(format(V, '05.2f') + 'V_tmp_sim_' + format(trapdepth_list[1], '0.2f') + 'mK_trapdepth_fine.dat','w')
	for f in freqlist:
		for item in [f, spec[k,h]]:
			outfile.write("%s " % item)	
		outfile.write("\n")
		k=k+1
	h=h+1
	k=0
	outfile.close()