using LinearAlgebra

"""
	sim_activity(J::Matrix{T}, g::T, h0::Vector{T}; dt=0.1, t_bound=1000.0,
		s=0.0, b=nothing, savestep=1, minimum_norm=0) where T <: AbstractFloat

Integrates a tanh rate network with coupling matrix `J` and strength scaling `g`
from initial condition `h0`, returning a `N x T` matrix of network unit activity.

When `minimum_norm > 0` the integration will be cut off when the network state
falls too close to the origin to prevent floating point underflow. The keyword
arguments `s` and `b` add self-coupling to `J` and add a bias term, respectively.
"""
function sim_activity(J::Matrix{T}, g::T, h0::Vector{T}; dt=0.1, t_bound=1000.0,
		s=0.0, b=nothing, savestep=1, minimum_norm=0) where T <: AbstractFloat

	N = size(J)[1]
	# Check network size and type consistency
	@assert N == size(J)[2] == length(h0)
	@assert T == typeof(dt) == typeof(t_bound) == typeof(s)

	if b isa Array{T} && (size(b) == (N,) || size(b) == (N, 1))
		use_bias = true
	else
		use_bias = false
	end

	maxstep = round(Int, t_bound / dt) # Discretized simulation length

	# Allocate memory for output
	out = Matrix{T}(undef, N, maxstep ÷ savestep)

	h = copy(h0) # Set initial condition
	r = Vector{T}(undef, N)

	J = g .* J + s * I

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

	lastindex = step ÷ savestep
	return out[:,1:lastindex]
end



"""
	sim_lyapunov(J::Matrix{T}, g::T, h0::Vector{T}, dh::Vector{T},
		Q0::Matrix{T}; s=0.0, b=nothing, dt=0.1, t_bound=1000.0, t_ONS=1.0,
		wait=0.0, return_max_LE=false, printstuff=false, scale_limit=100,
		p_factor=1e-1) where T <: AbstractFloat

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

Optionally returns the maximum Lyapunov exponent calculated using a direct
numerical simulation of a perturbed trajectory.
"""
function sim_lyapunov(J::Matrix{T}, g::T, h0::Vector{T}, dh::Vector{T},
		Q0::Matrix{T}; s=0.0, b=nothing, dt=0.1, t_bound=1000.0, t_ONS=1.0,
		wait=0.0, return_max_LE=false, printstuff=false, scale_limit=100,
		p_factor=1e-1) where T <: AbstractFloat

	N = size(J)[1]

	@assert(size(J) == (length(h0), length(dh)) == size(Q0) == (N,N),
			"Dimensions of the coupling matrix, IC, initial ONS, and " *
			"perturbation must be consistent"
	)
	@assert(T == typeof(s) == typeof(dt) == typeof(t_bound) == typeof(t_ONS)
			== typeof(wait), "Make sure all scalar parameters are the same type"
	)

	maxstep = round(Int64, t_bound / dt)
	startstep = round(Int64, wait / dt)
	tONSstep = round(t_ONS / dt)
	@assert(maxstep % tONSstep == startstep % tONSstep == 0,
		"t_ONS must exactly divide t_bound and wait"
	)

	# Whether or not to include a bias term in the network dynamics
	if b isa Array{T} && size(b) == (N,)
		use_bias = true
	elseif b isa Array{T} && size(b) == (N, 1)
		use_bias = true
		b = reshape(b, (N,))
	else
		use_bias = false
	end

	# Allocate memory for output
	spectrum = zeros(T, N)

	# Allocate vector for activity and firing rate of current state
	h = copy(h0)
	r = Vector{T}(undef, N)

	# Allocate memory for Jacobian calculation and orthonormal basis
	sech2(x::T)::T = (sech(x))^2
	sech2h = Matrix{T}(undef, 1, N)
	D = Matrix{T}(undef, N, N)
	Q = copy(Q0)
	dQ = Matrix{T}(undef, N, N)

	if return_max_LE
		max_LE = 0
		p_factor = convert(T, p_factor)
		dh = copy(dh) ./ norm(dh)
		h2 = h .+ p_factor .* dh
		r2 = Vector{T}(undef, N)
	end

	# The effective coupling matrix J
	J = g .* J + s * I

	step = 0
	while step < maxstep
		step += 1

		#Evolve the trajectory and the orthonormal basis in the tangent space
		map!(tanh, r, h)
		BLAS.gemv!('N', dt, J, r, 1 - dt, h)
		if use_bias
			h .+= dt .* b
		end
		map!(sech2, sech2h, h)
		copyto!(D, J)
		D .*= sech2h
		D[diagind(D)] .-= 1
		BLAS.gemm!('N', 'N', dt, D, Q, T(0), dQ)
		Q .+= dQ

		#Evolve the perturbed trajectory and calculate how much it has grown
		if return_max_LE
			map!(tanh, r2, h2)
			BLAS.gemv!('N', dt, J, r2, 1 - dt, h2)
			if use_bias
				h2 .+= dt .* b
			end
			dh .= (h2 .- h) ./ 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 > scale_limit || scale < 1 / scale_limit || step == startstep
				h2 .= h .+ dh ./ scale .* p_factor
				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)
			diags = view(F.factors, diagind(F.factors))
			if step > startstep  # forcing positive R-diagonal elements
				spectrum .+= log.(abs.(diags))
			end
			fill!(Q, 0)
			map!(sign, view(Q, diagind(Q)), diags)
			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
	true_time = dt * (maxstep - startstep)
	if return_max_LE
		max_LE += log(norm((h2 .- h) ./ p_factor))
		return spectrum ./ true_time, max_LE / true_time
	else
		return spectrum ./ true_time
	end
