import numpy as np

class IsingLattice:

#We define variables we will need to call later,e.g. one for n_cycle and four for when calculating the average values
    E = 0.0
    E2 = 0.0
    M = 0.0
    M2 = 0.0

    n_cycles = 0

    def __init__(self, n_rows, n_cols):
        self.n_rows = n_rows
        self.n_cols = n_cols
        self.lattice = np.random.choice([-1,1], size=(n_rows, n_cols))

    def energy(self):
        "Return the total energy of the current lattice configuration."
        #Creating new lattices: 1 lattice with the rows rolled by one and 1 lattice with the columns rolled by one
        Ollielattice=np.roll(self.lattice,1,axis=0)
        Nikilattice=np.roll(self.lattice,1,axis=1)
        
        #Multiply these lattices with the original lattice - this gives us lattices which include values of the products of spins on adjacent cells
        Ollielattice=np.multiply(self.lattice,Ollielattice)
        Nikilattice=np.multiply(self.lattice,Nikilattice)
        
        
        #Total Energy is the sum of all the products of adjacent spins
        energy = -np.sum(Ollielattice+Nikilattice)

        return energy

    def magnetisation(self):
        "Return the total magnetisation of the current lattice configuration."
        magnetisation = np.sum(self.lattice)
        
        return magnetisation

    def montecarlostep(self, T):
        energy1 = self.energy()

        # complete this function so that it performs a single Monte Carlo step

        #the following two lines will select the coordinates of the random spin for you
        random_i = np.random.choice(range(0, self.n_rows))
        random_j = np.random.choice(range(0, self.n_cols))
        
        #the following will flip the selected random coordinates
        self.lattice[random_i,random_j]=-self.lattice[random_i,random_j]       
        #the following will calculate the energy E1 of this new configuration
        
        energy2 = self.energy()
        Edif = energy2-energy1
        
        if Edif >= 0:
            random_number = np.random.random()
            Boltzfactor=np.exp(-Edif/T)
            if random_number > Boltzfactor:
                self.lattice[random_i,random_j]=-self.lattice[random_i,random_j]            

        # the function should end by calculating and returning both the energy and magnetisation:
        energy = self.energy()
        magnetisation = self.magnetisation()
        
        
        self.n_cycles=self.n_cycles+1
        
        if self.n_cycles > 2000:
            self.E = self.E + self.energy()
            self.E2 = self.E2 + (self.energy())**2
            self.M = self.M + self.magnetisation()
            self.M2 = self.M2 + (self.magnetisation())**2
        
        return energy, magnetisation
        
    def statistics(self):
        
        #Defined new variable (n_cycles2) to clarify code - 
        #to show clearly that we are averaging over energy values only after 2000 cycles, 
        #and that we exporting the number of cycles we are averaging over (n_cyles2) 
        #instead of the number of cycles in the simulation (n_cyles) to ILtemperaturerange
        
        n_cycles2=self.n_cycles-2000
        # complete this function so that it calculates the correct values for the averages of E, E*E (E2), M, M*M (M2), and returns them
        aveE = self.E / n_cycles2
        aveE2 = self.E2 / n_cycles2
        aveM = self.M / n_cycles2
        aveM2 = self.M2 / n_cycles2
        
        #Added return of n_cycles because we need it when using ILtemperaturerange
        return aveE, aveE2, aveM, aveM2, n_cycles2