import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
from random import random, seed
from itertools import product
from collections import defaultdict
from math import exp
from tabulate import tabulate


# ---------------------------------------------------------------
# VIBRATING STRINGS FUNCTION

# This function initializes the narray that you use for using solve_1d_string
# dimensions is a tuple with (N_t, N_x)
# init_cond is a tuple with in the first element the the function at the
# beginning, the second and third are the boundary conditions.

def initialize_string_wave(dimensions, init_cond):
    u = np.zeros(dimensions)
    u[0, :] = init_cond[0]
    u[:, 0] = init_cond[1]
    u[:, -1] = init_cond[2]
    return u


# solve_1d_string solve the 1d wave equation modifing fun
# fun is an narray created with initialize_string_wave
# derivative is the derivative in t = 0
# velocity is the velocity of the wave
# steps is a tuple with (h_t, h_x)

def solve_1d_string(fun, derivate, velocity, steps):
    compute_1st_time_step(fun, derivate, velocity, steps)
    ht, hx = steps
    k1 = ((ht * velocity) / hx) ** 2
    k2 = 2 * (1 - k1)
    for tj in range(1, fun.shape[0] - 1):
        fun[tj + 1, 1:-1] = k1 * (fun[tj, :-2] + fun[tj, 2:]) + \
            k2 * fun[tj, 1:-1] - fun[tj - 1, 1:-1]


# compute_1st_time_step computes the second row (t = h_x)
def compute_1st_time_step(fun, derivate, velocity, steps):
    ht, hx = steps
    k1 = ((ht * velocity) / hx) ** 2 / 2.
    k2 = (1 - k1 * 2)
    fun[1, 1:-1] = k1 * (fun[0, :-2] + fun[0, 2:]) + \
        k2 * fun[0, 1:-1] + ht * derivate

# ----------------------------------------------------------------
# WATER DROP FUNCTIONS


# It initilizes the funtion (narray wth shape = dimensions) used for solving
# the wave equation. init_cond it can be an (N_x, N_y) array of an number
# (in case of number the solution is trivial)
def initialize_water_w(dimensions, init_cond):
    u = np.zeros(dimensions)
    u[0, :, :] = init_cond
    return u


# compute_1st_time_step_2d does the analog of compute_1st_time_step for 2D
def compute_1st_time_step_2d(fun, derivate, velocity, steps):
    ht, hx, hy = steps
    kx = (velocity * ht / hx) ** 2
    ky = (velocity * ht / hy) ** 2
    s = slice(1, -1)  # for writing less dense code I used slices. s = 1:-1
    lf = slice(None, -2)  # lf = :-2 stands for left
    rg = slice(2, None)  # rg = 2: stands for right
    fun[1, s, s] = kx * (fun[0, rg, s] - 2 * fun[0, s, s] + fun[0, lf, s]) + \
        ky * (fun[0, s, rg] - 2 * fun[0, s, s] + fun[0, s, lf]) + \
        fun[0, s, s] + ht * derivate


def solve_water(fun, derivate, velocity, steps):
    compute_1st_time_step_2d(fun, derivate, velocity, steps)
    ht, hx, hy = steps
    kx = (velocity * ht / hx) ** 2 / 2.
    ky = (velocity * ht / hy) ** 2 / 2.
    s = slice(1, -1)
    lf = slice(None, -2)
    rg = slice(2, None)
    for t in range(1, fun.shape[0] - 1):
        fun[t + 1, s, s] = kx * (fun[t, rg, s] - 2 * fun[t, s, s] + fun[t, lf, s]) + \
            ky * (fun[t, s, rg] - 2 * fun[t, s, s] + fun[t, s, lf]) + \
            2 * fun[t, s, s] - fun[t - 1, s, s]


# This function is used for creating the initial condition for our specific
# problem.
def create_initial_cond(dim):
    Nt, Nx, Ny = dim
    x = np.tile(np.linspace(0, 1, Nx).reshape((1, Nx)), (Ny, 1))
    y = np.tile(np.linspace(0, 1, Ny).reshape((Ny, 1)), (1, Nx))
    return np.exp(-500 * ((x - .5) ** 2 + (y - .5) ** 2))

