"""
Contains functions that take in time-invariant network structure parameters,
such as the coupling matrix J and gain parameter g, and run simulations of
network activity. Includes functions that numerically calculate network
properties (such as Lyapunov spectrum) in parallel with the simulation.
"""

using LinearAlgebra

#-------------------------------------------------------------------------------
"""
	sim_activity(J::Matrix{T}, g::T, h0::Vector{T}; dt=0.1, t_bound=1000.0,
		s=0.0, 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`.

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.
"""
function sim_activity(J::Matrix{T}, g::T, h0::Vector{T}; dt=0.1, t_bound=1000.0,
		s=0.0, 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)

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

	# Allocate memory for output
	out = zeros((N, Int(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 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, dt=0.1, t_bound=1000.0, t_ONS=1.0, wait=0.0,
		return_max_LE=false, printstuff=false) 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, dt=0.1, t_bound=1000.0, t_ONS=1.0, wait=0.0,
		return_max_LE=false, printstuff=false) where T <: AbstractFloat

	N = size(J)[1]
	# Check dimensional consistency
	@assert size(J) == (length(h0), length(dh)) == size(Q0) == (N,N)
	@assert T == typeof(s) == typeof(dt) == typeof(t_bound) == typeof(t_ONS) == typeof(wait)

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

	maxstep = Int(t_bound / dt) # When to stop the simulation
	startstep = Int(wait / dt) # When to start tracking Lypunov exponents
	tONSstep = Int(t_ONS / dt) # Interval between re-orthonormalization
	try
		Int(t_bound/t_ONS)
		Int(wait/t_ONS)
	catch e
		if isa(e, InexactError)
			msg = "t_ONS must divide t_bound and wait"
			throw(ArgumentError(msg))
		end
	end

	# 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.0
		p_factor = T(1e-6)  # Order of the perturbation handled by the sim
		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)
		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)
			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 > 100 || scale < 0.01 || 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
	if return_max_LE
		max_LE += log(norm((h2 .- h) ./ p_factor))
		return (spectrum ./ (t_bound - wait), max_LE / (t_bound - wait))
	else
		return spectrum ./ (t_bound - wait)
	end
end



"""
	sim_direct_max_LE(J::Matrix{Float64}, g::Float64, h0::Vector{Float64},
		dh::Vector{Float64}; dt=0.1, t_bound=1000.0, wait=0.0)

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, dt=0.1, t_bound=1000.0, wait=0.0, printstuff=false) where T <: AbstractFloat

	N = size(J)[1]
	# Check dimensional consistency
	@assert size(J) == (length(h0), length(dh)) == (N,N)
	p_factor = 10^-10  # Order of the perturbation handled by the sim

	# Allocate memory for output
	max_LE = 0.0

	maxstep = Int(t_bound/dt) # When to stop the simulation
	startstep = Int(wait/dt) # When to start tracking Lypunov exponents

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

	r = Vector{Float64}(undef, N)
	r2 = Vector{Float64}(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*g, J, r, 1.0-dt, h)
		BLAS.gemv!('N', dt*g, J, r2, 1.0-dt, h2)
		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 > 10 || scale < 0.1 || step == startstep
			printstuff && println("Time step: ", step)
			copyto!(h2, h)
			dh ./= (scale / p_factor)
			h2 .+= dh
			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))
	return max_LE / (t_bound - wait)
end


"""
    sim_lyapunov(J::Matrix{Float64}, g::Float64, A::Matrix{Float64},
        onb::Matrix{Float64}; dt=0.1, t_bound=1000.0, t_ONS=1.0, wait=0.0,
		printstuff=false)

