text:
"""
    RotationFrameDiagonal(omega::Array{Float64},omega_drive::Float64)

# Argument:
- omega: transition frequencies
- omega_drive: driving frequency
- N: number of states

# Output:
- Diagonal matrix diag{0,omega[1]-omr,omega[2]-omr,...,omega[N-1]-omr}
"""
function RotationFrameDiagonal(omega::Array{Float64},omega_drive::Float64)
	diag_vec = 0.0
	for i = 1:length(omega)
		# potential troublemaker of auto-diff
		diag_vec = [diag_vec;diag_vec[end]+omega[i]-omega_drive]
	end
	HK = Diagonal(diag_vec)
    return HK
end

function RotationFrameDiagonal(omega,omega_drive)
	diag_vec = 0.0
	for i = 1:length(omega)
		# potential troublemaker of auto-diff
		diag_vec = [diag_vec;diag_vec[end]+omega[i]-omega_drive]
	end
	HK = Diagonal(diag_vec)
    return HK
end

"""
    RotationFrameDiagonal(gamma1::Array{Float64},gamma2::Array{Float64})

# Argument:
- gamma1: determine the decay operator in the Lindblad system
- gamma2: determine the dephasing operator in the Lindblad system

# Output:
- Ldecay: Ldecay[i][i+1] = sqrt(gamma1[i]), and zeros everywhere else
- Ldephase: Diagonal{0,sqrt(gamma2[1]),...)
"""
function RotationFrameLindblad(gamma1::Array{Float64},gamma2::Array{Float64})
	#L_decay = Bidiagonal(zeros(length(gamma1)+1),sqrt.(gamma1),:U)
	L_dephase = Diagonal([0;sqrt.(gamma2)])
	L_decay = [zeros(Float64,length(gamma1)) Diagonal(sqrt.(gamma1))]
	L_decay = [L_decay;zeros(Float64,1,length(gamma1)+1)]
	#L_decay = [0.0 sqrt(gamma1[1]) 0.0 0.0;
	#		   0.0 0.0      sqrt(gamma1[2]) 0.0;
	#		   0.0 0.0      0.0      sqrt(gamma1[3]);
	#		   0.0 0.0      0.0      0.0]
    return L_decay, L_dephase
end

function RotationFrameLindblad(gamma1,gamma2)
	#L_decay = Bidiagonal(zeros(length(gamma1)+1),sqrt.(gamma1),:U)
	L_dephase = Diagonal([0;sqrt.(gamma2)])
	L_decay = [zeros(Float64,length(gamma1)) Diagonal(sqrt.(gamma1))]
	L_decay = [L_decay;zeros(Float64,1,length(gamma1)+1)]
	#L_decay = [0.0 sqrt(gamma1[1]) 0.0 0.0;
	#		   0.0 0.0      sqrt(gamma1[2]) 0.0;
	#		   0.0 0.0      0.0      sqrt(gamma1[3]);
	#		   0.0 0.0      0.0      0.0]
    return L_decay, L_dephase
end

"""
    RotationFrameRamseyControl(abs_control::Float64,phase_control::Float64,N::Int64)

# Argument:
- abs_control: strength of the control ``\\Omega``
- phase_control: phase of the control signal ``\\theta``
- N: number of states

# Output:
- HCK,HCS: the real and imaginary part of the control operator, determined by
  ``p(a+a^\\dagger)+iq(a-a^\\dagger),``
  where ``a`` is the lowering operator, ``p=|\\Omega|\\cos(\\theta)`` and ``q=|\\Omega|\\sin(\\theta)``
"""
function RotationFrameRamseyControl(abs_control::Float64,phase_control::Float64,N::Int64)
    p = abs_control*cos(phase_control);
    q = abs_control*sin(phase_control);
	amat = [Matrix(Diagonal(sqrt.(collect(1:N-1))));zeros(Float64,1,N-1)]
	#amat = [Diagonal(sqrt.(collect(1:N-1)));zeros(Float64,1,N-1)]
	amat = [zeros(Float64,N) amat]
    #amat = Bidiagonal(zeros(N),sqrt.(collect(1:N-1)),:U) # trouble maker, maybe we can make a Bidiagonal adjoint
    HCK = p*(amat+transpose(amat))
    HCS = q*(amat-transpose(amat))
    return HCK,HCS
end
"""
    RotationFrameRamseyControl(N::Int64)

# Argument:
- abs_control: strength of the control ``\\Omega``
- phase_control: phase of the control signal ``\\theta``
- N: number of states

# Output:
- HCK,HCS: the real and imaginary part of the control operator, determined by
  ``(a+a^\\dagger)``
"""
function RotationFrameRamseyControl(N::Int64)
	amat = [Matrix(Diagonal(sqrt.(collect(1:N-1))));zeros(Float64,1,N-1)]
	amat = [zeros(Float64,N) amat]
    HCK = amat+transpose(amat)
	HCS = zeros(Float64,N,N)
    return HCK,HCS
end

