#!/Applications/Julia-1.7.app/Contents/Resources/julia/bin/julia

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::Float64
	magnetization::Int
	mc_steps::Int 
	accepted_moves::Int 
	energy_array::Vector{Float64}
	magnetization_array::Vector{Int}
    H::Float64
end 

Ising2D(l::Int, temperature::Float64, H=1.0) = 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 = Float64(-2 * n + 2 * H * n)
	magnetization = n
	return Ising2D(l, n, temperature, w, state, energy, magnetization, 0, 0, 
				   Int[], Int[], H)
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 

        changed_spins = 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)])
		de = 2 * changed_spins + 2 * ising.H * state[i, j]

		if de <= 0 || rand() < exp(-de / ising.temperature)
			accepted_moves += 1 
			new_spin = - state[i, j] # flip spin
			state[i, j] = new_spin

            # add the effects of the 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 find_m(H::Float64, l::Int, num::Int, T::Float64)
    m = Ising2D(l, T, H)
    steps!(m, num)
    print("T = $T\n")
    print("H = $H\n")
    print("Mean Energy: $(mean_energy(m))\n")
    print("Mean Magnetization: $(mean_magnetization(m))\n\n")
    return mean_magnetization(m)
end

function map_h_to_m(H_range::Vector{Float64}, l::Int, num::Int, T::Float64)
    m = []
    for H in H_range
        push!(m, find_m(H, l, num, T))
    end
    return m
end

function do_linear_regression(x::Vector{Float64}, y::Vector{Float64})
    n = length(x)
    x̄ = mean(x)
    ȳ = mean(y)
    Σxy = sum((x .- x̄) .* (y .- ȳ))
    Σx² = sum((x .- x̄) .^ 2)
    b = Σxy / Σx²
    a = ȳ - b * x̄
    return a, b
end

function plot_log_of_m_and_h(H_range::Vector{Float64}, l::Int, num::Int, T=2.27)
    m = map_h_to_m(H_range, l, num, T)
	p = scatter(H_range, m, label="M vs H", xlabel="H", ylabel="M", title="Magnetization (M) vs Field (B) for Ising Model at T_c", scale=:ln)

    # get the linear regression of the log
	log_h = log.(H_range)
	log_m = log.(m)
	a, b = do_linear_regression(log_h, log_m)
	println("a: $a, b: $b")
	# plot the linear regression
	plot!(p, H_range, exp.(a) .* H_range .^ b, label="linear regression = $(round(a, digits=3)) + $(round(b, digits=3))x", line=:dash, color=:red)

    return p
end

# textbook rec
h_range = .02:.02:.2
h_range = collect(h_range)
T = 2.27 # T_c for this system
side = 64
steps = 5000
plot_log_of_m_and_h(h_range, side, steps)