import time
import numpy as np
import matplotlib.pyplot as plt
import scipy
import itertools
from scipy.integrate import odeint
#
# if transitions in the Phase Plot are rounded in form
# then expand range: decrease xmin, and increase xmax
# if LSODA complains that too much accuracy is requested
# or overflows are encountered and python won't continue
# then narrow range: increase xmin, and decrease xmax so that doesn't occur
# or vastly scale down initial value and slope boundary condition values
#
# suggested improvements:
# parallel mapping of bsearch
# arbitrary precision computations
# improved graphics and output
#
#define potential energy function here
def U(x):
    return(1/2*abs(x)**1)

#this is Schr\"odinger's equation as two coupled first order eqns    
def shrek(x, state, e):
    y, yprime = state
    dyprime = 2*(U(x)-e)*y
    dy = yprime
    return [dy, dyprime]

print("===> Close the Plot Window to Proceed with Calculation")
# this code plots the "Phase Plot" for a given energy
def plot_phaseplot(ee):
    y0 = [.01, .01]
    xgrid = np.arange(xmin, xmax, 0.01)
    result = odeint(shrek, y0, xgrid, (ee,), tfirst=True)
    phaseplot = np.tanh(result[:,0])
    plt.clf()
    plt.axhline( )
    plt.plot(xgrid,phaseplot)
    plt.title('energy='+str(ee),loc='center')
    plt.show()
    return()
    
# serial calculation will start after plot window closed
starttime=time.time( )

#this code 'memo-izes' previously computed number-of-bound-states 
#values so they aren't unnecessarily recomputed -> 3x speed-up
#variable 'savednbs' is a python dictionary of key-value pairs
#nbs is another name for omega, number of bound states below energy
savednbs = {}
def nbs(ee):
    if (ee) in savednbs:
        return([ee,savednbs[ee]])
    y0 = [.1, .1]
    xgrid = np.arange(xmin,xmax, 0.01)
    result = odeint(shrek, y0, xgrid, (ee,), tfirst=True)
    t = np.diff(np.sign(np.tanh(result[:,0])))
    transitions = t[t!=0]/2
    nboundstates=len(transitions)
    savednbs[ee]=nboundstates
    return([ee, nboundstates])

def nstates(ee):
	return(nbs (ee)[1])

def bsearchnbs(intrvl):
    [low,high]=intrvl
    mid =(low+high)/2
    slow = nstates(low)
    shigh = nstates(high)
    smid = nstates(mid)
    newintervals = [ ]
    if slow < smid: 
        newintervals += [[low, mid]]
    if smid < shigh: 
        newintervals += [[mid, high]]
    return(newintervals)

def refine(intervals):
# this maps out the whole eigenvalue spectrum in an initial range 
# using an evolutionary tree-like binary search
# intervals=[[0.,20.]] example starting interval nested list
# the mapping over the list of intervals would best be done 
# in parallel as in mathematica version
    print("starting interval=",intervals)
    nit=40; # max number of iterations; ~30-40 iterations is OK
    it=0; 
    while it <= nit:
        phaseplot=list(map(bsearchnbs,intervals));
        intervals = list(itertools. chain(*phaseplot)); #flattens list
        if it%10==0:
           print("===>iteration=",it,";  nintervals=",\
           len(intervals),";  interval midpoints=")
           evals=list(map(np.mean,intervals))
           ngooddigits=7 #typical for 64 bit arithmetic
           roundit=lambda numb: np.round(numb,ngooddigits)
           print(list(map(roundit,evals)))
           print(" ")
        it=it+1
    phaseplot=list(map(np.mean,intervals))
    evals=list(map(roundit,phaseplot))
    return(evals)
 #============================
 # after executing the above code, the following commands 
 # can be entered into Jupyter (for example) to actually get results
 # if overflow occurs, increase xmin, decrease xmax
(xmin,xmax) = (-40,40)  
startinginterval=[eminsearch,emaxsearch]=[0.0,15.0]
nbs(emaxsearch)  
plot_phaseplot(emaxsearch); 
intervals=[startinginterval]; #tees up starting interval for search
evals=refine(intervals); #does the search 
evals  #gives result
print(time.time( )-starttime," seconds calculation time (serial)")