"""
	RamseyForwardSolve(rho_u0::Array{Float64},rho_v0::Array{Float64},
						  omega::Array{Float64},omega_drive::Float64,
						  gamma1::Array{Float64},gamma2::Array{Float64},
						  InitialState::Int64,
						  TC::Float64,tdark_times::Array{Float64},
					   	  N_states::Int64=0;
					      initial_type="states",
					   	  method="exponential",
					      DiffEqKwargs...)

# Argument:
- rho_u0,rho_v0: initial states ``\\rho_{u_0}-i\\rho_{v_0}``
- omega: transition frequencies
- omega_drive: driving frequency
- gamma1: the decay parameter  of the Lindblad operators, equal to ``1/T_1`` with ``T_1`` being the relaxation time
- gamma2: the dephasing parameter  of the Lindblad operators, equal to ``1/T_2`` with ``T_2`` being the dephasing time
- InitialState: initial state of the density matrix
- TC: control time of the control signal
- tdark_times: the dark time in the Ramsey experiment
- N_states: number of states. If it is not specified, the default value will be 0, and the code will automatically detect how many states are considered.
- initial_type: type of the initial condition
  1. initial values are state vectors (default choice)
  2. initial values are density matrices saving in flattened vectors
- method: method to solve the Lindblad system
 1. "exponential": exponential time integrator (default choice)
 2. "DiffEqDefault": a default choice made by DifferentialEquations.jl
 3. Other solvers available in DifferentialEquations.jl, for example, method = Trapezoid()
- DiffEqKwargs: keyword arguments feed to the ``solve function" of DifferentialEquations.jl

# Output:
- ``\\rho_u,\\rho_v``: real and imaginary part of the density matrix ``\\rho=\\rho_u-i\\rho_v`` at sampled dark times, and they are saved as flattened vectors
"""
function RamseyForwardSolve(rho_u0::Array{Float64},rho_v0::Array{Float64},
				   		  omega::Array{Float64},omega_drive::Float64,
				   		  gamma1::Array{Float64},gamma2::Array{Float64},
						  InitialState::Int64,
				   		  TC::Float64,t_dark_times::Array{Float64},
				   		  N_states::Int64=0;
						  initial_type="states",
						  method="exponential",
						  DiffEqKwargs...)
    if(N_states==0)
		N_states = length(omega)+1
	end
	#################################
	# Step 1: assemble operators
	#################################
	# Hamiltonian without control
	HK_free = RotationFrameDiagonal(omega,omega_drive)
	# Hamiltonain with control
	if(method=="exponential")
		Ω = 0.5*pi/(TC*sqrt(InitialState+1.0))
		θ = 0.0
		HK_control,HS_control = RotationFrameRamseyControl(Ω,θ,N_states)
	else
		HK_control,HS_control = RotationFrameRamseyControl(N_states)
	end
	# Lindblad operator
	L1,L2 = RotationFrameLindblad(gamma1,gamma2)
	# Assemble the operators for the vectorized system
	# Half pi pulse
	LK_half_pi,LS_half_pi,LD = make_lindblad_operator(Matrix(HK_control+HK_free),Matrix(HS_control),(L1,L2))
	# Free evolution
	LK_free = make_hamiltonian_operator(Matrix(HK_free))
	#################################
	# Step 2: Forward solves
	#################################
	# initial conditions
	if(initial_type == "states")
		rho_u_initial = (rho_u0*transpose(rho_u0)+rho_v0*transpose(rho_v0))[:]
		rho_v_initial = (rho_v0*transpose(rho_u0)-rho_u0*transpose(rho_v0))[:]
		rho_vec0 = [rho_u_initial;rho_v_initial]
	elseif(initial_type == "density")
		rho_vec0 = [rho_u0;rho_v0];
	else
		println("Error! initial_type must be \"density\" or \"states\"")
		return
	end
	N_dark_times = length(t_dark_times)
	if(method=="exponential")
	#if(false)
		# Ramsey experiment
		Half_pi_operator = exp(0.5*TC*[LD+LS_half_pi -LK_half_pi; LK_half_pi LD+LS_half_pi])
		FreeEvolution = [ LD -LK_free;
					  	  LK_free  LD]
		# half-pi pusle
		rho_vec1 = Half_pi_operator*rho_vec0
		# free propagation for a dark time
		Free_operator = exp(t_dark_times[1]*FreeEvolution)
		rho_vec2 = Free_operator*rho_vec1
		# half-pi pulse
		rho_ramsey = Half_pi_operator*rho_vec2
		for i = 2:N_dark_times
			# free propagation for a dark time
			Free_operator = exp(t_dark_times[i]*FreeEvolution)
			rho_vec2 = Free_operator*rho_vec1
			# half-pi pulse
			rho_ramsey = [rho_ramsey Half_pi_operator*rho_vec2]
		end
	else
		function GaussianPulse(_time,_amplitude,_width)
			σ = _width/sqrt(2.0*pi)
			return _amplitude*exp(-(_time-0.5*TC)^2/(2.0*σ*σ))
		end
		width = TC/2.5
		GaussianAmplitude = 0.5*pi/(width*sqrt(InitialState+1.0))
		function ControlledODE!(_drho,_rho,_p,_t)
			# Half pi pulse
			_LK_half_pi,_LS_half_pi,_LD = make_lindblad_operator(HK_control*GaussianPulse(_t,GaussianAmplitude,width)+HK_free,
															     HS_control*GaussianPulse(_t,GaussianAmplitude,width),
															     (L1,L2))
			SystemOperator  = [_LD+_LS_half_pi -_LK_half_pi;
					  		   _LK_half_pi      _LD+_LS_half_pi]
			_drho .= SystemOperator*_rho
		end
		# DiffEq interface for the controlled problem
		controlled_problem = ODEProblem(ControlledODE!,rho_vec0,(0.0,0.5*TC),[1.0])
		# Free propagation part
		FreeEvolution = [ LD -LK_free;
						  LK_free  LD]
		# half-pi pusle
		if(method=="DiffEqDefault")
			sol = solve(controlled_problem;DiffEqKwargs...)
		else
			sol = solve(controlled_problem,method;DiffEqKwargs...)
		end
		rho_vec1 = sol[:,end]
		# free propagation for a dark time
		Free_operator = exp(t_dark_times[1]*FreeEvolution)
		rho_vec2 = Free_operator*rho_vec1
		# second half-pi pulse
		controlled_problem = remake(controlled_problem,u0=rho_vec2)
		if(method=="DiffEqDefault")
			sol = solve(controlled_problem;DiffEqKwargs...)
		else
			sol = solve(controlled_problem,method;DiffEqKwargs...)
		end
		rho_ramsey = sol[:,end]
		# free propagation for a dark time + second half-pi pulse
		n_dark_times = length(t_dark_times)
		for i = 2:n_dark_times
			# free propagation
			Free_operator = exp(t_dark_times[i]*FreeEvolution)
			rho_vec2 = Free_operator*rho_vec1
			# second half pi pulse
			controlled_problem = remake(controlled_problem,u0=rho_vec2)
			if(method=="DiffEqDefault")
				sol = solve(controlled_problem;DiffEqKwargs...)
			else
				sol = solve(controlled_problem,method;DiffEqKwargs...)
			end
			rho_ramsey = [ rho_ramsey sol[:,end] ]
		end
	end
	return rho_ramsey[1:N_states^2,:],rho_ramsey[N_states^2+1:end,:]
