import numpy as np
from sympy import *
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib.ticker import LinearLocator, FormatStrFormatter

def initOneSoliton(off, v, L, N, S):
    psi = np.zeros(N, dtype=np.complex_)
    t = 0
    h = L/N
    a = np.sqrt(2)
    B = 3/2 - 2 * S
    for i in range(N):
        x = i * h - off
        f = (2 * a * exp(a * x)) / (1 + B * exp(2 * a * x))
        psi[i] = f * exp(I * t + I * v * x)
    return psi

def initTwoSoliton(x1off, v1, x2off, v2, L, N, S):
    psi = np.zeros(N, dtype=np.complex_)
    t = 0
    h = L/N
    a = np.sqrt(2)
    B = 3/2 - 2 * S
    for i in range(N):
        x1 = i * h - x1off
        f1 = ( 2 * a * exp(a * x1)) / (1 + B * exp(2 * a * x1) )
        x2 = i * h - x2off
        f2 = ( 2 * a * exp(a * x2)) / (1 + B * exp(2 * a * x2) )
        psi[i] = ( f1 * exp(I * t + I * v1 * x1) + f2 * 
                    exp(I * t + I * v2 * x2) )
    return psi

def computeN(psi1, psi2, L, N):
    N1 = np.trapz(abs(psi1), dx=L/N)
    N2 = np.trapz(abs(psi2), dx=L/N)
    return np.abs(N1 - N2)

def plot2D(y, L, N):
    plt.plot(np.linspace(0, L, num=N), abs(y))
    plt.show()

def plot3D(psiEv, L, N, T, tau):
    fig = plt.figure(figsize=(20,10))
    ax = fig.gca(projection='3d')

    # Make data.
    X = np.arange(0, L, L/N)
    Y = np.arange(0, T, tau)
    X, Y = np.meshgrid(X, Y)

    # Plot the surface.
    surf = ax.plot_surface(X, Y, psiEv, cmap=cm.jet, linewidth=10,
       antialiased=True, rstride=1, cstride=1)

    # Customize the z axis.
    ax.set_zlim(0, 2)
    ax.zaxis.set_major_locator(LinearLocator(10))
    #ax.zaxis.set_major_formatter(FormatStrFormatter('%.1f
    #ax.view_init(30, 190)

    # Add a color bar which maps values to colors.
    fig.colorbar(surf, shrink=0.5, aspect=10)

    plt.show()
