"""
Contains functions for generating matrices with correlated or motif structure
"""

using LinearAlgebra
import Random.MersenneTwister

#------------------------------------------------------------------------------


"""
	generate_nv(N::Int64, T::Type; seed=nothing)

Returns a vector with entries i.i.d. on a zero-mean, unit-variance Gaussian
"""
function generate_nv(N::Int64, T::Type; seed=nothing)

	@assert T<:AbstractFloat "T must be a floating point type"
	rng = MersenneTwister(seed)

	return randn(rng, T, N)
end


"""
	generate_ONB(N::Int64, T::Type; seed=nothing)

Returns a random unitary matrix
"""
function generate_ONB(N::Int64, T::Type; seed=nothing)

	@assert T<:AbstractFloat "T must be a floating point type"
	rng = MersenneTwister(seed)
	Z = randn(rng, T, N, N)
	F = qr!(Z)

	return F.Q * diagm(0 => sign.(diag(F.R)))
end



"""
	generate_J(N::Int64, eta::T; seed=nothing) where T <: AbstractFloat

Returns a `N x N` Gaussian matrix with correlation between entries opposite the
diagonal equal to `eta` and diagonal entries themselves equal to `0`.

Entries of the matrix are scaled `sqrt(N)`.
"""
function generate_J(N::Int64, eta::T; seed=nothing) where T <: AbstractFloat

	rng = MersenneTwister(seed)

	J = zeros(T, N, N)
	M = [convert(T, 1)     convert(T, 0);
		 eta               sqrt(1-eta^2)]
	for j in 1:N-1
		for i in j+1:N
			(J[i,j], J[j,i]) = M * randn(rng, T, 2)
		end
	end

	return J ./ convert(T, sqrt(N))
end



"""
	generate_motif_P(N::Int64, kcon::T, kdiv::T, kchn::T; seed=nothing,
		no_autapses=true) where T <: AbstractFloat

Returns a rank-2 perturbation matrix that introduces converging, diverging, and
chain motifs when added to a dense random Gaussian matrix.

In dense Gaussian networks, motifs are interpreted as covariances between
synapses that share presynaptic or postsynaptic units. Use care when adding
perurbations to partially symmetric or antisymmetric matrices, as the final
symmetry depends on both the eta and kchn parameters.
"""
function generate_motif_P(N::Int64, kcon::T, kdiv::T, kchn::T; seed=nothing,
		no_autapses=true) where T <: AbstractFloat

	@assert kchn^2<=kcon*kdiv "κchn^2 cannot exceed κdiv*κcon"
	@assert kcon>=0 "Negative κcon is not allowed"
	@assert kdiv>=0 "Negative κdiv is not allowed"
	@assert kchn>=0 "Negative κchn is not allowed"

	rng = MersenneTwister(seed)

	if kcon == 0
		M = randn(rng, T, 2, N)
		a = zeros(N)
		b = sqrt(kdiv) .* reshape(M[2,:], (1, N))
	else
		M = [sqrt(kcon)          convert(T, 0);
			 kchn/sqrt(kcon)     sqrt(max(0, kdiv - kchn^2/kcon))]
		M *= randn(rng, T, 2, N)
		a = M[1,:]
		b = reshape(M[2,:], (1, N))
	end

	P = a .+ b

	if no_autapses
		P[diagind(P)] .= 0
	end

	return P ./ convert(T, sqrt(N))
end



"""
	generate_outlier_P(N::Int64, re::T, im::T, alpha::T; seed=nothing,
		no_autapses=true) where T <: AbstractFloat

Returns a rank-2 perturbation matrix that, when added to a random Gaussian
matrix, introduces a conjugate pair of outlier eigenvalues with given expected
real and imaginary parts.

The `alpha` parameter controls the strength of the uncorrelated components of
the outer product perturbations. This does not affect the expected value of the
real and imaginary parts.
"""
function generate_outlier_P(N::Int64, re::T, im::T, alpha::T; seed=nothing,
		no_autapses=true) where T <: AbstractFloat

	@assert re>0 "Real part of outliers must be positive"

	gamma = 2 * sqrt(1 /(1 + im^2 / re^2))
	rho = sqrt(2 * re / gamma)

	rng = MersenneTwister(seed)
	x1 = randn(rng, T, N)
	x2 = randn(rng, T, N)
	x3 = randn(rng, T, N)
	x4 = randn(rng, T, N)
	y1 = randn(rng, T, N)
	y2 = randn(rng, T, N)

	m1 = alpha .* x1 .+ rho .* y1
	m2 = alpha .* x2 .+ rho .* y2
	n1 = alpha .* x3 .+ rho .* y2 .+ rho * gamma .* y1
	n2 = alpha .* x4 .- rho .* y1

	P = m1 * n1' .+ m2 * n2'

	if no_autapses
		P[diagind(P)] .= 0
	end

	return P ./ N
end