end

# Defined for the compatability of adjoint solves in Turing.jl to hold their data structure
function RamseyForwardSolve(rho_u0,rho_v0,
				   		  omega,omega_drive,
				   		  gamma1,gamma2,
						  InitialState,
				   		  TC,t_dark_times,
				   		  N_states::Int64=0;
						  initial_type="states",
						  method="exponential",
						  DiffEqKwargs...)
    if(N_states==0)
		N_states = length(omega)+1
	end
	#################################
	# Step 1: assemble operators
	#################################
	# Hamiltonian without control
	HK_free = RotationFrameDiagonal(omega,omega_drive)
	# Hamiltonain with control
	if(method=="exponential")
		Ω = 0.5*pi/(TC*sqrt(InitialState+1.0))
		θ = 0.0
		HK_control,HS_control = RotationFrameRamseyControl(Ω,θ,N_states)
	else
		HK_control,HS_control = RotationFrameRamseyControl(N_states)
	end
	# Lindblad operator
	L1,L2 = RotationFrameLindblad(gamma1,gamma2)
	# Assemble the operators for the vectorized system
	# Half pi pulse
	LK_half_pi,LS_half_pi,LD = make_lindblad_operator(Matrix(HK_control+HK_free),Matrix(HS_control),(L1,L2))
	# Free evolution
	LK_free = make_hamiltonian_operator(Matrix(HK_free))
	#################################
	# Step 2: Forward solves
	#################################
	# initial conditions
	if(initial_type == "states")
		rho_u_initial = (rho_u0*transpose(rho_u0)+rho_v0*transpose(rho_v0))[:]
		rho_v_initial = (rho_v0*transpose(rho_u0)-rho_u0*transpose(rho_v0))[:]
		rho_vec0 = [rho_u_initial;rho_v_initial]
	elseif(initial_type == "density")
		rho_vec0 = [rho_u0;rho_v0];
	else
		println("Error! initial_type must be \"density\" or \"states\"")
		return
	end
	N_dark_times = length(t_dark_times)
	if(method=="exponential")
	#if(false)
		# Ramsey experiment
		Half_pi_operator = exp(0.5*TC*[LD+LS_half_pi -LK_half_pi; LK_half_pi LD+LS_half_pi])
		FreeEvolution = [ LD -LK_free;
					  	  LK_free  LD]
		# half-pi pusle
		rho_vec1 = Half_pi_operator*rho_vec0
		# free propagation for a dark time
		Free_operator = exp(t_dark_times[1]*FreeEvolution)
		rho_vec2 = Free_operator*rho_vec1
		# half-pi pulse
		rho_ramsey = Half_pi_operator*rho_vec2
		for i = 2:N_dark_times
			# free propagation for a dark time
			Free_operator = exp(t_dark_times[i]*FreeEvolution)
			rho_vec2 = Free_operator*rho_vec1
			# half-pi pulse
			rho_ramsey = [rho_ramsey Half_pi_operator*rho_vec2]
		end
	else
		function GaussianPulse(_time,_amplitude,_width)
			σ = _width/sqrt(2.0*pi)
			return _amplitude*exp(-(_time-0.5*TC)^2/(2.0*σ*σ))
		end
		width = TC/2.5
		GaussianAmplitude = 0.5*pi/(width*sqrt(InitialState+1.0))
		function ControlledODE!(_drho,_rho,_p,_t)
			# Half pi pulse
			_LK_half_pi,_LS_half_pi,_LD = make_lindblad_operator(HK_control*GaussianPulse(_t,GaussianAmplitude,width)+HK_free,
															     HS_control*GaussianPulse(_t,GaussianAmplitude,width),
															     (L1,L2))
			SystemOperator  = [_LD+_LS_half_pi -_LK_half_pi;
					  		   _LK_half_pi      _LD+_LS_half_pi]
			_drho .= SystemOperator*_rho
		end
		# DiffEq interface for the controlled problem
		controlled_problem = ODEProblem(ControlledODE!,rho_vec0,(0.0,0.5*TC),[omega;gamma1;gamma2])
		# Free propagation part
		FreeEvolution = [ LD -LK_free;
						  LK_free  LD]
		# half-pi pusle
		if(method=="DiffEqDefault")
			sol = solve(controlled_problem;DiffEqKwargs...)
		else
			sol = solve(controlled_problem,method;DiffEqKwargs...)
		end
		rho_vec1 = sol[:,end]
		# free propagation for a dark time
		Free_operator = exp(t_dark_times[1]*FreeEvolution)
		rho_vec2 = Free_operator*rho_vec1
		# second half-pi pulse
		controlled_problem = remake(controlled_problem,u0=rho_vec2)
		if(method=="DiffEqDefault")
			sol = solve(controlled_problem;DiffEqKwargs...)
		else
			sol = solve(controlled_problem,method;DiffEqKwargs...)
		end
		rho_ramsey = sol[:,end]
		# free propagation for a dark time + second half-pi pulse
		n_dark_times = length(t_dark_times)
		for i = 2:n_dark_times
			# free propagation
			Free_operator = exp(t_dark_times[i]*FreeEvolution)
			rho_vec2 = Free_operator*rho_vec1
			# second half pi pulse
			controlled_problem = remake(controlled_problem,u0=rho_vec2)
			if(method=="DiffEqDefault")
				sol = solve(controlled_problem;DiffEqKwargs...)
			else
				sol = solve(controlled_problem,method;DiffEqKwargs...)
			end
			rho_ramsey = [ rho_ramsey sol[:,end] ]
		end
	end
	return rho_ramsey[1:N_states^2,:],rho_ramsey[N_states^2+1:end,:]
