from itertools import product
from tabulate import tabulate
import matplotlib.pyplot as plt
import numpy as np


def Ising_2x2():
    table = []
    head = ['S1', 'S2', 'S3', 'S4', 'Eij(J)', 'E(J)', 'M']
    for s1, s2, s3, s4 in product([-1, 1], repeat=4):
        S = np.array([[s1, s2], [s3, s4]])
        M = sum(sum(S))
        energies = calculate_energies(S)

        table.append([s1, s2, s3, s4, tuple(energies), sum(energies), M])
    return tabulate(table, headers=head)


def calculate_energies(S):
    energies = []
    n, m = S.shape
    for row in range(n):
        for col in range(m):
            if col + 1 != m:
                energies.append(-S[row][col] * S[row][col + 1])
            if row + 1 != n:
                energies.append(-S[row][col] * S[row + 1][col])
    return energies


def calculate_total_E_M(S):
    E = 0
    M = sum(sum(S))
    n, m = S.shape
    for row in range(n):
        for col in range(m):
            '''
            row_1 = (row + 1) % n
            col_1 = (col + 1) % m
            '''
            if col + 1 != m:
                E += -S[row][col] * S[row][col + 1]
            if row + 1 != n:
                E += -S[row][col] * S[row + 1][col]
    return E, M


def calculate_energy_change(S, index):
    i, j = index
    n, m = S.shape
    dE = 0
    s = S[i][j]
    if i > 0:
        dE += 2 * s * S[i - 1][j]
    if i < n - 1:
        dE += 2 * s * S[i + 1][j]
    if j > 0:
        dE += 2 * s * S[i][j - 1]
    if j < n - 1:
        dE += 2 * s * S[i][j + 1]
    return dE


def test_2x2():
    table = []
    head = ['S1', 'S2', 'S3', 'S4', 'E(J)', 'M']
    for s1, s2, s3, s4 in product([-1, 1], repeat=4):
        S = np.array([[s1, s2], [s3, s4]])
        E, M = calculate_total_E_M(S)
        table.append([s1, s2, s3, s4, E, M])
    print tabulate(table, headers=head)


def test_examplo_0_19():
    S = np.resize(np.genfromtxt('exemplo-0-19.dat',
                                usecols=-1), (20, 20))  # [:, ::-1]
    # print tabulate(S)
    E, M = calculate_total_E_M(S)
    print 'E(J)  M'
    print E, ' ', M


def test_energy_change():
    S = np.resize(np.genfromtxt('exemplo-0-19.dat',
                                usecols=-1), (20, 20))  # [::-1, :]
    for index in [(0, 0), (0, 12), (14, 14), (10, 10), (7, 9), (5, 11)]:
        # i, j = index
        print 'Spin flipped:', index, 'previous spin', S[index]
        print calculate_energy_change(S, index)


def problem_1():
    print 'Problem 1:\n'
    print Ising_2x2()


def problem_2():
    print 'Problem 2:\n'
    print 'a) Test the 2x2 ising model:'
    test_2x2()
    print '\nb) Test the exaple given (20x20):'
    test_examplo_0_19()
    print '\nTesting the flipping of spin:'
    test_energy_change()


if __name__ == '__main__':
    problem_1()
    print '\n'
    problem_2()
