# FOR PROBLEM 8.1
# author: sotech117

# parameters
z = 4

function equation(s, T)
    return s - tanh(z * s / T)
end
function derivative(s, T)
    return 1 - z / T * sech(z * s / T)^2
end


# apply newton solver to find the solution
function newton_solver(f, df, x0, tol, max_iter, T = 1)
    x = x0
    for i in 1:max_iter
        if abs(df(x, T)) < 1e-6
            println("derivative is too small")
            return x
        end
        x = x - f(x, T) / df(x, T)
        if abs(f(x, T)) < tol
            return x
        end
    end
    return x
end

# initial guess
s0 = .5

# solve the equation
t = newton_solver(equation, derivative, s0, 1e-6, 100, 1)
println("s(T) = ", t)

# go over Ts from 0 - 8 and plot the solutions
using Plots
Ts = .001:0.001:8
ss = [newton_solver(equation, derivative, s0, 1e-6, 100, T) for T in Ts]
scatter(
    Ts, ss, xlabel="T (units of J / k_b)", ylabel="Magnetization", label="Numerical Solution", title="Magnetization vs Temp (Mean Field Theory)",
    msw=0, ms=1.5, mc=:blue, lc=:red, lw=1.5
)

# plot the analytical solution
function analytical_solution(T)
    n1 = (3 / T) * (T / z) ^ 3
    n2 = z - T
    if n1 < 0 || n2 < 0
        return 0
    end
    return sqrt(n1) * sqrt(n2)
end
plot!(
    Ts, [analytical_solution(T) for T in Ts], label="Analytical Solution",
    lc=:green, lw=2
)

# make a plot of the difference between the analytical and numerical solutions
diffs = [abs(analytical_solution(T) - newton_solver(equation, derivative, s0, 1e-6, 100, T)) for T in Ts]
plot!(Ts, diffs, label="Difference in Solutions", lc=:red, lw=1.5, ls=:dot)

# plot a vertical line where T = 3.75
vline!([3.9], label="T = 3.9", lc=:orange, lw=1.5, ls=:dash)

# save the plot
savefig("8-1.png")