using PyPlot, Random

# getting the class in Bins
function classInBins(x::Float64)
    # class 0 implies that x is not in [x_min, x_max)
    # class i implies that x is in [x_min + dx*(i-1), x_min + dx*i)
    x_min, x_max, dx = -3.0, 3.0, 0.01
    if !(x_min <= x < x_max)
        return 0
    end
    y = (x - x_min)/dx
    return Int64(floor(y))+1
end

# getting the centre of mass from the vector z
function get_CM(z::Array{Float64,1})
    return sum(z)/length(z)
end

# the tent function, the theoretical solution for N=2 and N=infty
function tent(x::Float64)
    return abs(x)<=1 ? 1-abs(x) : 0
end

# running a simulation to get the empirical steady PDF
function get_OB_avg(seed::Int64=10, N::Int64=5)
    # initialization
    x_min, x_max, dx = -3.0, 3.0, 0.01
    rng = MersenneTwister(seed)
    z, ξ, L = zeros(N), zeros(N), 2.0
    OB, n_sample = zeros(600), 0
    
    dt = 0.0001/(N/2.0)^0.5 
    t_end = 10000
    
    ts = collect(-20.0:dt:t_end)    
    T, dt = ts[end], ts[2]-ts[1]
    dt_sqrt = sqrt(dt)
    
    # simulation loop starts
    for t in ts
        # the samples during t in [-20.0, 0.0) is discorded for the initial run
        # we take the samples during [0.0, t_end)
        if t >= 0.0
            z_CM = get_CM(z)
            classes = classInBins.(z.-z_CM)
            for class in classes
                if class >=1
                    OB[class] += 1.0
                end
            end
            n_sample += N
        end
        # transaction detection and resubmission jump
        i_max, i_min = argmax(z), argmin(z)
        if z[i_max]-z[i_min] >= L
            z[i_max] = z[i_min] = (z[i_max]+z[i_min])/2.0
        end
        # random walks
        randn!(rng,ξ)
        z += ξ*dt_sqrt
    end        
    return OB / n_sample / dx
end

# plot function
function plot()
    r_min, r_max, dr = -3.0, 3.0, 0.01
    rs = collect(r_min+dr/2.0:dr:r_max)
    skip = 1

    fig, axes = subplots(2,2,figsize=(10,6))
    axes[1].set_xlim(-1.5,1.5)
    axes[3].set_xlim(0.8,1.2)
    axes[3].set_ylim(-0.001,0.2)
    for (N, c) in zip([2,3,4,5,6,7], ["o","v","^","8","s","x"])
        OB = get_OB_avg(N,N)
        axes[1].plot(rs,OB,label="\$N=\$$N")
        axes[3].plot(rs,OB,c,label="\$N=\$$N")
    end
    axes[1].plot(rs,tent.(rs),color="lightskyblue",linestyle="dashed",label="Tent")
    axes[3].plot(rs,tent.(rs),color="lightskyblue",linestyle="dashed",label="Tent")
    axes[1].set_xlabel("Relative price \$r\$")
    axes[1].set_ylabel("\$\\phi(r)\$")
    axes[3].set_xlabel("Relative price \$r\$")
    axes[3].set_ylabel("\$\\phi(r)\$")
    axes[1].legend(loc="upper left")
    axes[3].legend(loc="upper right")

    axes[2].set_xlim(-1.5,1.5)
    axes[4].set_xlim(0.8,1.2)
    axes[4].set_ylim(-0.001,0.2)
    for (N, c) in zip([8,9,10,20,40,80], ["o","v","^","8","s","x"])
        OB = get_OB_avg(N,N)
        axes[2].plot(rs,OB,label="\$N=\$$N")
        axes[4].plot(rs,OB,c,label="\$N=\$$N")
    end
    axes[2].plot(rs,tent.(rs),color="lightskyblue",linestyle="dashed",label="Tent")
    axes[4].plot(rs,tent.(rs),color="lightskyblue",linestyle="dashed",label="Tent")
    axes[2].set_xlabel("Relative price \$r\$")
    axes[2].set_ylabel("\$\\phi(r)\$")
    axes[4].set_xlabel("Relative price \$r\$")
    axes[4].set_ylabel("\$\\phi(r)\$")
    axes[2].legend(loc="upper left")
    axes[4].legend(loc="upper right")
    tight_layout()
    savefig("OB_avg.pdf")
end

plot()