end
"""
    RamseyParityForwardSolve(rho_u0::Array{Float64},rho_v0::Array{Float64},
					   omega::Array{Float64},omega_drive:Float64,
					   charge_noise::Array{Float64},
					   gamma1::Array{Float64},gamma2::Array{Float64},
					   InitalState::Int64,
					   TC::Float64,tdark_times::Array{Float64},
					   N_states::Int64=0;
					   initial_type="states",
					   method="exponential",
					   DiffEqKwargs...)

# Argument:
- rho_u0,rho_v0: initial states ``\\rho_{u_0}-i\\rho_{v_0}``
- omega: transition frequencies
- omega_drive: driving frequency
- gamma1: the decay parameter  of the Lindblad operators, equal to ``1/T_1`` with ``T_1`` being the relaxation time
- gamma2: the dephasing parameter  of the Lindblad operators, equal to ``1/T_2`` with ``T_2`` being the dephasing time
- Initial State: initial state of the density matrix
- TC: control time of the control signal
- dark_times: the dark time in the Ramsey experiment
- N_states: number of states. If it is not specified, the default value will be 0, and the code will automatically detect how many states are considered.
- initial_type: type of the initial condition
  1. initial values are state vectors (default choice)
  2. initial values are density matrices saving in flattened vectors
- method: method to solve the Landlady system
 1. "exponential": exponential time integrator (default choice)
 2. "DiffEqDefault": a default choice made by DifferentialEquations.jl
 3. Other solvers available in DifferentialEquations.jl, for example, method = Trapezoid()
- DiffEqKwargs: keyword arguments feed to the ``solve function" of DifferentialEquations.jl

# Parity event
- Transition frequency will take equal probability to be omega``\\pm``0.5``\\times``charge_noise

# Output:
- ``\\rho_u,\\rho_v``: real and imaginary part of the density matrix ``\\rho=\\rho_u-i\\rho_v`` at sampled dark times, and they are saved as flattened vectors
"""
function RamseyParityForwardSolve(rho_u0::Array{Float64},rho_v0::Array{Float64},
				   		  omega::Array{Float64},omega_drive::Float64,
						  charge_noise::Array{Float64},
				   		  gamma1::Array{Float64},gamma2::Array{Float64},
						  InitialState::Int64,
				   		  TC::Float64,t_dark_times::Array{Float64},
				   		  N_states::Int64=0;
						  initial_type="states",
						  method="exponential",
						  DiffEqKwargs...)

	rho_u1, rho_v1 = RamseyForwardSolve(rho_u0,rho_v0,
					   		  omega-0.5.*charge_noise,omega_drive,
					   		  gamma1,gamma2,
							  InitialState,
					   		  TC,t_dark_times,
					   		  N_states;
							  initial_type,
							  method,
							  DiffEqKwargs...)
	rho_u2, rho_v2 = RamseyForwardSolve(rho_u0,rho_v0,
							  omega+0.5.*charge_noise,omega_drive,
							  gamma1,gamma2,
							  InitialState,
							  TC,t_dark_times,
							  N_states;
							  initial_type,
							  method,
							  DiffEqKwargs...)
	return 0.5.*(rho_u1+rho_u2),0.5.*(rho_v1+rho_v2)
end

