using Color
using Gadfly
using Interact
using ProgressMeter

n=500
M=randn(n, n)
M=(M+M')/√2n

xg = linspace(-2, 2, 100)
plot(Coord.Cartesian,
    Guide.xlabel("Eigenvalue"), Guide.ylabel("Spectral density"),
    layer(x=xg, y=√abs(4 - xg.^2)/(2π), Geom.line,
        Theme(default_color=color("black"))
    ),
layer(x=eigvals(M), Geom.histogram(bincount=round(Int,√n), density=true))
)

function wigner55(n::Int, N::Int=n, v::Real=1.0)
    T = promote_type(typeof(v*1), typeof(n))
    H = zeros(T, 2n+1, 2n+1)
    for i=1:2n+1
        H[i,i] = i-1-n
    end
    for m = 1:N, j=1:2n+1-m
        H[j,j+m] = v*(2randbool()-1)
    end
    Symmetric(H, :U)
end
wigner55(2, 1, 1)

function demo1(n::Int=20)
    @manipulate for v=0.0:0.25:5.0, k=1:2n+1
        E = eigfact!(wigner55(n, 1, v))
        plot(Coord.Cartesian(ymin=-1, ymax=1),
            Guide.ylabel("Eigenvector"), Guide.xlabel("Position"),
            layer(x=-n:n, y=E[:vectors][:,k],
                Geom.point, Theme(default_color=color("black"))),
            layer(x=-n:n, y=map(m->abs(besselj(m-k+(n+1), 2v)), -n:n),
                Geom.line, Theme(default_color=color("lightgrey"))),
            layer(x=-n:n, y=map(m->-abs(besselj(m-k+(n+1), 2v)), -n:n),
                Geom.line, Theme(default_color=color("lightgrey"))),
        )
    end
end
demo1()

function wigner55_2(n::Int, N::Int=n, v::Real=1.0)
    T = promote_type(typeof(v*1), typeof(n))
    H = zeros(T, 2n+1, 2n+1) #Diagonals are now 0
    for m = 1:N, j=1:2n+1-m
        H[j,j+m] = v*(2randbool()-1)
    end
    Symmetric(H, :U)
end
wigner55_2(3, 6, 1)

t = 1000 #Number of trials
n = 30 
v = 4 #Magnitude of off-diagonal
evals = zeros(2n+1, t)

@showprogress for i=1:t
    H1 = wigner55_2(n, 2n, v)
    evals[:, i] = eigvals(H1)
end

xg = linspace(-v*√8n, v*√8n, 100)
plot(Coord.Cartesian(aspect_ratio=φ),
    Guide.xlabel("Eigenvalue"), Guide.ylabel("Spectral density"),
    layer(x=xg, y=√abs(8n*v^2 - xg.^2)/(4π*n*v^2), Geom.line,
        Theme(default_color=color("black"))
    ),
    layer(x=evals, Geom.histogram(bincount=round(Int,√(t*(2n+1))), density=true))
)

q = 100
n = 50
N = 2n
t = 10 #Number of trials
evals = zeros(2n+1, t)

@showprogress for i=1:t
    H = wigner55(n, N, √(q*N))
    evals[:, i] = eigvals(H)
end

v=√(q*N)
xg = linspace(-v*√4N, v*√4N, 100)
plot(Coord.Cartesian(aspect_ratio=φ),
    Guide.xlabel("Eigenvalue"), Guide.ylabel("Spectral density"),
    layer(x=xg, y=√abs(4N*v^2 - xg.^2)/(2π*N*v^2), Geom.line,
        Theme(default_color=color("black"))
    ),
    layer(x=evals, Geom.histogram(bincount=round(Int,√(t*(2n+1))), density=true))
)

@manipulate for logq = -5:0.1:5
    
    q = 10.0^logq
    n = 50
    N = 2n
    t = 10 #Number of trials
    evals = zeros(2n+1, t)

    for i=1:t
        H = wigner55(n, N, √(q*N))
        evals[:, i] = eigvals(H)
    end    

    v=√(q*N)
    xg = linspace(-v*√4N, v*√4N, 100)
    plot(Coord.Cartesian(aspect_ratio=φ),
        Guide.xlabel("Eigenvalue"), Guide.ylabel("Spectral density"),
        layer(x=xg, y=√abs(4N*v^2 - xg.^2)/(2π*N*v^2), Geom.line,
    Theme(default_color=color("grey"))
        ),
    layer(x=[-n; -n:n; n], y=[0; fill(1/(2n+1), 2n+1); 0], Geom.line,
    Theme(default_color=color("grey"))
        ),
        layer(x=evals, Geom.histogram(bincount=round(Int,√(t*(2n+1))), density=true))
    )
end

function samplefree(n, N, t, q)
    v=√(q*N)

    evals1 = zeros(2n+1, t)
    evals2 = zeros(2n+1, t)
    
    #Sample eigenvalues
    @showprogress for i=1:t
        𝐤 = Diagonal(1.0*[-n:n;])
        𝐯 = wigner55_2(n,2n,v) |> full
        Q = full(qrfact!(randn(2n+1, 2n+1))[:Q])
        evals1 = eigvals(𝐤 + 𝐯)
        evals2 = eigvals(Q*full(𝐤)*Q' + 𝐯)
    end
    
    #Compute normalized histograms
    xv, yv1 = hist(evals1, round(Int, √(n*√t)))
    yv1 /= sum(yv1)*diff(xv)[1]
    xv, yv2 = hist(evals2, xv)
    yv2 /= sum(yv2)*diff(xv)[1]

    midpoints(xv), yv1, yv2
end

function plotdensities(xv, yv1, yv2)
    cm = map(x->AlphaColorValue(x, 0.5), palette("Set1", 3)[1:2])
    plot(Guide.xlabel("Eigenvalue"), Guide.ylabel("Density"),
        layer(x=xv, y=yv1, Geom.bar, Theme(default_color=cm[1])),
        layer(x=xv, y=yv2, Geom.bar, Theme(default_color=cm[2])),
        Guide.manual_color_key("", ["exact", "free"], cm), 
    )
end

@manipulate for logq = -5:0.5:5
    xv, yv1, yv2 = samplefree(30, 2n, 100, 10.0^logq)
    plotdensities(xv, yv1, yv2)
end