# my neural net from 12.11

# first, store the memory of A as a lattice of 1, -1, where 1 is the memory and -1 is the absence of memory
A = [
	-1 -1 -1 -1 1 1 -1 -1 -1 -1;
	-1 -1 -1 1 1 1 1 -1 -1 -1;
	-1 -1 1 1 -1 -1 1 1 -1 -1;
	-1 1 1 -1 -1 -1 -1 1 1 -1;
	1 1 -1 -1 -1 -1 -1 -1 1 1;
	1 1 -1 -1 -1 -1 -1 -1 1 1;
	-1 1 1 -1 -1 -1 -1 1 1 -1;
	-1 -1 1 1 -1 -1 1 1 -1 -1;
	-1 -1 -1 1 1 1 1 -1 -1 -1;
	-1 -1 -1 -1 1 1 -1 -1 -1 -1
]

B = [
	1 1 1 1 1 -1 -1 -1 -1 -1;
	1 1 -1 -1 1 1 -1 -1 -1 -1;
	1 1 -1 -1 1 1 -1 -1 -1 -1;
	1 1 -1 -1 1 1 -1 -1 -1 -1;
	1 1 1 1 1 -1 -1 -1 -1 -1;
	1 1 1 1 1 -1 -1 -1 -1 -1;
	1 1 -1 -1 1 1 -1 -1 -1 -1;
	1 1 -1 -1 1 1 -1 -1 -1 -1;
	1 1 -1 -1 1 1 -1 -1 -1 -1;
	1 1 1 1 1 -1 -1 -1 -1 -1
]

C = [
	-1 1 1 1 1 1 1 -1 -1 -1;
	1 1 1 1 1 1 1 -1 -1 -1;
	1 1 1 -1 -1 -1 -1 -1 -1 -1;
	1 1 -1 -1 -1 -1 -1 -1 -1 -1;
	1 1 -1 -1 -1 -1 -1 -1 -1 -1;
	1 1 -1 -1 -1 -1 -1 -1 -1 -1;
	1 1 -1 -1 -1 -1 -1 -1 -1 -1;
	1 1 1 -1 -1 -1 -1 -1 -1 -1;
	1 1 1 1 1 1 1 -1 -1 -1;
	-1 1 1 1 1 1 1 -1 -1 -1
]

D = [
	1 1 1 1 1 -1 -1 -1 -1 -1;
	1 1 -1 -1 1 1 -1 -1 -1 -1;
	1 1 -1 -1 1 1 -1 -1 -1 -1;
	1 1 -1 -1 1 1 -1 -1 -1 -1;
	1 1 -1 -1 1 1 -1 -1 -1 -1;
	1 1 -1 -1 1 1 -1 -1 -1 -1;
	1 1 -1 -1 1 1 -1 -1 -1 -1;
	1 1 -1 -1 1 1 -1 -1 -1 -1;
	1 1 -1 -1 1 1 -1 -1 -1 -1;
	1 1 1 1 1 -1 -1 -1 -1 -1
]

E_letter = [
	1 1 1 1 1 1 1 1 1 1;
	1 1 -1 -1 -1 -1 -1 -1 -1 -1;
	1 1 -1 -1 -1 -1 -1 -1 -1 -1;
	1 1 1 1 1 1 1 1 1 1;
	1 1 -1 -1 -1 -1 -1 -1 -1 -1;
	1 1 -1 -1 -1 -1 -1 -1 -1 -1;
	1 1 1 1 1 1 1 1 1 1;
	1 1 -1 -1 -1 -1 -1 -1 -1 -1;
	1 1 -1 -1 -1 -1 -1 -1 -1 -1;
	1 1 1 1 1 1 1 1 1 1
]

F = [
	1 1 1 1 1 1 1 1 1 1;
	1 1 -1 -1 -1 -1 -1 -1 -1 -1;
	1 1 -1 -1 -1 -1 -1 -1 -1 -1;
	1 1 1 1 1 1 1 1 1 1;
	1 1 -1 -1 -1 -1 -1 -1 -1 -1;
	1 1 -1 -1 -1 -1 -1 -1 -1 -1;
	1 1 -1 -1 -1 -1 -1 -1 -1 -1;
	1 1 -1 -1 -1 -1 -1 -1 -1 -1;
	1 1 -1 -1 -1 -1 -1 -1 -1 -1;
	1 1 -1 -1 -1 -1 -1 -1 -1 -1
]