"""
    EchoForwardSolve(rho_u0::Array{Float64},rho_v0::Array{Float64},
					   omega::Array{Float64},omega_drive::Float64,
					   gamma1::Array{Float64},gamma2::Array{Float64},
					   InitialState::Int64,
					   TC::Float64,tdark_times::Array{Float64},
					   N_states::Int64=0;
					   initial_type="states",
					   method="exponential",
					   DiffEqKwargs...)

# Argument:
- rho_u0,rho_v0: initial states ``\\rho_{u_0}-i\\rho_{v_0}``
- omega: transition frequencies
- omega_drive: driving frequency
- gamma1: the decay parameter  of the Lindblad operators, equal to ``1/T_1`` with ``T_1`` being the relaxation time
- gamma2: the dephasing parameter  of the Lindblad operators, equal to ``1/T_2`` with ``T_2`` being the dephasing time
- InitialState: initial state of the density matrix
- TC: control time of the control signal
- tdark_times: the dark time in the Ramsey experiment
- N_states: number of states. If it is not specified, thedefault value will be 0, and the code will automatically detect how many states are considered.
- initial_type: type of the initial condition
  1. initial values are state vectors (default choice)
  2. initial values are density matrices saving in flattened vectors
- method: method to solve the Lindblad system
 1. "exponential": exponential time integrator (default choice)
 2. "DiffEqDefault": a default choice made by DifferentialEquations.jl
 3. Other solvers available in DifferentialEquations.jl, for example, method = Trapezoid()
- DiffEqKwargs: keyword arguments feed to the ``solve function" of DifferentialEquations.jl

# Output:
- ``\\rho_u,\\rho_v``: real and imaginary part of the density matrix ``\\rho=\\rho_u-i\\rho_v`` at sampled dark times, and they are saved as flattened vectors
"""
function EchoForwardSolve(rho_u0::Array{Float64},rho_v0::Array{Float64},
				   omega::Array{Float64},omega_drive::Float64,
				   gamma1::Array{Float64},gamma2::Array{Float64},
				   InitialState::Int64,
				   TC::Float64,t_dark_times::Array{Float64},
				   N_states::Int64=0;
				   initial_type="states",
				   method="exponential",
				   DiffEqKwargs...)

	if(N_states==0)
		N_states = length(omega)+1
	end
	#################################
	# Step 1: assemble operators
	#################################
	# Hamiltonian without control
	HK_free = RotationFrameDiagonal(omega,omega_drive)
	# Hamiltonain with control
	if(method=="exponential")
		Ω = 0.5*pi/(TC*sqrt(InitialState+1.0))
		θ = 0.0
		HK_control,HS_control = RotationFrameRamseyControl(Ω,θ,N_states)
	else
		HK_control,HS_control = RotationFrameRamseyControl(N_states)
	end
	# Lindblad operator
	L1,L2 = RotationFrameLindblad(gamma1,gamma2)
	# Assemble the operators for the vectorized system
	# Half pi pulse
	LK_half_pi,LS_half_pi,LD = make_lindblad_operator(Matrix(HK_control+HK_free),Matrix(HS_control),(L1,L2))
	# Free evolution
	LK_free = make_hamiltonian_operator(Matrix(HK_free))
	#################################
	# Step 2: Forward solves
	#################################
	# initial conditions
	if(initial_type == "states")
		rho_u_initial = (rho_u0*transpose(rho_u0)+rho_v0*transpose(rho_v0))[:]
		rho_v_initial = (rho_v0*transpose(rho_u0)-rho_u0*transpose(rho_v0))[:]
		rho_vec0 = [rho_u_initial;rho_v_initial]
	elseif(initial_type == "density")
		rho_vec0 = [rho_u0;rho_v0];
	else
		println("Error! initial_type must be \"density\" or \"states\"")
		return
	end
	N_dark_times = length(t_dark_times)
	if(method=="exponential")
		# Assemble operators
		Half_pi_operator = exp(0.5*TC*[LD+LS_half_pi -LK_half_pi; LK_half_pi LD+LS_half_pi])
		Pi_operator = Half_pi_operator*Half_pi_operator
		FreeEvolution = [ LD -LK_free;
					      LK_free  LD]
					  # half-pi pusle
		rho_vec1 = Half_pi_operator*rho_vec0
		# free propagation for a dark time
		Free_operator = exp(t_dark_times[1]*FreeEvolution)
		rho_vec2 = Free_operator*rho_vec1
		# pi pulse
		rho_vec2 = Pi_operator*rho_vec2
		# free propagation for the same dark time
		rho_vec2 = Free_operator*rho_vec2
		# half-pi pulse
		rho_echo = Half_pi_operator*rho_vec2
		for i = 2:N_dark_times
			# free propagation for a dark time
			Free_operator = exp(t_dark_times[i]*FreeEvolution)
			rho_vec2 = Free_operator*rho_vec1
			# pi pulse
			rho_vec2 = Pi_operator*rho_vec2
			# free propagation for the same dark time
			rho_vec2 = Free_operator*rho_vec2
			# half-pi pulse
			rho_echo = [rho_echo Half_pi_operator*rho_vec2]
		end
	else
		function GaussianPulse(_time,_amplitude,_width)
			_tot_time = _width*2.5
			σ = _width/sqrt(2.0*pi)
			return _amplitude*exp(-(_time-0.5*_tot_time)^2/(2.0*σ*σ))
		end
		width = TC/2.5
		GaussianAmplitude = 0.5*pi/(width*sqrt(InitialState+1.0))
		function ControlledODE!(_drho,_rho,_p,_t)
			# Half pi pulse
			_LK_half_pi,_LS_half_pi,_LD = make_lindblad_operator(HK_control*GaussianPulse(_t,GaussianAmplitude,width)+HK_free,
															     HS_control*GaussianPulse(_t,GaussianAmplitude,width),
															     (L1,L2))
			SystemOperator  = [_LD+_LS_half_pi -_LK_half_pi;
					  		   _LK_half_pi      _LD+_LS_half_pi]
			_drho .= SystemOperator*_rho
		end
		# DiffEq interface for the controlled problem
		half_pi_pulse = ODEProblem(ControlledODE!,rho_vec0,(0.0,0.5*TC),[1.0])
		# Free propagation part
		FreeEvolution = [ LD -LK_free;
						  LK_free  LD]
		# half-pi pusle
		if(method=="DiffEqDefault")
			sol = solve(half_pi_pulse;DiffEqKwargs...)
		else
			sol = solve(half_pi_pulse,method;DiffEqKwargs...)
		end
		rho_vec1 = sol[:,end]
		# free propagation for a dark time
		Free_operator = exp(t_dark_times[1]*FreeEvolution)
		rho_vec2 = Free_operator*rho_vec1
		# pi pulse
		pi_pulse = ODEProblem(ControlledODE!,rho_vec2,(0.0,TC),[1.0])
		if(method=="DiffEqDefault")
			sol = solve(pi_pulse;DiffEqKwargs...)
		else
			sol = solve(pi_pulse,method;DiffEqKwargs...)
		end
		# free propagation for the same dark time
		rho_vec2 = Free_operator*sol[:,end]
		# half pi pulse
		half_pi_pulse = remake(half_pi_pulse,u0=rho_vec2)
		if(method=="DiffEqDefault")
			sol = solve(half_pi_pulse;DiffEqKwargs...)
		else
			sol = solve(half_pi_pulse,method;DiffEqKwargs...)
		end
		rho_echo = sol[:,end]
		# free propagation for a dark time + second half-pi pulse
		n_dark_times = length(t_dark_times)
		for i = 2:n_dark_times
			Free_operator = exp(t_dark_times[i]*FreeEvolution)
			# second half pi
			rho_vec2 = Free_operator*rho_vec1
			# pi pulse
			pi_pulse = remake(pi_pulse,u0=rho_vec2)
			if(method=="DiffEqDefault")
				sol = solve(pi_pulse;DiffEqKwargs...)
			else
				sol = solve(pi_pulse,method;DiffEqKwargs...)
			end
			# free propagation for the same dark time
			rho_vec2 = Free_operator*sol[:,end]
			# half pi pulse
			half_pi_pulse = remake(half_pi_pulse,u0=rho_vec2)
			if(method=="DiffEqDefault")
				sol = solve(half_pi_pulse;DiffEqKwargs...)
			else
				sol = solve(half_pi_pulse,method;DiffEqKwargs...)
			end
			rho_echo = [rho_echo sol[:,end]]
		end

	end

	return rho_echo[1:N_states^2,:],rho_echo[N_states^2+1:end,:]
end