# -------------------------------------------------------------------------------
# ISING MODEL FUNCTIONS


# It calculates E and M for a 2D narray rapresenting spins.
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):
            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


# It calculates the energy change if you flip an element in index = (i, j)
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


# It returns 2 array with same dimensions, one with the energy and one with
# at the same index the degeneration of a (N, N) system of spin cells
def get_degeration(N):
    degen = defaultdict(lambda: 0, {})
    for s in product([-1, 1], repeat=N * N):
        S = np.array(s).reshape((N, N))
        E, M = calculate_total_E_M(S)
        degen[E] += 1
    E = np.array(degen.keys())
    W_E = np.array(degen.values())
    return E, W_E


# It calculates the probability of a given energy state with a given T
def get_prob_energy(E, W_E, T):
    return (W_E * np.exp(- E / T)) / sum(W_E * np.exp(- E / T))


# It initilizes an NxN array with random or all aligned spins depending on the
# temperature
def initialize_spins(N, temperature):
    if temperature == 'high':
        return np.random.choice([1, -1], N * N).reshape((N, N))
    elif temperature == 'low':
        return np.ones((N, N)) * -1


# It runs the metropolis alorithm for getting the frequency with which
# some energy state appears
def ising_metropolis_degen(S, n0, ns, T):
    degen = defaultdict(lambda: 0, {})  # degen is a misleading variable name
    E, M = calculate_total_E_M(S)
    for i in range(ns):
        E, M = metro_step(S, T, E, M)
        if i > n0:
            degen[E] += 1
    E = np.array(degen.keys())
    p = np.array(degen.values(), dtype=float) / float(ns - n0)
    return E, p


# It runs through all spins and check which to spin using metropolis
def metro_step(S, T, E, M):
    Nx, Ny = S.shape
    for x in range(Nx):
        for y in range(Ny):
            dE = calculate_energy_change(S, (x, y))
            if exp(- dE / T) > random():
                S[x, y] = - S[x, y]
                E += dE
                M += 2 * S[x, y]
    return E, M


# It gives an array of energy and magnetization for each iteration
# with no previous termalization n0 = 0
def ising_macrostate_each_step(S, ns, T):
    E, M = calculate_total_E_M(S)
    Es = np.empty((ns,))
    Ms = np.empty((ns,))
    E_tot = M_tot = 0
    for i in range(ns):
        E, M = metro_step(S, T, E, M)
        E_tot += E
        M_tot += abs(M)
        Es[i] = E_tot / float(i + 1)
        Ms[i] = M_tot / float(i + 1)
    return Es, Ms


# It gives the avareges of M E and C chi over the ns - n0 iteration
def ising_get_avareges(S, n0, ns, T):
    sumE = sumE2 = sumM = sumM2 = 0
    E, M = calculate_total_E_M(S)
    for i in range(ns):
        E, M = metro_step(S, T, E, M)
        if i > n0:
            sumE += E
            sumE2 += E ** 2
            sumM += M
            sumM2 += M ** 2
    na = float(ns - n0)
    E = sumE / na
    E2 = sumE2 / na
    M = sumM / na
    M2 = sumM2 / na
    C = (E2 - E ** 2) / T ** 2
    chi = (M2 - M ** 2) / T
    return E, M, C, chi