G = [
	1 1 1 1 1 1 1 1 1 1;
	1 1 -1 -1 -1 -1 -1 -1 -1 -1;
	1 1 -1 -1 -1 -1 -1 -1 -1 -1;
	1 1 -1 -1 -1 -1 -1 -1 -1 -1;
	1 1 -1 -1 -1 -1 -1 1 1 1;
	1 1 -1 -1 -1 -1 -1 -1 1 1;
	1 1 -1 -1 -1 -1 -1 -1 1 1;
	1 1 -1 -1 -1 -1 -1 -1 1 1;
	-1 1 1 1 1 1 1 1 1 1;
	-1 -1 1 1 1 1 1 1 1 1
]

encodings = [A, B, C, D, E_letter, F, G]

spin_numbers(row, col) = (row - 1) * 10 + col

# create strength of interactions between ith and jth spinds in Ji,j
J = zeros(100, 100)
for m in 1:100 # row
	for n in 1:100 # col
		i = (m - 1) % 10 + 1
		j = (n - 1) % 10 + 1
		k = (m - 1) ÷ 10 + 1
		l = (n - 1) ÷ 10 + 1

		J[m, n] = 0
		for encoding in encodings
			J[m, n] += encoding[i, k] * encoding[j, l]
		end
		J[m, n] /= length(encodings)
	end
end

# the energy function
function energy(s, J)
	E = 0
	for i in 1:100
		for j in 1:100
			E += J[i, j] * s[i] * s[j]
		end
	end
	return -E
end

# the update function (uses monte carlo steps)
function monte_carlo(s, J)
	for i in 1:100 # systematically go through each point in the lattice
		# calculate the energy of the system
		E = energy(s, J)
		# randomly flip a spin
		s[i] = -s[i]
		# calculate the new energy of the system
		E_new = energy(s, J)
		# calculate the change in energy
		dE = E_new - E
		# if the change in energy is positive, flip the spin back
		if dE > 0
			s[i] = -s[i]
		end
	end
	return s
end

# create the main function
function main(s, J, nsteps)
	E = energy(s, J)
	for i in 1:nsteps
		s = monte_carlo(s, J)
		E = energy(s, J)
	end
	return s, E
end

# run the main function

# randomly change some values in A to see if NN works
function produce_test_arr(enc, prob_change = 0.1)
	tmp = zeros(10, 10)
	for i in 1:10
		for j in 1:10
			if rand() < prob_change
				tmp[i, j] = rand([-1, 1]) # set some random values
			else
				tmp[i, j] = enc[i, j]
			end
		end
	end
	return tmp
end

function check_if_same(s, enc)
	for i in 1:10
		for j in 1:10
			if s[i, j] != enc[i, j]
				println("The neural net did not work")
				return false
			end
		end
	end
	# println("The neural net worked")
	return true
end

function apply_damage_to_J(J, prob_damage = 0.8)
	println("Applying damage to J with probability ", prob_damage)
	ret = copy(J)
	for i in 1:100
		for j in 1:100
			if rand() < prob_damage
				ret[i, j] = 0
			else
				ret[i, j] = J[i, j]
			end
		end
	end
	return ret
end

function run_tests(num_times, Js, MC_steps = 2)
	num_correct = 0
	num_total = 0

	for i in 1:num_times
		# randomly select a memory
		enc = encodings[rand(1:length(encodings))]
		# produce a test array
		s = produce_test_arr(enc)
		# run theNN 
		s, E = main(s, Js, MC_steps)
		# check if the NN worked
		if check_if_same(s, enc)
			num_correct += 1
		end
		num_total += 1
	end

	println("Number of correct tests: ", num_correct)

	return num_correct / num_total
end

function run_tests_for_damage_range(damages, num_times, MC_steps = 2, init_J = J)
	res = []

	for damage in damages
		damaged_J = apply_damage_to_J(init_J, damage)
		p = run_tests(num_times, damaged_J, MC_steps)

		push!(res, p)
	end

	return res
end

damage_range = collect(0.0:0.1:0.15)
res = run_tests_for_damage_range(damage_range, 10)

#plot the results
using Plots
plot(damage_range, res, xlabel = "Damage", ylabel = "Success rate", title = "NN success rate vs damage probability", marker = :circle, label = "MC steps = 2")

res = run_tests_for_damage_range(damage_range, 10, 10)
plot!(damage_range, res, xlabel = "Damage", ylabel = "Success rate", title = "NN success rate vs damage probability", marker = :circle, label = "MC steps = 10")

res = run_tests_for_damage_range(damage_range, 10, 100)
plot!(damage_range, res, xlabel = "Damage", ylabel = "Success rate", title = "NN success rate vs damage probability", marker = :circle, label = "MC steps = 100")

res = run_tests_for_damage_range(damage_range, 10, 1000)
plot!(damage_range, res, xlabel = "Damage", ylabel = "Success rate", title = "NN success rate vs damage probability", marker = :circle, label = "MC steps = 1000")

savefig("12-12-all.png")