"""
    EchoParityForwardSolve(rho_u0::Array{Float64},rho_v0::Array{Float64},
					   omega::Array{Float64},omega_drive::Float64,
					   charge_noise::Array{Float64},
					   gamma1::Array{Float64},gamma2::Array{Float64},
					   InitialState::Int64,
					   TC::Float64,tdark_times::Array{Float64},
					   N_states::Int64=0;
					   initial_type="states",
					   method="exponential",
					   DiffEqKwargs...)

# Argument:
- rho_u0,rho_v0: initial states ``\\rho_{u_0}-i\\rho_{v_0}``
- omega: transition frequencies
- omega_drive: driving frequency
- gamma1: the decay parameter  of the Lindblad operators, equal to ``1/T_1`` with ``T_1`` being the relaxation time
- gamma2: the dephasing parameter  of the Lindblad operators, equal to ``1/T_2`` with ``T_2`` being the dephasing time
- InitialState: initial state of the density matrix
- TC: control time of the control signal
- tdark_times: the dark time in the Ramsey experiment
- N_states: number of states. If it is not specified, the default value will be 0, and the code will automatically detect how many states are considered.
- initial_type: type of the initial condition
  1. initial values are state vectors (default choice)
  2. initial values are density matrices saving in flattened vectors
- method: method to solve the Lindblad system
 1. "exponential": exponential time integrator (default choice)
 2. "DiffEqDefault": a default choice made by DifferentialEquations.jl
 3. Other solvers available in DifferentialEquations.jl, for example, method = Trapezoid()
- DiffEqKwargs: keyword arguments feed to the ``solve function" of DifferentialEquations.jl

# Parity event
- Transition frequency will take equal probability to be omega``\\pm``0.5``\\times``charge_noise

# Output:
- ``\\rho_u,\\rho_v``: real and imaginary part of the density matrix ``\\rho=\\rho_u-i\\rho_v`` at sampled dark times, and they are saved as flattened vectors
"""
function EchoParityForwardSolve(rho_u0::Array{Float64},rho_v0::Array{Float64},
				   omega::Array{Float64},omega_drive::Float64,
				   charge_noise::Array{Float64},
				   gamma1::Array{Float64},gamma2::Array{Float64},
				   InitialState::Int64,
				   TC::Float64,t_dark_times::Array{Float64},
				   N_states::Int64=0;
				   initial_type="states",
				   method="exponential",
				   DiffEqKwargs...)

	rho_u1, rho_v1 = EchoForwardSolve(rho_u0,rho_v0,
							  omega-0.5.*charge_noise,omega_drive,
							  gamma1,gamma2,
							  InitialState,
							  TC,t_dark_times,
							  N_states;
							  initial_type,
							  method,
							  DiffEqKwargs...)
	rho_u2, rho_v2 = EchoForwardSolve(rho_u0,rho_v0,
							  omega+0.5.*charge_noise,omega_drive,
							  gamma1,gamma2,
							  InitialState,
							  TC,t_dark_times,
							  N_states;
							  initial_type,
							  method,
							  DiffEqKwargs...)
	return 0.5.*(rho_u1+rho_u2),0.5.*(rho_v1+rho_v2)
end