# it plots the data produced in the c++ program 'ising_model.cpp'
def make_plot_ising(Ls, name):

    f, (ax1, ax2, ax3, ax4) = plt.subplots(4, 1, sharex=True, figsize=(6, 12))
    data = np.genfromtxt('Isi2DH0.dat')
    T = data[:, 0]
    E = data[:, 1]
    M = data[:, 2]
    C = data[:, 3]
    ax1.plot(T, E, color='black', label='Theoretical result')
    ax2.plot(T, C, color='black')
    ax3.plot(T, M, color='black')

    for i in Ls:

        data = np.genfromtxt("ising_" + str(i) + '.txt', delimiter=',')

        E = data[:, 0]
        M = np.abs(data[:, 1])
        C = data[:, 2]
        chi = data[:, 3]
        T = (np.arange(len(E)) + 1) * .2
        markersize = 3
        ax1.plot(T, E, 'o', label=str(i) + 'x' + str(i), ms=markersize)
        ax2.plot(T, C, 'o', ms=markersize)
        ax3.plot(T, M, 'o', ms=markersize)
        ax4.plot(T, chi, 'o', ms=markersize)

    ax1.grid(linestyle='dotted')
    ax2.grid(linestyle='dotted')
    ax3.grid(linestyle='dotted')
    ax4.grid(linestyle='dotted')
    ax1.set_ylabel('<E> / N [J]')
    ax2.set_ylabel('C / N [k_B]')
    ax3.set_ylabel('<|M|> / N')
    ax4.set_ylabel('Susceptibility / N')
    ax4.set_xlabel('Temperature [J / k_B]')
    ax1.legend()

    f.savefig(name)

# --------------------------------------------------
# --------------------------------------------------


def problem_2():
    f, ((ax1, ax2), (ax3, ax4), (ax5, ax6)) = plt.subplots(
        3, 2, sharex='col', sharey='row')
    axis = (ax1, ax2, ax3, ax4, ax5, ax6)
    ax5.set_xlabel('Length [m]')
    ax3.set_ylabel('Dislocation [m]')
    dim = (1200, 1200)

    velocity = (0.25) ** .5
    for n in range(1, 4):
        u_t0 = 2.5 * np.sin(n * np.pi * np.linspace(0, 1, dim[0]))
        cond = (u_t0, 0, 0)
        step = np.linspace(0, 1, dim[0])[1]
        wave = initialize_string_wave(dim, cond)
        solve_1d_string(wave, 0, velocity, (step, step))
        for ax, t in zip(axis, range(0, dim[0], dim[0] / 6)):
            ax.plot(np.linspace(0, 1, dim[0]), wave[t, :])
            ax.grid(True)
            str_time = '%.2f' % (step * t)
            ax.set_title('t = ' + str_time + 's')
    f.savefig('plot/vibrating_string.png')


def problem_5():
    dim = (301, 120, 120)
    velocity = .5
    init_cond = create_initial_cond(dim)
    fun = initialize_water_w(dim, init_cond)
    stept = np.linspace(0, 3, dim[0])[1]
    step = np.linspace(0, 1, dim[1])[1]
    solve_water(fun, 0, velocity, (stept, step, step))

    fig = plt.figure(figsize=(10, 15))
    X = np.linspace(0, 1, dim[1])
    Y = np.linspace(0, 1, dim[2])
    X, Y = np.meshgrid(X, Y)
    for i in range(1, 6 + 1):
        ax = fig.add_subplot(3, 2, i, projection='3d')
        ax.plot_surface(X, Y, fun[i * 40, :, :], rstride=1, cstride=1,
                        cmap=cm.coolwarm, linewidth=0, antialiased=False)
        ax.set_xlabel('x [m]')
        ax.set_ylabel('y [m]')
        str_time = '%.2f' % (stept * i * 40)
        ax.set_title('t = ' + str_time)
    fig.savefig('plot/3Dwave.png')


def problem_6_plot():
    E, W_E = get_degeration(3)
    plt.figure()
    plt.bar(E, W_E)
    plt.xlabel('E(J)')
    plt.ylabel('W(E)')
    plt.grid(linestyle='--')
    plt.savefig('plot/degen.png')

    for T in [1., 100.]:

        p = get_prob_energy(E, W_E, T)
        print 'Z(' + str(T) + ' J / k_b) =', sum(W_E * np.exp(- E / T))

        plt.figure()
        plt.bar(E + 0.4, p, log=False, label='Brute force')

        S = initialize_spins(3, 'low')
        E_sim, p = ising_metropolis_degen(S, 1000, 11000, T)
        plt.bar(E_sim - 0.4, p, log=False, label='Simulated')

        plt.legend()
        plt.xlabel('E(J)')
        plt.ylabel('p(E)')
        plt.grid(linestyle='--')
        plt.savefig('plot/ising_prob_T_' + str(int(T)) + '.png')