Calculates the Lyapunov spectrum of a tanh rate network with coupling matrix
`J` and gain parameter `g` by integrating an orthonormal system according to
the Jacobians in the tangent space along a trajectory given by activity
matrix `A`.
"""
function sim_lyapunov(J::Matrix{Float64}, g::Float64, A::Matrix{Float64},
		onb::Matrix{Float64}; dt=0.1, t_bound=1000.0, t_ONS=1.0, wait=0.0,
		printstuff=false)

	N = size(J)[1]
	# Check dimensional consistency
	@assert(size(J)==(size(A,1),N)==size(onb)==(N,N))

	# Ranges that index the diagonal entries of useful matrix sizes
	diagonals = diagind(J)

	# Allocate memory for output
	spectrum = zeros(N)

	maxstep = Int(t_bound/dt) # When to stop the simulation
	startstep = Int(wait/dt) # When to start tracking Lypunov exponents
	tONSstep = Int(t_ONS/dt) # Interval between re-orthonormalization
	@assert maxstep==size(A,2) "Simulation length must match time parameters"
	try
		Int(t_bound/t_ONS)
		Int(wait/t_ONS)
	catch e
		if isa(e, InexactError)
			msg = "t_ONS must divide t_bound and wait"
			throw(ArgumentError(msg))
		end
	end

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

	J = g .* J

	step = 0
	while step < maxstep
		step += 1

		#Evolve the trajectory and the orthonormal basis in the tangent space
		map!(sech2, sech2h, view(A, :, step))
		copyto!(D, J)
		D .*= sech2h
		D[diagonals] .-= 1.0
		BLAS.gemm!('N', 'N', dt, D, Q, 0.0, dQ)
		Q .+= dQ

		# 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, diagonals)
			if step > startstep  # forcing positive R-diagonal elements
				spectrum .+= log.(abs.(diags))
			end
			fill!(Q, 0.0)
			map!(sign, view(Q, diagonals), 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
	return spectrum ./ (t_bound - wait)
end


"""
	sim_lyapunov_halfspec(J::Matrix{Float64}, g::Float64,
		A::Matrix{Float64}, onb::Matrix{Float64}; dt=0.1, t_bound=1000.0,
		t_ONS=1.0, wait=0.0, printstuff=false)

Calculates the Lyapunov spectrum of a tanh rate network with coupling matrix
`J` and gain parameter `g` by integrating an orthonormal system according to
the Jacobians in the tangent space along a trajectory given by activity
matrix `A`. Saves time by only calculating half of the Lyapunov spectrum
"""
function sim_lyapunov_halfspec(J::Matrix{Float64}, g::Float64,
		A::Matrix{Float64}, onb::Matrix{Float64}; dt=0.1, t_bound=1000.0,
		t_ONS=1.0, wait=0.0, printstuff=false)

	N = size(J)[1]
	@assert N%2==0 "Calculating half the spectrum only works for even N"
	Nh = Int(N/2)
	# Check dimensional consistency
	@assert(size(J)==(size(A,1),N)==size(onb)==(N,N))

	# Ranges that index the diagonal entries of useful matrix sizes
	diagonals = diagind(J)
	diags_1halfrng = range(1, step=N+1, length=Nh)
	diags_2halfrng = range(1 + Nh * (N + 1), step=N+1, length=Nh)

	# Allocate memory for output
	spectrum = zeros(Nh)

	maxstep = Int(t_bound/dt) # When to stop the simulation
	startstep = Int(wait/dt) # When to start tracking Lypunov exponents
	tONSstep = Int(t_ONS/dt) # Interval between re-orthonormalization
	@assert maxstep==size(A,2) "Simulation length must match time parameters"
	try
		Int(t_bound/t_ONS)
		Int(wait/t_ONS)
	catch e
		if isa(e, InexactError)
			msg = "t_ONS must divide t_bound and wait"
			throw(ArgumentError(msg))
		end
	end

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

	J = g .* J

	step = 0
	while step < maxstep
		step += 1

		#Evolve the orthonormal basis in the tangent space
		map!(sech2, sech2h, view(A, :, step))
		copyto!(D, J)
		D .*= sech2h
		D[diagonals] .-= 1.0
		BLAS.gemm!('N', 'N', dt, D, Q, 0.0, dQ)
		Q .+= dQ

		# 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!(Qh, view(Q, :, 1:Nh))
			F = qr!(Qh)
			diags = view(F.factors, diags_1halfrng)
			if step > startstep  # forcing positive R-diagonal elements
				spectrum .+= log.(abs.(diags))
			end
			fill!(Q, 0.0)
			map!(sign, view(Q, diags_1halfrng), diags)
			Q[diags_2halfrng] .= 1.0
			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
	return spectrum ./ (t_bound - wait)
end