"""
    T1ForwardSolve(rho_u0::Array{Float64},rho_v0::Array{Float64},
					   omega::Array{Float64},omega_drive::Float64,
					   gamma1::Array{Float64},gamma2::Array{Float64},
					   InitialState::Int64,
					   TC::Float64,tdark_times::Array{Float64},
					   N_states::Int64=0;
					   initial_type="states",
					   method="exponential",
					   DiffEqKwargs...)

# Argument:
- rho_u0,rho_v0: initial states ``\\rho_{u_0}-i\\rho_{v_0}``
- omega: transition frequencies
- omega_drive: driving frequency
- gamma1: the decay parameter  of the Lindblad operators, equal to ``1/T_1`` with ``T_1`` being the relaxation time
- gamma2: the dephasing parameter  of the Lindblad operators, equal to ``1/T_2`` with ``T_2`` being the dephasing time
- InitialState: initial state of the density matrix
- TC: control time of the control signal
- tdark_times: the dark time in the Ramsey experiment
- N_states: number of states. If it is not specified, the default value will be 0, and the code will automatically detect how many states are considered.
- initial_type: type of the initial condition
  1. initial values are state vectors (default choice)
  2. initial values are density matrices saving in flattened vectors
- method: method to solve the Lindblad system
 1. "exponential": exponential time integrator (default choice)
 2. "DiffEqDefault": a default choice made by DifferentialEquations.jl
 3. Other solvers available in DifferentialEquations.jl, for example, method = Trapezoid()
- DiffEqKwargs: keyword arguments feed to the ``solve function" of DifferentialEquations.jl

# Output:
- ``\\rho_u,\\rho_v``: real and imaginary part of the density matrix ``\\rho=\\rho_u-i\\rho_v`` at sampled dark times, and they are saved as flattened vectors
"""
function T1ForwardSolve(rho_u0::Array{Float64},rho_v0::Array{Float64},
				   omega::Array{Float64},omega_drive::Float64,
				   gamma1::Array{Float64},gamma2::Array{Float64},
				   InitialState::Int64,
				   TC::Float64,t_dark_times::Array{Float64},
				   N_states::Int64=0;
				   initial_type="states",
				   method="exponential",
				   DiffEqKwargs...)

	if(N_states==0)
		N_states = length(omega)+1
	end
	#################################
	# Step 1: assemble operators
	#################################
	# Hamiltonian without control
	HK_free = RotationFrameDiagonal(omega,omega_drive)
	# Hamiltonain with control
	# Hamiltonain with control
	if(method=="exponential")
		Ω = 0.5*pi/(TC*sqrt(InitialState+1.0))
		θ = 0.0
		HK_control,HS_control = RotationFrameRamseyControl(Ω,θ,N_states)
	else
		HK_control,HS_control = RotationFrameRamseyControl(N_states)
	end
	# Lindblad operator
	L1,L2 = RotationFrameLindblad(gamma1,gamma2)
	# Assemble the operators for the vectorized system
	# Half pi pulse
	LK_half_pi,LS_half_pi,LD = make_lindblad_operator(Matrix(HK_control+HK_free),Matrix(HS_control),(L1,L2))
	# Free evolution
	LK_free = make_hamiltonian_operator(Matrix(HK_free))
	#################################
	# Step 2: Forward solves
	#################################
	# initial conditions
	if(initial_type == "states")
		rho_u_initial = (rho_u0*transpose(rho_u0)+rho_v0*transpose(rho_v0))[:]
		rho_v_initial = (rho_v0*transpose(rho_u0)-rho_u0*transpose(rho_v0))[:]
		rho_vec0 = [rho_u_initial;rho_v_initial]
	elseif(initial_type == "density")
		rho_vec0 = [rho_u0;rho_v0];
	else
		println("Error! initial_type must be \"density\" or \"states\"")
		return
	end
	N_dark_times = length(t_dark_times)
	if(method=="exponential")
		# Assemble operators
		Pi_operator = exp(TC*[LD+LS_half_pi -LK_half_pi; LK_half_pi LD+LS_half_pi])
		FreeEvolution = [ LD -LK_free;
					  	  LK_free  LD]
		# pi pusle
		rho_vec1 = Pi_operator*rho_vec0
		# free propagation for a dark time
		Free_operator = exp(t_dark_times[1]*FreeEvolution)
		rho_T1 = Free_operator*rho_vec1
		for i = 2:N_dark_times
			# free propagation for a dark time
			Free_operator = exp(t_dark_times[i]*FreeEvolution)
			rho_T1 = [rho_T1 Free_operator*rho_vec1]
		end
	else
		function GaussianPulse(_time,_amplitude,_width)
			σ = _width/sqrt(2.0*pi)
			return _amplitude*exp(-(_time-0.5*TC)^2/(2.0*σ*σ))
		end
		width = TC/2.5
		GaussianAmplitude = 0.5*pi/(width*sqrt(InitialState+1.0))
		function ControlledODE!(_drho,_rho,_p,_t)
			# Half pi pulse
			_LK_half_pi,_LS_half_pi,_LD = make_lindblad_operator(HK_control*GaussianPulse(_t,GaussianAmplitude,width)+HK_free,
																 HS_control*GaussianPulse(_t,GaussianAmplitude,width),
																 (L1,L2))
			SystemOperator  = [_LD+_LS_half_pi -_LK_half_pi;
							   _LK_half_pi      _LD+_LS_half_pi]
			_drho .= SystemOperator*_rho
		end
		# DiffEq interface for the controlled problem
		pi_pulse = ODEProblem(ControlledODE!,rho_vec0,(0.0,TC),[1.0])
		# Free propagation part
		FreeEvolution = [ LD -LK_free;
						  LK_free  LD]
		# pi pusle
		if(method=="DiffEqDefault")
			sol = solve(pi_pulse;DiffEqKwargs...)
		else
			sol = solve(pi_pulse,method;DiffEqKwargs...)
		end
		rho_vec1 = sol[:,end]
		Free_operator = exp(t_dark_times[1]*FreeEvolution)
		rho_T1 = Free_operator*rho_vec1
		for i = 2:N_dark_times
			# free propagation for a dark time
			Free_operator = exp(t_dark_times[i]*FreeEvolution)
			rho_T1 = [rho_T1 Free_operator*rho_vec1]
		end
	end
	return rho_T1[1:N_states^2,:],rho_T1[N_states^2+1:end,:]

end

"""
    T1ParityForwardSolve(rho_u0::Array{Float64},rho_v0::Array{Float64},
					   omega::Array{Float64},omega_drive::Float64,
					   charge_noise::Array{Float64},
					   gamma1::Array{Float64},gamma2::Array{Float64},
					   InitialState::Int64,
					   TC::Float64,tdark_times::Array{Float64},
					   N_states::Int64=0;
					   initial_type="states",
					   method="exponential",
					   DiffEqKwargs...)

# Argument:
- rho_u0,rho_v0: initial states ``\\rho_{u_0}-i\\rho_{v_0}``
- omega: transition frequencies
- omega_drive: driving frequency
- gamma1: the decay parameter  of the Lindblad operators, equal to ``1/T_1`` with ``T_1`` being the relaxation time
- gamma2: the dephasing parameter  of the Lindblad operators, equal to ``1/T_2`` with ``T_2`` being the dephasing time
- InitialState: initial state of the density matrix
- TC: control time of the control signal
- tdark_times: the dark time in the Ramsey experiment
- N_states: number of states. If it is not specified, the default value will be 0, and the code will automatically detect how many states are considered.
- initial_type: type of the initial condition
  1. initial values are state vectors (default choice)
  2. initial values are density matrices saving in flattened vectors
- method: method to solve the Lindblad system
 1. "exponential": exponential time integrator (default choice)
 2. "DiffEqDefault": a default choice made by DifferentialEquations.jl
 3. Other solvers available in DifferentialEquations.jl, for example, method = Trapezoid()
- DiffEqKwargs: keyword arguments feed to the ``solve function" of DifferentialEquations.jl

# Parity event
- Transition frequency will take equal probability to be omega``\\pm``0.5``\\times``charge_noise

# Output:
- ``\\rho_u,\\rho_v``: real and imaginary part of the density matrix ``\\rho=\\rho_u-i\\rho_v`` at sampled dark times, and they are saved as flattened vectors
"""
function T1ParityForwardSolve(rho_u0::Array{Float64},rho_v0::Array{Float64},
				   omega::Array{Float64},omega_drive::Float64,
				   charge_noise::Array{Float64},
				   gamma1::Array{Float64},gamma2::Array{Float64},
				   InitialState::Int64,
				   TC::Float64,t_dark_times::Array{Float64},
				   N_states::Int64=0;
				   initial_type="states",
				   method="exponential",
				   DiffEqKwargs...)

	rho_u1, rho_v1 = T1ForwardSolve(rho_u0,rho_v0,
							  omega-0.5.*charge_noise,omega_drive,
							  gamma1,gamma2,
							  InitialState,
							  TC,t_dark_times,
							  N_states;
							  initial_type,
							  method,
							  DiffEqKwargs...)
	rho_u2, rho_v2 = T1ForwardSolve(rho_u0,rho_v0,
							  omega+0.5.*charge_noise,omega_drive,
							  gamma1,gamma2,
							  InitialState,
							  TC,t_dark_times,
							  N_states;
							  initial_type,
							  method,
							  DiffEqKwargs...)
	return 0.5.*(rho_u1+rho_u2),0.5.*(rho_v1+rho_v2)