def problem_6():
    E, W_E = get_degeration(3)
    print tabulate([[y, x] for (y, x) in sorted(zip(E, W_E))], headers=['E', 'W(E)'])
    S = initialize_spins(3, 'low')

    E_sim, p = ising_metropolis_degen(S, 1000, 11000, 1.)

    W_sim = 10**4 * p
    W_sim = [int(x) for x in W_sim]
    plt.figure()
    plt.bar(E_sim, W_sim)
    plt.xlabel('E [J]')
    plt.ylabel('W(E)')
    plt.grid(linestyle='dotted')
    plt.savefig('plot/degen.png')

    print '\nMetropolis T = 1'
    print tabulate([[y, x] for (y, x) in sorted(zip(E_sim, p))], headers=['E', 'W(E)'])

    p = get_prob_energy(E, W_E, 1.)
    print '\nBoltzmann T = 1'
    print tabulate([[y, x] for (y, x) in sorted(zip(E, p))], headers=['E', 'W(E)'])

    E_sim, p = ising_metropolis_degen(S, 1000, 11000, 5.)
    print '\nMetropolis T = 5'
    print tabulate([[y, x] for (y, x) in sorted(zip(E_sim, p))], headers=['E', 'W(E)'])

    p = get_prob_energy(E, W_E, 5.)
    print '\nBoltzmann T = 5'
    print tabulate([[y, x] for (y, x) in sorted(zip(E, p))], headers=['E', 'W(E)'])


def problem_7():
    S = initialize_spins(6, 'low')
    E, M = ising_macrostate_each_step(S, 10000, 2)

    fig, (ax1, ax2) = plt.subplots(2, sharex=True)
    ax1.plot(np.arange(len(E)), E)
    ax1.set_ylabel('Energy [J]')
    ax1.grid()

    ax2.plot(np.arange(len(M)), M)
    ax2.set_ylabel('Magnetization')
    ax2.set_xlabel('Iteration')
    ax2.grid()
    fig.savefig('plot/no_random_spin.png')

    S = initialize_spins(6, 'high')
    E, M = ising_macrostate_each_step(S, 10000, 2)

    fig, (ax1, ax2) = plt.subplots(2, sharex=True)
    ax1.plot(np.arange(len(E)), E)
    ax1.set_ylabel('Energy [J]')
    ax1.grid()

    ax2.plot(np.arange(len(M)), M)
    ax2.set_ylabel('Magnetization')
    ax2.set_xlabel('Iteration')
    ax2.grid()
    fig.savefig('plot/random_spin.png')


def problem_8_9():
    make_plot_ising([6, 8, 16, 40], 'plot/ising_varing.png')


def problem_10():
    x = np.genfromtxt('spins.txt')
    plt.matshow(x[:40, :], cmap=plt.cm.gray)
    plt.title('T = 1.8 J / k_B')
    plt.savefig('plot/syst_18.png')
    plt.matshow(x[41:80, :], cmap=plt.cm.gray)
    plt.title('T = 2.2 J / k_B')
    plt.savefig('plot/syst_22.png')
    plt.matshow(x[81:, :], cmap=plt.cm.gray)
    plt.title('T = 5 J / k_B')
    plt.savefig('plot/syst_50.png')


if __name__ == '__main__':
    print 'Problem 2 running ...'
    problem_2()
    print 'Problem 5 running ...'
    problem_5()
    print 'Problem 6 running ...'
    problem_6()
    print 'Problem 7 running ...'
    problem_7()
    print 'Problem 8 and 9 running ...'
    problem_8_9()
    print 'Problem 10 running  ...'
    problem_10()