end



"""
	sim_direct_max_LE(J::Matrix{T}, g::T, h0::Vector{T}, dh::Vector{T};
		s=0.0, b=nothing, dt=0.1, t_bound=1000.0, wait=0.0, printstuff=false,
		scale_limit=100, p_factor=1e-1) where T <: AbstractFloat

Calculates the maximum Lyapunov exponent of a tanh rate network given by
coupling matrix J, dependent on a given initial condition and initial
perturbation, using a direct method that tracks the growth of the perturbed
trajectory relative to the unperturbed one.
"""
function sim_direct_max_LE(J::Matrix{T}, g::T, h0::Vector{T}, dh::Vector{T};
		s=0.0, b=nothing, dt=0.1, t_bound=1000.0, wait=0.0, printstuff=false,
		scale_limit=100, p_factor=1e-1) where T <: AbstractFloat

	N = size(J)[1]
	@assert(size(J) == (length(h0), length(dh)) == (N,N), "Dimensions of the " *
			"coupling matrix, IC, and perturbation must be consistent"
	)

	# Whether or not to include a bias term in the network dynamics
	if b isa Array{T} && size(b) == (N,)
		use_bias = true
	elseif b isa Array{T} && size(b) == (N, 1)
		use_bias = true
		b = reshape(b, (N,))
	else
		use_bias = false
	end

	maxstep = round(Int64, t_bound / dt)
	startstep = round(Int64, wait / dt)

	# Effective coupling matrix
	J = g .* J + s * I

	# Order of the perturbation handled by the sim
	p_factor = convert(T, p_factor)
	# Allocate memory for output
	max_LE = 0

	# Set the initial condition and perturbed IC
	h = copy(h0)
	dh = copy(dh) ./ norm(dh)
	h2 = h .+ p_factor .* dh

	r = Vector{T}(undef, N)
	r2 = Vector{T}(undef, N)

	step = 0
	while step < maxstep
		step += 1

		# Evolve the trajectory and its perturbation
		map!(tanh, r, h)
		map!(tanh, r2, h2)
		BLAS.gemv!('N', dt, J, r, 1 - dt, h)
		BLAS.gemv!('N', dt, J, r2, 1 - dt, h2)
		if use_bias
			h .+= dt .* b
			h2 .+= dt .* b
		end
		dh .= (h2 .- h) ./ 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 > scale_limit || scale < 1 / scale_limit || step == startstep
			printstuff && println("Time step: ", step, ", Scale: ", scale)
			h2 .= h .+ dh ./ scale .* p_factor
			if step > startstep
				max_LE += log(scale)
			end
		end
	end

	# Update the maximum Lyapunov exponent after the simulation has ended so
	# that the full simulation time is used for the calculation
	max_LE += log(norm((h2 .- h) ./ p_factor))
	true_time = dt * (maxstep - startstep)
	return max_LE / true_time
end