end


text:
module X

using CompEcon
include("../../ModelAbstraction.jl")
using .ModelAbstraction

immutable BFZStateSpace <: AbstractStateSpace{2}
    k::EndogenousState{Float64}
    x::ExogenousState{Float64}
    g::ExogenousState{Float64}
    grid::Matrix{Float64}
    grid_transpose::Matrix{Float64}
    basis::Basis{2}
    basis_struct::BasisStructure{Direct}
end

immutable BFZModel
    agent::EpsteinZinAgent
    producer::CESProducer
    exog::ConstantVolatility
    state_space::BFZStateSpace
end

function BFZcv(;ρ::Real=-1, α::Real=-9, β::Real=0.99,
              η::Real=1/3, ν::Float64=0.1, δ::Real=0.025,
              lgbar::Real=0.004, A::Real=0., B::Real=1., vbar::Real=0.015^2,
              nk::Int=20, nx::Int=8, nϵ1::Int=3)
    # (ρ, α, β, η, ν, δ, lgbar, A, B, vbar, nk, nx, nϵ1) = (-1., -9., .99, .33, .1, .025, .004, 0., 1., .015^2, 20, 8, 9)
    # Create the Consumer
    BFZagent = EpsteinZinAgent(ρ, α, β)

    # Create means of Production
    BFZprod = CESProducer(η, ν, δ)

    # Create Exogenous Process
    exog = ConstantVolatility(A, B, vbar, nϵ1)

    #
    # Create the State Space
    #
    # First Create Basis
    basis = Basis(SplineParams(collect(linspace(15.0, 30.0, nk)), 0, 3),
                  SplineParams(collect(linspace(-.05, .05, nx)), 0, 1))
    basis_struc = BasisStructure(basis, Direct())

    # Get the implied grids
    grid, (kgrid, xgrid) = nodes(basis)
    grid_transpose = grid'
    ggrid = exp(lgbar + xgrid)

    xpgrid = Array(Float64, nϵ1, nx)
    gpgrid = Array(Float64, nϵ1, nx)
    for i=1:nx
        xpgrid[:, i] = step(exog, xgrid[i], exog.ϵ)
        gpgrid[:, i] = exp(lgbar + xpgrid[:, i])
    end

    # Create the states
    kstate = EndogenousState(kgrid)
    xstate = ExogenousState(xgrid, xpgrid)
    gstate = ExogenousState(ggrid, gpgrid)

    # Finally create space
    statespace = BFZStateSpace(kstate, xstate, gstate, grid, grid_transpose,
                               basis, basis_struc)

    return BFZModel(BFZagent, BFZprod, exog, statespace)
end

function envelope_method(bfz::BFZModel; maxiter=500, tol=1e-4)

    # Useful parameters
    dist = 10.
    iter = 1
    (ρ, α, β) = bfz.agent.ρ, bfz.agent.α, bfz.agent.β
    (η, ν, δ) = bfz.producer.η, bfz.producer.ν, bfz.producer.δ
    ρinv = 1./ρ

    # Useful function
    _Y(k) = _production(bfz.producer, k)

    # Pull information out of type
    kgrid = bfz.state_space.k.grid
    xgrid = bfz.state_space.x.grid
    stategrid = bfz.state_space.grid
    nk, nx = length(kgrid), length(xgrid)

    # Initialize Value Function and Interpolant
    JT = (((1 - β) * (_Y(stategrid[:, 1])+(1 - δ)*stategrid[:, 1]).^(ρ)).^(ρinv)) ./ 500
    Jvals = repeat(JT, inner=[1, nx])
    Jupd = copy(Jvals)
    Jcoeffs = funfitxy(bfz.state_space.basis, bfz.state_space.basis_struct)
    χpolicy = Array(Float64, nk, nx)

    # Solve this bad boy
    while dist>tol && iter<maxiter

        iter += 1
        # Given the current vf, find the χ that solves the env condition
        for xind=1:nx
            # Pull out current x
            xt = xgrid[xind]

            # Create a 1 dimensional interpolant so I can take derivative
            Jvalsi = Jvals[:, xind]
            Jkspl = Spline1D(kgrid, Jvalsi, k=1, bc="nearest", s=0.0)

            for kind=1:nk
                # Pull out current k
                kt = kgrid[kind]

                # Get the pieces that we care about
                yt = _Y(kt)
                dJ_k = derivative(Jkspl, kt)
                Jt = evaluate(Jspl, kt, xt)
                ct = (dJ_k/((1-β)*Jt^(1-ρ)*(yt^(1-ν)*η*kt^(ν-1)+1-δ)))^(1./(ρ-1))
                χt = yt + (1 - δ)*kt - ct
                if χt < 0
                    warn("Negative chi")
                end
                χpolicy[kind, xind] = χt
                expterm = 0.

                for shockind=1:nshocks
                    xtp1 = xpgrid[shockind, xind]
                    gtp1 = gpgrid[shockind, xind]
                    jtp1 = evaluate(Jspl, χt/gtp1, xtp1)
                    expterm += (jtp1 * gtp1)^α * Π[shockind]
                end
                μ = expterm^αinv

                Jupd[kind, xind] = ((1 - β)*ct^ρ + β*μ^ρ)^(ρinv)

            end
        end

        dist = maxabs(Jvals - Jupd)
        copy!(Jvals, Jupd)
        Jspl = Spline2D(kgrid, xgrid, Jvals; kx=1, ky=1, s=0.0)
        @show iter, dist

    end

    return Jvals, χpolicy
end

end  # End Module

ext: jl

max_stars_repo_name: NYUEcon/GrowthModels

max_issues_repo_name: NYUEcon/GrowthModels

max_forks_repo_name: NYUEcon/GrowthModels