# FOR PROBLEM 8.11
# author: sotech117

using Statistics 
using Plots

function wrap_index(i::Int, l::Int)::Int
	wrap = (i - 1) % l + 1
	return (wrap <= 0) ? l + wrap : wrap
end

mutable struct Ising2D
	l::Int 
	n::Int 
	temperature::Float64
	w::Vector{Float64} # Boltzmann weights
	state::Matrix 
	energy::Int
	magnetization::Int
	mc_steps::Int 
	accepted_moves::Int 
	energy_array::Vector{Int}
	magnetization_array::Vector{Int}
end 

Ising2D(l::Int, temperature::Float64) = begin	
	n = l^2 
	w = zeros(9)
	w[9] = exp(-8.0 / temperature)
	w[5] = exp(-4.0 / temperature)
	state = ones(Int, l, l) # initially all spins up
	energy = -2 * n
	magnetization = n
	return Ising2D(l, n, temperature, w, state, energy, magnetization, 0, 0, 
				   Int[], Int[])
end

function reset!(ising::Ising2D)
	ising.mc_steps = 0 
	ising.accepted_moves = 0 
	ising.energy_array = Int[] 
	ising.magnetization_array = Int[] 
end

function mc_step!(ising::Ising2D) 
	l::Int = ising.l 
	n::Int = ising.n 
	w = ising.w 

	state = ising.state 
	accepted_moves = ising.accepted_moves
	energy = ising.energy 
	magnetization = ising.magnetization

	random_positions = l * rand(2 * n)
	random_array = rand(n) 

	for k in 1:n
		i = trunc(Int, random_positions[2 * k - 1]) + 1 
		j = trunc(Int, random_positions[2 * k]) + 1 

		de = 2 * state[i, j] * (state[i % l + 1, j] + 
				state[wrap_index(i - 1, l), j] + state[i, j % l + 1] + 
				state[i, wrap_index(j - 1, l)])

		if de <= 0 || w[de + 1] > random_array[k]
			accepted_moves += 1 
			new_spin = - state[i, j] # flip spin
			state[i, j] = new_spin 
			energy += de 
			magnetization += 2 * new_spin
		end

	end

	ising.state = state 
	ising.accepted_moves = accepted_moves
	ising.energy = energy 
	ising.magnetization = magnetization 

	append!(ising.energy_array, ising.energy) 
	append!(ising.magnetization_array, ising.magnetization)
	ising.mc_steps = ising.mc_steps + 1
end

function steps!(ising::Ising2D, num::Int=100)
	for i in 1:num
		mc_step!(ising)
	end 
end

function mean_energy(ising::Ising2D)
	return mean(ising.energy_array) / ising.n
end 

function specific_heat(ising::Ising2D)
	return (std(ising.energy_array) / ising.temperature) ^ 2 / ising.n
end

function mean_magnetization(ising::Ising2D)
	return mean(ising.magnetization_array) / ising.n
end

function susceptibility(ising::Ising2D)
	return (std(ising.magnetization_array)) ^ 2 / (ising.temperature * ising.n)	
end

function observables(ising::Ising2D)
	printstyled("Temperature\t\t", bold=true)
	print(ising.temperature); print("\n")

	printstyled("Mean Energy\t\t", bold=true)
	print(mean_energy(ising)); print("\n")

	printstyled("Mean Magnetiz.\t\t", bold=true)
	print(mean_magnetization(ising)); print("\n")

	printstyled("Specific Heat\t\t", bold=true)
	print(specific_heat(ising)); print("\n")

	printstyled("Susceptibility\t\t", bold=true)
	print(susceptibility(ising)); print("\n")

	printstyled("MC Steps\t\t", bold=true)
	print(ising.mc_steps); print("\n")
	printstyled("Accepted Moves\t\t", bold=true)
	print(ising.accepted_moves); print("\n")
end


function plot_ising(state::Matrix{Int})
	pos = Tuple.(findall(>(0), state))
	neg = Tuple.(findall(<(0), state))
	scatter(pos, markersize=5)
	scatter!(neg, markersize=5)
end


function get_magnetization(T, n=1000)
    m = Ising2D(64, T)
    steps!(m, n)

    println("done with T = $T")
    return mean_magnetization(m)
end

Ts = 0:.1:5
ms = [abs(get_magnetization(T)) for T in Ts]

println("done with calculating magnetizations")

function linear_regression(x, y)
    n = length(x)
    x̄ = sum(x) / n
    ȳ = sum(y) / n
    a = sum((x .- x̄) .* (y .- ȳ)) / sum((x .- x̄).^2)
    b = ȳ - a * x̄
    return (a, b)
end

# plot M^(8) over T
betas = .001:.001:1
residuals = []
for i in 1:length(betas)
    b = betas[i]
    m = ms .^ (1 / b)
    # filter out the zero values


    s = scatter(p,
        Ts, m, xlabel="T (units of J / k_b)", ylabel="Magnetization", label="$b-beta", title="Magnetization vs Temp (Ising Monte Carlo)",
        msw=0, ms=1.5, mc=:red, lc=:red, lw=1.5, legend=:bottomleft
    )


    # do a linear regression
    a, b = linear_regression(Ts, m)

    # plot a linear regression line
    plot!(s, Ts, a*Ts .+ b, label="Linear Regression", lw=1.2, color=:red, linestyle=:dash)

    # calculate the residuals
    push!(residuals, sum((m .- (a*Ts .+ b)).^2))


    savefig(s, "hw6/b/8-2-$i.png")
end

# plot the residuals over beta
plot(betas, residuals, xlabel="beta", ylabel="Squared Distance", label="Residuals", title="Error from Linear Regression of M^(1/Beta)", lw=1.2, lc=:red, legend=:topright)
# find the min on the first half of the residuals
min_residuali = argmin(residuals[1:div(length(residuals), 2)])
min_residual = betas[min_residuali]
println("Minimum Residual: ", min_residual)
vline!([min_residual], label="Minimum Point @ Beta = $min_residual", lc=:orange, lw=1.5, ls=:dash)
# plot the analityical beta of .125
vline!([.125], label="Analytical Minimum Beta = .125", lc=:green, lw=1.5, ls=:dot)
savefig("hw6/8-2-residuals-100.png")