using LinearAlgebra
using CUDA

"""
	sim_activity_gpu(J::Matrix{Float32}, g::Float32, h0::Vector{Float32};
		dt=1f-1, t_bound=1f3, s=0f0, savestep=1, minimum_norm=0)

Integrates a tanh rate network with coupling matrix J and gain param g from
initial condition h0 and returns the matrix of single unit activities.
"""
function sim_activity_gpu(J::Matrix{Float32}, g::Float32, h0::Vector{Float32};
		dt=1f-1, t_bound=1f3, s=0f0, savestep=1, minimum_norm=0)

	N = size(J)[1]
	# Check dimensional consistency
	@assert(N==size(J)[2]==length(h0))

	maxstep = round(Int64, t_bound / dt)

	# Allocate memory for output
	out = CuArray{Float32,2}(undef, N, maxstep ÷ savestep)

	h = CuArray{Float32,1}(h0) # Set initial condition
	r = CuArray{Float32,1}(undef, N)

	J = CuArray{Float32,2}(g .* J + s * I)

	step = 0
	while step < maxstep
		step += 1
		map!(tanh, r, h)
		mul!(h, J, r, dt, 1 - dt)
		if step % savestep == 0 # Save every specified time interval
			savept = step ÷ savestep
			out[:,savept] .= h
		end
		if norm(h) < minimum_norm
			D = Matrix(-I + J .* ((sech.(h)).^2)')
			if all(real(eigvals(D)) .< 0)
				break
			end
		end
	end

	return Matrix{Float32}(out)
end


"""
	sim_lyapunov_gpu(J::Matrix{Float32}, g::Float32, h0::Vector{Float32},
		dh::Vector{Float32}, Q0::Matrix{Float32}; dt=1f-1, t_bound=1f3,
		t_ONS=1f0, wait=0f0, return_max_LE=false, s=0f0, printstuff=false)

Integrates a tanh network with a given weight matrix and gain parameter from a
specified initial condition, and returns a Lyapunov spectrum calculated from
Jacobians along the solution trajectory.

Can also return the max Lyapunov exponent calculated from a direct numerical
simulation of a perturbed trajectory.
"""
function sim_lyapunov_gpu(J::Matrix{Float32}, g::Float32, h0::Vector{Float32},
		dh::Vector{Float32}, Q0::Matrix{Float32}; dt=1f-1, t_bound=1f3,
		t_ONS=1f0, wait=0f0, return_max_LE=false, s=0f0, printstuff=false)

	# Check dimensional consistency
	N = size(J, 1)
	@assert size(J) == (length(h0), length(dh)) == size(Q0) == (N,N)

	# Discretization of time
	maxstep = round(Int64, t_bound / dt)
	startstep = round(Int64, wait / dt)
	tONSstep = round(Int64, t_ONS / dt)
	@assert(maxstep % tONSstep == 0 && startstep % tONSstep == 0,
			"The reorthonormalization interval must divide both the total " *
			"simulation time and wait period"
	)

	# Allocate memory for output
	spectrum = CUDA.zeros(N)

	# Allocate vector for activity and firing rate of current state
	h = CuArray{Float32,1}(h0)
	r = CuArray{Float32,1}(undef, N)

	# Load network structure and ICs to GPU
	J = CuArray{Float32,2}(g .* J + s * I)
	Q = CuArray{Float32,2}(Q0)

	# Allocate mem for calculating Jacobian and evolving the orthonormal basis
	sech2(x::Float32)::Float32 = 1f0 / (CUDA.cosh(x))^2
	sech2h = CuArray{Float32,2}(undef, 1, N)
	D = CuArray{Float32,2}(undef, N, N)
	dQ = CuArray{Float32,2}(undef, N, N)

	# Things for performing a direct numerical sim of a perturbation
	if return_max_LE
		max_LE = 0.0f0
		p_factor = 1f-6  # Order of the perturbation handled by the sim
		dh = CuArray{Float32,1}(dh ./ norm(dh))
		h2 = h .+ p_factor .* dh
		r2 = CuArray{Float32,1}(undef, N)
	end

	step = 0
	while step < maxstep
		step += 1

		#Evolve the trajectory and the orthonormal basis in the tangent space
		map!(tanh, r, h)
		mul!(h, J, r, dt, 1 - dt)

		map!(sech2, sech2h, h)
		copyto!(D, J)
		D .*= sech2h
		D[diagind(D)] .-= 1.0

		mul!(dQ, D, Q, dt, 0f0)
        Q .+= dQ

		#Evolve the perturbed trajectory and calculate how much it has grown
		if return_max_LE
			map!(tanh, r2, h2)
			mul!(h2, J, r2, dt, 1 - dt)
			copyto!(dh, h2)
			dh .-= h
			dh ./= p_factor
			scale = norm(dh)

			# If the perturbation scale has grown or shrunk too much, then
			# reset it to 1 and track the cumulative sum of log scales
			if scale > 100 || scale < 0.01 || step == startstep
				copyto!(h2, h)
				dh ./= (scale / p_factor)
				h2 .+= dh
				if step > startstep
					max_LE += log(scale)
				end
			end
		end

		# At each re-orthonormalization step, the log diagonal entries of R are
		# saved in a cumulative sum
		if step % tONSstep == 0
			printstuff && println("Time step: ", step)
			copyto!(dQ, Q)
			F = qr!(dQ)
			fill!(Q, 0)
			diags = view(F.factors, diagind(F.factors))
			map!(sign, view(Q, diagind(Q)), diags)
			if step > startstep  # forcing positive R-diagonal elements
				spectrum .+= log.(abs.(diags))
			end
            lmul!(F.Q, Q)
		end
	end

	# Update the maximum Lyapunov exponent after the simulation has ended so
	# that the full simulation time is used for the calculation
	actual_time = dt * (maxstep - startstep)
	if return_max_LE
		max_LE += log(norm((h2 .- h) ./ p_factor))
		spectrum ./= actual_time
		max_LE /= actual_time
		return (Vector{Float32}(spectrum), max_LE)
	else
		spectrum ./= actual_time 
		return Vector{Float32}(spectrum)
	end
end
