from __future__ import division
import numpy as np
from numpy import pi
import matplotlib.pyplot as plt
import matplotlib
from matplotlib import ticker
import os,sys
plt.rc('text', usetex=True)
plt.rc('font', family='serif',serif='Times',size=50)

def initwavefunc(N1,v,delta_x,x0,x):
    p1=2/N1
    p2=2/(1-N1)
    n1=N1/2
    n2=(1-N1)/2
    u1=n1/(np.cosh((x-x0-delta_x)/p1))*np.exp(1j*v*x) 
    u2=n2/(np.cosh((x-x0-delta_x)/p2))*np.exp(1j*v*x)
    return u1,u2

def trappingpotentialfunc(e,s,wx,x):
    v0=0.5*wx**2
    V=e/(np.sqrt(2*pi)*s)*np.exp(-0.5*((x)/s)**2)+v0*x**2
    return V

def int_n(n,N,h):
    areaj=0
    for j in range(N-1):
        areaj=areaj+n[j]+n[j+1]
    return areaj*h/2

def splitstepmethod(g,v,wx,e,s,u10,u20,expon,M,x,L,k):
    N=2**expon
    if wx<1e-6:
        dt=1.20*L/(2*v*M)
    if v<1e-6:
        dt=1.20*pi/(wx*M)
    u1,u2=u10,u20
    V=trappingpotentialfunc(e,s,wx,x)
    Tmatrix=k**2/2
    n1,n2=np.zeros((N,M)),np.zeros((N,M))
    for m in range(M):
        u1,u2=np.exp(-0.5*dt*1j*(V-u1*np.conj(u1)-g*u2*np.conj(u2)))*u1,np.exp(-0.5*dt*1j*(V-u2*np.conj(u2)-g*u1*np.conj(u1)))*u2
        u1=np.fft.fftshift(np.fft.fft(u1,norm='ortho'))
        u1=np.exp(-1j*Tmatrix*dt)*u1
        u1=np.fft.ifft(np.fft.fftshift(u1),norm='ortho')
        u2=np.fft.fftshift(np.fft.fft(u2,norm='ortho'))
        u2=np.exp(-1j*Tmatrix*dt)*u2
        u2=np.fft.ifft(np.fft.fftshift(u2),norm='ortho')
        u1,u2=np.exp(-0.5*dt*1j*(V-u1*np.conj(u1)-g*u2*np.conj(u2)))*u1,np.exp(-0.5*dt*1j*(V-u2*np.conj(u2)-g*u1*np.conj(u1)))*u2
        n1[:,m],n2[:,m]=u1*np.conj(u1),u2*np.conj(u2)
    return n1,n2,dt

def imag_splitstepmethod(N1,g,wx,delta_x,absx0,expon,x,h,k):
    u1,u2=initwavefunc(N1,0,delta_x,-absx0,x)
    V=trappingpotentialfunc(0,1,wx,x)
    dtau=0.1
    Tmatrix=np.exp(-dtau*(k**2)/4)
    N=2**expon
    for m in range(M):
        u1=np.fft.ifft(np.fft.fftshift(Tmatrix*np.fft.fftshift(np.fft.fft(u1))))*np.sqrt(N1/int_n(u1*np.conj(u1),N,h))
        u2=np.fft.ifft(np.fft.fftshift(Tmatrix*np.fft.fftshift(np.fft.fft(u2))))*np.sqrt((1-N1)/int_n(u2*np.conj(u2),N,h))
        u1,u2=np.exp(-dtau*(V-(u1*np.conj(u1)+g*u2*np.conj(u2))))*u1*np.sqrt(N1/int_n(u1*np.conj(u1),N,h)),np.exp(-dtau*(V-(u2*np.conj(u2)+g*u1*np.conj(u1))))*u2*np.sqrt((1-N1)/int_n(u2*np.conj(u2),N,h))
        u1=np.fft.ifft(np.fft.fftshift(Tmatrix*np.fft.fftshift(np.fft.fft(u1))))*np.sqrt(N1/int_n(u1*np.conj(u1),N,h))
        u2=np.fft.ifft(np.fft.fftshift(Tmatrix*np.fft.fftshift(np.fft.fft(u2))))*np.sqrt((1-N1)/int_n(u2*np.conj(u2),N,h))
    return u1*np.conj(u1),u2*np.conj(u2)

def Tfunc(n_0,total,N,h):
    area_t=0
    for j in range(int(N/2)-1):
        area_t=area_t+n_0[int(j+N/2)]+n_0[int(j+1+N/2)]
    return h*area_t/(2*total)

binarypaperplots=True
if binarypaperplots:
    vs=np.array([0.155,0.16])
    for i in range(2):
        N1=0.3
        g=0.2
        s=0.4
        e=0.07
        wx=0
        absx0=30
        expon=10
        M=4000
        psi_pad=105
        L=4*absx0
        N=2**expon
        h=L/N
        hk=2*pi/L
        n=np.linspace(-N/2,N/2-1,N)
        x=n*h
        k=n*hk
        n10,n20=imag_splitstepmethod(N1,g,0,0,absx0,expon,x,h,k)
        n1,n2,dt=splitstepmethod(g,vs[i],wx,e,s,np.sqrt(n10)*np.exp(1j*vs[i]*x),np.sqrt(n20)*np.exp(1j*vs[i]*x),expon,M,x,L,k)
        ts=np.arange(M+1)*dt
        fsize=(20,12)
        fig=plt.figure(figsize=fsize)
        cmap=plt.get_cmap('inferno')
        plotting_number_N=4
        plotting_number_M=8
        ax1=fig.add_subplot(211)
        im1=ax1.pcolormesh(ts[::plotting_number_M],x[::plotting_number_N],n1[::plotting_number_N,::plotting_number_M],cmap=cmap,shading='flat')
        cbr=fig.colorbar(im1)
        ax1.set_xlabel("$t$")
        ax1.set_ylabel("$x$",labelpad=-20,rotation=0)
        cbr.ax.tick_params(labelsize=50)
        cbr.set_label('$\\vert \\psi_1 \\vert ^2$',fontsize=50,labelpad=psi_pad,rotation=0)
        tick_locator = ticker.MaxNLocator(nbins=5)
        cbr.locator = tick_locator
        cbr.update_ticks()
        ax2=fig.add_subplot(212)
        im2=ax2.pcolormesh(ts[::plotting_number_M],x[::plotting_number_N],n2[::plotting_number_N,::plotting_number_M],cmap=cmap,shading='flat')
        cbr2=fig.colorbar(im2)
        ax2.set_xlabel("$t$")
        ax2.set_ylabel("$x$",labelpad=-20,rotation=0)
        cbr2.ax.tick_params(labelsize=50)
        cbr2.set_label('$\\vert \\psi_2 \\vert ^2$',fontsize=50,labelpad=psi_pad,rotation=0)
        cbr2.locator = tick_locator
        cbr2.update_ticks()
        plt.tight_layout()
        if i==0:
            plt.savefig("traj1_eps.eps",format='eps')
        if i==1:
            plt.savefig("traj2_eps.eps",format='eps')
