import numpy as np
from matplotlib.image import imread
from pathlib import Path

from MCFF.ising_model import all_up_state, all_down_state, random_state
from MCFF.ising_model import energy, energy_numpy

# create some test states
all_up = np.ones([100,100])
all_down = -np.ones([100,100])
random = np.random.choice([-1, 1], size = (100,100))
custom = (1 - 2*imread(Path(__file__).parents[2]/'learning/data/test_state.png')[:, :, 0]) #load a 100x100 png, take the red channel, remap 0,1 to -1,1

states = [all_up, all_down, random, custom]

def E_prediction_all_the_same(L):
    "The exact energy in for the case where all spins are up or down"
    return -(4*(L - 2)**2 + 12*(L-2) + 8) / L**2

def test_exact_energies():
    for state in [all_up, all_down]:
        L = state.shape[0]
        assert energy(state) == E_prediction_all_the_same(L)

def test_energy_implementations():
    for state in states:
        assert np.allclose(energy(state), energy_numpy(state))