#!/usr/bin/env python3

import numpy as np
import sys
# integration
from scipy.integrate import ode
from scipy.interpolate import interp1d

if sys.version_info[0] < 3:
    raise Exception("Must be using Python 3")

# my "package"
sys.path.append('/home/mart/Documents/KybI/2019/python/hw-burgers')
# my stuff
from utils.nonuniform_grid import nonuniform_grid
from hwbasics.HwBasics import Pn_nu, Pnx_nu, Hm_nu
from utils.reprint import reprint
from utils.utils import plot2D, plot_grid, plot3D

# cluster or not
from getpass import getuser

from dynnu.adaptive_grid import AdaptiveGrid, AdaptiveGridType
    
def solve_SG(J=3, c=1-1/1e4, x0=1/4, fineWidth=3/16, bHO=True, widthTol=1/25, borders=1, a=0.9, bc=[0, 2 * np.pi], bDebug=True):
    # calculate tf
    tf = (1 - 2 * x0)/c

    M2 = 2 * 2 ** J
    adaptiveGrid = AdaptiveGrid(J, AdaptiveGridType.DERIV_NU_PLUS_BORDERS, x0, borders, 
                    a=a, hw=fineWidth, onlyMax=True, bc=bc, Nper=int(M2/2)-borders)
    # adaptiveGrid = AdaptiveGrid(J, AdaptiveGridType.STATIONARY, x0, borders, gap=.11)
    # first, get uniform grid
    X, Xg = nonuniform_grid(J, 1)
    u0x = get_exact_x(X, 0, c, x0)
    # then 
    Xg, X = adaptiveGrid.get_grid(X, u0x)
    # print('MIN grid, coll:', np.min(np.diff(Xg)),np.min(np.diff(X)))
    # plot_grid(Xg, X)

    u0 = get_exact(X, 0, c, x0)
    v0 = get_exact_t(X, 0, c, x0)
    # ax = plot2D(X, u0, bShow=False)
    # plot2D(X, v0, ax=ax, legend=("u0", v0))

    # H, P
    H_and_P = get_H_and_P(J, Xg, X, bHO)
    Ps, Pbs = H_and_P[0], H_and_P[1]
    R2, R1, R0 = get_Rs(X, bHO)
    S2, S1, S0 = get_Ss(X, bHO) 
    Dx = get_Dxs(R2, R1, R0, X, Ps, Pbs)
    Sc2 = S2(X, Pbs)
    Sc1 = S1(X, Pbs)
    Sc0 = S0(X, Pbs)
    # in cluster or not
    cluster = getuser() == "mart.ratas"

    r = ode(fun).set_integrator('lsoda', nsteps=int(1e9))
    r.set_initial_value(np.hstack((u0.flatten(), v0.flatten())), 0).set_f_params(X, M2, Dx, Sc2, Sc1, Sc0)
    dt = 1e-3#tf/1e3 # TODO different? adaptive?
    xres = [np.hstack((0, X.flatten(), 1))]
    tres = [0]
    ures = [np.hstack((bc[0], u0.flatten(), bc[1]))]
    ue = [np.hstack((bc[0], u0.flatten(), bc[1]))]
    mdiff = [0]
    xmaxCur = x0
    mids = [x0,]
    rmids1 = [get_cur_mid(X, get_exact(X, 0, c, x0)),]
    exact_mid = lambda t: c * t + x0
    rmids2 = [exact_mid(0),]
    swaps = []
    while r.t < tf:
        r.integrate(r.t+dt)
        toPrint = "tf=%g" % r.t
        if cluster:
            print(toPrint)
        else:
            reprint(toPrint)
        # print('r.y', r.y)
        if np.any(r.y != r.y):
            print("'nan's in the result! CANNOT have this")
            print(r.y)
            break
        # find new (if needed) 
        xres.append(np.hstack((0, X.flatten(), 1)))
        tres.append(r.t)       # TODO - need to get the X axes!
        ucur = r.y[:M2]
        vcur = r.y[M2:]
        # if len(mids) < 10:
        #     vstart = vcur[0]
        #     vmax = np.max(np.abs(vcur))
        #     vstartRel = vstart/vmax
        #     toPrint = "%s : %g, %g"%(toPrint, vstart, vstartRel)
        #     plot2D(X, vcur, ax=plot2D(X, ucur, bShow=False), legend=("u", "v"), title=toPrint)
        ures.append(np.hstack((bc[0], ucur, bc[1]))) # reshape?
        ue.append(np.hstack((bc[0], get_exact(X, r.t, c, x0).flatten(), bc[1])))
        mdiff.append(np.max(np.abs(ue[-1]-ures[-1])))
        # plot2D(X, r.y, bShow=False)
        xMid = get_cur_mid(X, vcur) # TODO - finding middle is not trivial...
        realMid1 = X[0, np.argmax(get_exact_x(X, r.t, c, x0))] 
        realMid2 = exact_mid(r.t)
        mids.append(xMid)
        rmids1.append(realMid1)
        rmids2.append(realMid2)
        if abs(xMid - xmaxCur) > widthTol:
            # uxcur = r.y.reshape(1, M2) @ (np.linalg.lstsq(R2(X, Ps, Pbs), R1(X, Ps, Pbs))[0])
            # break; # TODO - REMOVE (just not going fowrad here)
            uxcur = np.diff(ucur)
            Xog, Xo = Xg, X
            Ps_old, Pbs_old = Ps, Pbs
            # Xg, X = adaptiveGrid.get_grid(Xo, uxcur)
            Xg, X = adaptiveGrid.get_grid((Xo[0,:-1] + Xo[0,1:])/2, uxcur)
            if np.any(Xog != Xg) or np.any(Xo != X): # if there is an actual change in the grid
                # print("DIFF:\n", Xog != Xg)
                H_and_P = get_H_and_P(J, Xg, X, bHO)
                Ps, Pbs = H_and_P[0], H_and_P[1]
                # plot_grid(Xg, X)
                try:
                    ucur = change_grid(J, ucur, Ps_old, Pbs_old, X, Xog, Xo, R2, bHO, bc=bc)
                except ValueError as e:
                    print(xMid)
                    print(Xo, Xog, "\n", X, Xg)
                    print(e)
                    break
                # if len(swaps) < 2:
                #     ax = plot2D(Xo, ucur, bShow=False)
                #     plot2D(X, ucur, ax=ax, bShow=False)
                #     Xop = np.hstack((0, Xo.flatten(), 1))
                #     plot2D(Xop, ue[-1], ax=ax, legend=("PREV", "NEW", "EXACT"))
                try:
                    vcur = change_grid(J, vcur, Ps_old, Pbs_old, X, Xog, Xo, R2, bHO, bc=[0, 0], bInterpolate=True) # NEED to interpolate
                except ValueError as e:
                    print(xMid)
                    print(Xo, Xog, "\n", X, Xg)
                    print(e)
                    break
                Dx = get_Dxs(R2, R1, R0, X, Ps, Pbs)
                Sc2 = S2(X, Pbs)
                Sc1 = S1(X, Pbs)
                Sc0 = S0(X, Pbs)
                # plot2D(Xo, r.y,bShow=False),plot2D(X, ucur) # TODO - remove
                r.set_initial_value(np.hstack((ucur.flatten(), vcur.flatten())), r.t)
                r.set_f_params(X, M2, Dx, Sc2, Sc1, Sc0)
                xmaxCur = xMid
                swaps.append(r.t)
    print("DONE:", toPrint)
    T = np.array(tres).reshape(len(tres), 1)
    xx, tt = np.meshgrid(np.hstack((0, X.flatten(), 1)), T)
    xx = np.array(xres)
    print ('TF=', tres[-1], "(target tf=%g)"%tf)
    U = np.array(ures)
    Ue = np.array(ue)
    print('SHAPES:', U.shape, Ue.shape)
    print('MAX diff:', np.max(mdiff))
    if bDebug:
        from matplotlib import pyplot as plt
        plt.plot(T.flatten(), mids)
        plt.plot(T.flatten(), rmids1)
        plt.plot(T.flatten(), rmids2)
        plt.plot(T.flatten(), mdiff)
        plt.plot(swaps, swaps, 'o')
        plt.legend(("CALC", "Real1", "REALREL", "ERROR", "swaps"))
        plt.ylim((0, 1))
    return xx, tt, U, Ue

# def get_cur_mid(X, ux, target=0.5):
#     # iclosest = np.argmin(np.abs(ux - target))
#     # x1 = X[0, iclosest]
#     # # print(x1, iclosest)
#     # val1 = ux[iclosest]
#     # if (val1 > target):
#     #     x2 = X[0, iclosest + 1]
#     #     val2 = ux[iclosest + 1]
#     # else:
#     #     x2 = X[0, iclosest - 1]
#     #     val2 = ux[iclosest - 1]
#     # diff1 = abs(val1 - target)
#     # diff2 = abs(val2 - target)
#     # return (diff1 * x2 + diff2 * x1)/(diff1 + diff2)
#     return X[0, np.argmin(ux)]

# def get_cur_mid(X, u):
    # Xc = (X[0,:-1] + X[0,1:])/2
    # ux = np.diff(u)
    # return Xc[np.argmin(ux)]
def get_cur_mid(X, v):
    return X[0, np.argmin(v)]


def get_H_and_P(J, Xg, X, bHO):
    H = Hm_nu(Xg)
    P2 = Pn_nu(J, 2, Xg, X.flatten())
    P2b = Pnx_nu(J, 2, 1, Xg) 
    P1 = Pn_nu(J, 1, Xg, X.flatten())
    if bHO:
        P4 = Pn_nu(J, 4, Xg, X.flatten())
        P4b = Pnx_nu(J, 4, 1, Xg)
        P3 = Pn_nu(J, 3, Xg, X.flatten())
        return (H, P1, P2, P3, P4), (None, None, P2b, None, P4b)
    else:
        return (H, P1, P2)        , (None, None, P2b)

def change_grid(J, ucur, Pso, Pbso, X, Xog, Xo, R2, bHO, bInterpolate=True, bc=None):
    if not bInterpolate:
        Px2 = np.hstack([Pnx_nu(J, 2, x, Xog) for x in X.flatten()])
        if bHO:
            Px4 = np.hstack([Pnx_nu(J, 4, x, Xog) for x in X.flatten()])
        PsMID = [None, None, Px2] if not bHO else [None, None, Px2, None, Px4]
        return ucur @ np.linalg.lstsq(R2(Xo, Pso, Pbso), R2(X, PsMID, Pbso))[0]
    else:
        filledUcur = np.hstack((bc[0], ucur.flatten(), bc[1]))
        filledX = np.hstack((0, Xo.flatten(), 1))
        interpolant = interp1d(filledX, filledUcur, 3)
        unew = interpolant(X.flatten())
        return unew.reshape(*ucur.shape)

def get_Rs(X, bHO):
    if not bHO:
        c1 = lambda Pbs: 2 * np.pi - Pbs[2]
        c2 = lambda Pbs: 0 * Pbs[2] + 1
        R2 = lambda X, Ps, Pbs: Ps[2] + c1(Pbs) @ X      + c2(Pbs)
        R1 = lambda X, Ps, Pbs: Ps[1] + c1(Pbs) 
        R0 = lambda X, Ps, Pbs: Ps[0]
        return R2, R1, R0
    else:
        c1 = lambda Pbs: - Pbs[2]
        c2 = lambda Pbs: 0 * Pbs[2]
        c3 = lambda Pbs: 1/6 * (- 6 * Pbs[4] + Pbs[2]) # 1/16 * (12 * np.pi) = 2 * np.pi -> S
        c4 = lambda Pbs: 0 * Pbs[2]
        R4 = lambda X, Ps, Pbs: Ps[4] + c1(Pbs) @ X**3/6  + c2(Pbs) @ X**2/2 + c3(Pbs) @ X      + c4(Pbs)
        R3 = lambda X, Ps, Pbs: Ps[3] + c1(Pbs) @ X**2/2  + c2(Pbs) @ X      + c3(Pbs)
        R2 = lambda X, Ps, Pbs: Ps[2] + c1(Pbs) @ X       + c2(Pbs)
        return R4, R3, R2


def get_Ss(X, bHO):
    if not bHO:
        c1 = lambda Pbs: 2 * np.pi
        c2 = lambda Pbs: 0
        S2 = lambda X, Pbs: c1(Pbs) * X + c2(Pbs)
        S1 = lambda X, Pbs: c1(Pbs) * X**0
        S0 = lambda X, Pbs: 0 * Pbs[2] * X**0
        return S2, S1, S0
    else:
        c1 = lambda Pbs: 0
        c2 = lambda Pbs: 0
        c3 = lambda Pbs: 2 * np.pi
        c4 = lambda Pbs: 0
        S4 = lambda X, Pbs: c1(Pbs) * X**3/6 + c2(Pbs) * X**2/2 + c3(Pbs) * X + c4(Pbs)
        S3 = lambda X, Pbs: c1(Pbs) * X**2/2 + c2(Pbs) * X      + c3(Pbs)
        S2 = lambda X, Pbs: c1(Pbs) * X      + c2(Pbs) 
        return S4, S3, S2


def get_Dxs(R2, R1, R0, X, Ps, Pbs):
    R2c = R2(X, Ps, Pbs)
    Dx = np.linalg.lstsq(R2c, R0(X, Ps, Pbs))[0]
    return Dx


def get_exact(X, T, c, x0=1/4):
    return 4 * np.arctan(np.exp((X - c * T - x0)/np.sqrt(1 - c**2)))


def get_exact_x(X, T, c, x0=1/4):
    # D(arctan(y),x) = y'/(1 + y^2)
    y = np.exp((X - c * T - x0)/np.sqrt(1 - c**2))
    yp = y/np.sqrt(1 - c**2)
    return 4 * yp/(1 + y**2)

def get_exact_t(X, T, c, x0=1/4):
    # D(arctan(y),t) = y'/(1 + y^2)
    y = np.exp((X - c * T - x0)/np.sqrt(1 - c**2))
    yp = -c * y/np.sqrt(1 - c**2)
    return 4 * yp/(1 + y**2)

def fun(t, uv, X, M2, Dx, S2, S1, S0):
    u = uv[:M2].reshape(1, M2)
    v = uv[M2:] # = dudt
    uxx = (u - S2) @ Dx + S0
    dvdt = uxx - np.sin(u)
    # if t < 4.4e-7:
    #     plot2D(X, v, ax=plot2D(X, u, bShow=False), legend=("u", "v"), title="t=%g"%t, xlims=[0, 0.4], ylims=[-.01,.01])
    return np.hstack((v, dvdt.flatten()))



if __name__ == '__main__':
    x0 = .3
    # widthTol = 1/35
    # fineWidth = .1
    c = 1 - 1e-4
    tf = (1 - 2 * x0)/c
    nrOfBorders = 1
    JRange = [4,]#5,6,7]
    for J in JRange:
        bests = []
        for fineWidth in np.arange(.2, .5, .01):
        # for fineWidth in [0.37]:
        #     for widthTol in [0.095]:
            for widthTol in np.arange(.04, .15, .005):
            # for widthTol in [.055]:#np.arange(.01, .1, .005):
                if J == 4:
                    aValues = np.arange(.881, .906, .001)
                    # aValues = np.arange(.892, .906, .001)
                    # aValues = [.896,]
                elif J == 5:
                    # aValues = np.arange(.963, .976, .001)
                    aValues = np.arange(.960, .980, .001)
                    # aValues = [.974,]
                elif J == 6:
                    aValues = np.arange(1, 1.0001, .001)
                elif J == 7:
                    aValues = [1.,]
                for a in aValues:
                    mStr = "J=%d, HOHWM, fineWidth = %g, widthTol=%g, a=%g"%(J, fineWidth, widthTol, a)
                    print(mStr)
                    # X, T, U, Ue = solve_kdv(J, alpha=alpha, beta=beta, c=c, tf=tf, bHO=True, x0=x0, fineWidth=fineWidth, widthTol=widthTol, borders=nrOfBorders, a=a)
                    try:
                        X, T, U, Ue = solve_SG(J, c=c, x0=x0, a=a, widthTol=widthTol,fineWidth=fineWidth, borders=nrOfBorders)
                    except ValueError as e:
                        print('Got exception (continuing on next)', e)
                    md = np.max(np.abs(U - Ue))
                    bests.append((fineWidth, widthTol, a, np.max(T), md))
                    #mdiffs.append(np.max(np.abs(U - Ue)))
                    # plot3D(X, T, U, bShow=False, title=mStr, zlims=[0, 2*np.pi]),plot3D(X,T,Ue, bShow=False),plot3D(X, T, U-Ue)
                # iMin = np.argmin(mdiffs)
                # print("best result: ", mdiffs[iMin], "at", aValues[iMin])
                    # bests.append((aValues[iMin], mdiffs[iMin]))
        print ("BEST:\n", bests)
        minDiff = 1e10
        minFW = "N/A"
        minWT = "N/A"
        minA = "N/A"
        for fw, wt, a, ctf, md in bests:
            if (ctf >= tf) and (md < minDiff):
                minDiff = md
                minFW = fw
                minWT = wt
                minA = a
        print ("Best at fw=", minFW, "wt=", minWT, "a=", minA, " with max diff of ", minDiff)
85c2a36e831c8444bb9a68feb8a1670b22fcf0d9
<mask token>


def solve_SG(J=3, c=1 - 1 / 10000.0, x0=1 / 4, fineWidth=3 / 16, bHO=True,
    widthTol=1 / 25, borders=1, a=0.9, bc=[0, 2 * np.pi], bDebug=True):
    tf = (1 - 2 * x0) / c
    M2 = 2 * 2 ** J
    adaptiveGrid = AdaptiveGrid(J, AdaptiveGridType.DERIV_NU_PLUS_BORDERS,
        x0, borders, a=a, hw=fineWidth, onlyMax=True, bc=bc, Nper=int(M2 / 
        2) - borders)
    X, Xg = nonuniform_grid(J, 1)
    u0x = get_exact_x(X, 0, c, x0)
    Xg, X = adaptiveGrid.get_grid(X, u0x)
    u0 = get_exact(X, 0, c, x0)
    v0 = get_exact_t(X, 0, c, x0)
    H_and_P = get_H_and_P(J, Xg, X, bHO)
    Ps, Pbs = H_and_P[0], H_and_P[1]
    R2, R1, R0 = get_Rs(X, bHO)
    S2, S1, S0 = get_Ss(X, bHO)
    Dx = get_Dxs(R2, R1, R0, X, Ps, Pbs)
    Sc2 = S2(X, Pbs)
    Sc1 = S1(X, Pbs)
    Sc0 = S0(X, Pbs)
    cluster = getuser() == 'mart.ratas'
    r = ode(fun).set_integrator('lsoda', nsteps=int(1000000000.0))
    r.set_initial_value(np.hstack((u0.flatten(), v0.flatten())), 0
        ).set_f_params(X, M2, Dx, Sc2, Sc1, Sc0)
    dt = 0.001
    xres = [np.hstack((0, X.flatten(), 1))]
    tres = [0]
    ures = [np.hstack((bc[0], u0.flatten(), bc[1]))]
    ue = [np.hstack((bc[0], u0.flatten(), bc[1]))]
    mdiff = [0]
    xmaxCur = x0
    mids = [x0]
    rmids1 = [get_cur_mid(X, get_exact(X, 0, c, x0))]
    exact_mid = lambda t: c * t + x0
    rmids2 = [exact_mid(0)]
    swaps = []
    while r.t < tf:
        r.integrate(r.t + dt)
        toPrint = 'tf=%g' % r.t
        if cluster:
            print(toPrint)
        else:
            reprint(toPrint)
        if np.any(r.y != r.y):
            print("'nan's in the result! CANNOT have this")
            print(r.y)
            break
        xres.append(np.hstack((0, X.flatten(), 1)))
        tres.append(r.t)
        ucur = r.y[:M2]
        vcur = r.y[M2:]
        ures.append(np.hstack((bc[0], ucur, bc[1])))
        ue.append(np.hstack((bc[0], get_exact(X, r.t, c, x0).flatten(), bc[1]))
            )
        mdiff.append(np.max(np.abs(ue[-1] - ures[-1])))
        xMid = get_cur_mid(X, vcur)
        realMid1 = X[0, np.argmax(get_exact_x(X, r.t, c, x0))]
        realMid2 = exact_mid(r.t)
        mids.append(xMid)
        rmids1.append(realMid1)
        rmids2.append(realMid2)
        if abs(xMid - xmaxCur) > widthTol:
            uxcur = np.diff(ucur)
            Xog, Xo = Xg, X
            Ps_old, Pbs_old = Ps, Pbs
            Xg, X = adaptiveGrid.get_grid((Xo[0, :-1] + Xo[0, 1:]) / 2, uxcur)
            if np.any(Xog != Xg) or np.any(Xo != X):
                H_and_P = get_H_and_P(J, Xg, X, bHO)
                Ps, Pbs = H_and_P[0], H_and_P[1]
                try:
                    ucur = change_grid(J, ucur, Ps_old, Pbs_old, X, Xog, Xo,
                        R2, bHO, bc=bc)
                except ValueError as e:
                    print(xMid)
                    print(Xo, Xog, '\n', X, Xg)
                    print(e)
                    break
                try:
                    vcur = change_grid(J, vcur, Ps_old, Pbs_old, X, Xog, Xo,
                        R2, bHO, bc=[0, 0], bInterpolate=True)
                except ValueError as e:
                    print(xMid)
                    print(Xo, Xog, '\n', X, Xg)
                    print(e)
                    break
                Dx = get_Dxs(R2, R1, R0, X, Ps, Pbs)
                Sc2 = S2(X, Pbs)
                Sc1 = S1(X, Pbs)
                Sc0 = S0(X, Pbs)
                r.set_initial_value(np.hstack((ucur.flatten(), vcur.flatten
                    ())), r.t)
                r.set_f_params(X, M2, Dx, Sc2, Sc1, Sc0)
                xmaxCur = xMid
                swaps.append(r.t)
    print('DONE:', toPrint)
    T = np.array(tres).reshape(len(tres), 1)
    xx, tt = np.meshgrid(np.hstack((0, X.flatten(), 1)), T)
    xx = np.array(xres)
    print('TF=', tres[-1], '(target tf=%g)' % tf)
    U = np.array(ures)
    Ue = np.array(ue)
    print('SHAPES:', U.shape, Ue.shape)
    print('MAX diff:', np.max(mdiff))
    if bDebug:
        from matplotlib import pyplot as plt
        plt.plot(T.flatten(), mids)
        plt.plot(T.flatten(), rmids1)
        plt.plot(T.flatten(), rmids2)
        plt.plot(T.flatten(), mdiff)
        plt.plot(swaps, swaps, 'o')
        plt.legend(('CALC', 'Real1', 'REALREL', 'ERROR', 'swaps'))
        plt.ylim((0, 1))
    return xx, tt, U, Ue


def get_cur_mid(X, v):
    return X[0, np.argmin(v)]


def get_H_and_P(J, Xg, X, bHO):
    H = Hm_nu(Xg)
    P2 = Pn_nu(J, 2, Xg, X.flatten())
    P2b = Pnx_nu(J, 2, 1, Xg)
    P1 = Pn_nu(J, 1, Xg, X.flatten())
    if bHO:
        P4 = Pn_nu(J, 4, Xg, X.flatten())
        P4b = Pnx_nu(J, 4, 1, Xg)
        P3 = Pn_nu(J, 3, Xg, X.flatten())
        return (H, P1, P2, P3, P4), (None, None, P2b, None, P4b)
    else:
        return (H, P1, P2), (None, None, P2b)


def change_grid(J, ucur, Pso, Pbso, X, Xog, Xo, R2, bHO, bInterpolate=True,
    bc=None):
    if not bInterpolate:
        Px2 = np.hstack([Pnx_nu(J, 2, x, Xog) for x in X.flatten()])
        if bHO:
            Px4 = np.hstack([Pnx_nu(J, 4, x, Xog) for x in X.flatten()])
        PsMID = [None, None, Px2] if not bHO else [None, None, Px2, None, Px4]
        return ucur @ np.linalg.lstsq(R2(Xo, Pso, Pbso), R2(X, PsMID, Pbso))[0]
    else:
        filledUcur = np.hstack((bc[0], ucur.flatten(), bc[1]))
        filledX = np.hstack((0, Xo.flatten(), 1))
        interpolant = interp1d(filledX, filledUcur, 3)
        unew = interpolant(X.flatten())
        return unew.reshape(*ucur.shape)


def get_Rs(X, bHO):
    if not bHO:
        c1 = lambda Pbs: 2 * np.pi - Pbs[2]
        c2 = lambda Pbs: 0 * Pbs[2] + 1
        R2 = lambda X, Ps, Pbs: Ps[2] + c1(Pbs) @ X + c2(Pbs)
        R1 = lambda X, Ps, Pbs: Ps[1] + c1(Pbs)
        R0 = lambda X, Ps, Pbs: Ps[0]
        return R2, R1, R0
    else:
        c1 = lambda Pbs: -Pbs[2]
        c2 = lambda Pbs: 0 * Pbs[2]
        c3 = lambda Pbs: 1 / 6 * (-6 * Pbs[4] + Pbs[2])
        c4 = lambda Pbs: 0 * Pbs[2]
        R4 = lambda X, Ps, Pbs: Ps[4] + c1(Pbs) @ X ** 3 / 6 + c2(Pbs
            ) @ X ** 2 / 2 + c3(Pbs) @ X + c4(Pbs)
        R3 = lambda X, Ps, Pbs: Ps[3] + c1(Pbs) @ X ** 2 / 2 + c2(Pbs
            ) @ X + c3(Pbs)
        R2 = lambda X, Ps, Pbs: Ps[2] + c1(Pbs) @ X + c2(Pbs)
        return R4, R3, R2


def get_Ss(X, bHO):
    if not bHO:
        c1 = lambda Pbs: 2 * np.pi
        c2 = lambda Pbs: 0
        S2 = lambda X, Pbs: c1(Pbs) * X + c2(Pbs)
        S1 = lambda X, Pbs: c1(Pbs) * X ** 0
        S0 = lambda X, Pbs: 0 * Pbs[2] * X ** 0
        return S2, S1, S0
    else:
        c1 = lambda Pbs: 0
        c2 = lambda Pbs: 0
        c3 = lambda Pbs: 2 * np.pi
        c4 = lambda Pbs: 0
        S4 = lambda X, Pbs: c1(Pbs) * X ** 3 / 6 + c2(Pbs) * X ** 2 / 2 + c3(
            Pbs) * X + c4(Pbs)
        S3 = lambda X, Pbs: c1(Pbs) * X ** 2 / 2 + c2(Pbs) * X + c3(Pbs)
        S2 = lambda X, Pbs: c1(Pbs) * X + c2(Pbs)
        return S4, S3, S2


def get_Dxs(R2, R1, R0, X, Ps, Pbs):
    R2c = R2(X, Ps, Pbs)
    Dx = np.linalg.lstsq(R2c, R0(X, Ps, Pbs))[0]
    return Dx


<mask token>


def get_exact_x(X, T, c, x0=1 / 4):
    y = np.exp((X - c * T - x0) / np.sqrt(1 - c ** 2))
    yp = y / np.sqrt(1 - c ** 2)
    return 4 * yp / (1 + y ** 2)


<mask token>


def fun(t, uv, X, M2, Dx, S2, S1, S0):
    u = uv[:M2].reshape(1, M2)
    v = uv[M2:]
    uxx = (u - S2) @ Dx + S0
    dvdt = uxx - np.sin(u)
    return np.hstack((v, dvdt.flatten()))


<mask token>

<mask token>


def solve_SG(J=3, c=1 - 1 / 10000.0, x0=1 / 4, fineWidth=3 / 16, bHO=True,
    widthTol=1 / 25, borders=1, a=0.9, bc=[0, 2 * np.pi], bDebug=True):
    tf = (1 - 2 * x0) / c
    M2 = 2 * 2 ** J
    adaptiveGrid = AdaptiveGrid(J, AdaptiveGridType.DERIV_NU_PLUS_BORDERS,
        x0, borders, a=a, hw=fineWidth, onlyMax=True, bc=bc, Nper=int(M2 / 
        2) - borders)
    X, Xg = nonuniform_grid(J, 1)
    u0x = get_exact_x(X, 0, c, x0)
    Xg, X = adaptiveGrid.get_grid(X, u0x)
    u0 = get_exact(X, 0, c, x0)
    v0 = get_exact_t(X, 0, c, x0)
    H_and_P = get_H_and_P(J, Xg, X, bHO)
    Ps, Pbs = H_and_P[0], H_and_P[1]
    R2, R1, R0 = get_Rs(X, bHO)
    S2, S1, S0 = get_Ss(X, bHO)
    Dx = get_Dxs(R2, R1, R0, X, Ps, Pbs)
    Sc2 = S2(X, Pbs)
    Sc1 = S1(X, Pbs)
    Sc0 = S0(X, Pbs)
    cluster = getuser() == 'mart.ratas'
    r = ode(fun).set_integrator('lsoda', nsteps=int(1000000000.0))
    r.set_initial_value(np.hstack((u0.flatten(), v0.flatten())), 0
        ).set_f_params(X, M2, Dx, Sc2, Sc1, Sc0)
    dt = 0.001
    xres = [np.hstack((0, X.flatten(), 1))]
    tres = [0]
    ures = [np.hstack((bc[0], u0.flatten(), bc[1]))]
    ue = [np.hstack((bc[0], u0.flatten(), bc[1]))]
    mdiff = [0]
    xmaxCur = x0
    mids = [x0]
    rmids1 = [get_cur_mid(X, get_exact(X, 0, c, x0))]
    exact_mid = lambda t: c * t + x0
    rmids2 = [exact_mid(0)]
    swaps = []
    while r.t < tf:
        r.integrate(r.t + dt)
        toPrint = 'tf=%g' % r.t
        if cluster:
            print(toPrint)
        else:
            reprint(toPrint)
        if np.any(r.y != r.y):
            print("'nan's in the result! CANNOT have this")
            print(r.y)
            break
        xres.append(np.hstack((0, X.flatten(), 1)))
        tres.append(r.t)
        ucur = r.y[:M2]
        vcur = r.y[M2:]
        ures.append(np.hstack((bc[0], ucur, bc[1])))
        ue.append(np.hstack((bc[0], get_exact(X, r.t, c, x0).flatten(), bc[1]))
            )
        mdiff.append(np.max(np.abs(ue[-1] - ures[-1])))
        xMid = get_cur_mid(X, vcur)
        realMid1 = X[0, np.argmax(get_exact_x(X, r.t, c, x0))]
        realMid2 = exact_mid(r.t)
        mids.append(xMid)
        rmids1.append(realMid1)
        rmids2.append(realMid2)
        if abs(xMid - xmaxCur) > widthTol:
            uxcur = np.diff(ucur)
            Xog, Xo = Xg, X
            Ps_old, Pbs_old = Ps, Pbs
            Xg, X = adaptiveGrid.get_grid((Xo[0, :-1] + Xo[0, 1:]) / 2, uxcur)
            if np.any(Xog != Xg) or np.any(Xo != X):
                H_and_P = get_H_and_P(J, Xg, X, bHO)
                Ps, Pbs = H_and_P[0], H_and_P[1]
                try:
                    ucur = change_grid(J, ucur, Ps_old, Pbs_old, X, Xog, Xo,
                        R2, bHO, bc=bc)
                except ValueError as e:
                    print(xMid)
                    print(Xo, Xog, '\n', X, Xg)
                    print(e)
                    break
                try:
                    vcur = change_grid(J, vcur, Ps_old, Pbs_old, X, Xog, Xo,
                        R2, bHO, bc=[0, 0], bInterpolate=True)
                except ValueError as e:
                    print(xMid)
                    print(Xo, Xog, '\n', X, Xg)
                    print(e)
                    break
                Dx = get_Dxs(R2, R1, R0, X, Ps, Pbs)
                Sc2 = S2(X, Pbs)
                Sc1 = S1(X, Pbs)
                Sc0 = S0(X, Pbs)
                r.set_initial_value(np.hstack((ucur.flatten(), vcur.flatten
                    ())), r.t)
                r.set_f_params(X, M2, Dx, Sc2, Sc1, Sc0)
                xmaxCur = xMid
                swaps.append(r.t)
    print('DONE:', toPrint)
    T = np.array(tres).reshape(len(tres), 1)
    xx, tt = np.meshgrid(np.hstack((0, X.flatten(), 1)), T)
    xx = np.array(xres)
    print('TF=', tres[-1], '(target tf=%g)' % tf)
    U = np.array(ures)
    Ue = np.array(ue)
    print('SHAPES:', U.shape, Ue.shape)
    print('MAX diff:', np.max(mdiff))
    if bDebug:
        from matplotlib import pyplot as plt
        plt.plot(T.flatten(), mids)
        plt.plot(T.flatten(), rmids1)
        plt.plot(T.flatten(), rmids2)
        plt.plot(T.flatten(), mdiff)
        plt.plot(swaps, swaps, 'o')
        plt.legend(('CALC', 'Real1', 'REALREL', 'ERROR', 'swaps'))
        plt.ylim((0, 1))
    return xx, tt, U, Ue


def get_cur_mid(X, v):
    return X[0, np.argmin(v)]


def get_H_and_P(J, Xg, X, bHO):
    H = Hm_nu(Xg)
    P2 = Pn_nu(J, 2, Xg, X.flatten())
    P2b = Pnx_nu(J, 2, 1, Xg)
    P1 = Pn_nu(J, 1, Xg, X.flatten())
    if bHO:
        P4 = Pn_nu(J, 4, Xg, X.flatten())
        P4b = Pnx_nu(J, 4, 1, Xg)
        P3 = Pn_nu(J, 3, Xg, X.flatten())
        return (H, P1, P2, P3, P4), (None, None, P2b, None, P4b)
    else:
        return (H, P1, P2), (None, None, P2b)


def change_grid(J, ucur, Pso, Pbso, X, Xog, Xo, R2, bHO, bInterpolate=True,
    bc=None):
    if not bInterpolate:
        Px2 = np.hstack([Pnx_nu(J, 2, x, Xog) for x in X.flatten()])
        if bHO:
            Px4 = np.hstack([Pnx_nu(J, 4, x, Xog) for x in X.flatten()])
        PsMID = [None, None, Px2] if not bHO else [None, None, Px2, None, Px4]
        return ucur @ np.linalg.lstsq(R2(Xo, Pso, Pbso), R2(X, PsMID, Pbso))[0]
    else:
        filledUcur = np.hstack((bc[0], ucur.flatten(), bc[1]))
        filledX = np.hstack((0, Xo.flatten(), 1))
        interpolant = interp1d(filledX, filledUcur, 3)
        unew = interpolant(X.flatten())
        return unew.reshape(*ucur.shape)


def get_Rs(X, bHO):
    if not bHO:
        c1 = lambda Pbs: 2 * np.pi - Pbs[2]
        c2 = lambda Pbs: 0 * Pbs[2] + 1
        R2 = lambda X, Ps, Pbs: Ps[2] + c1(Pbs) @ X + c2(Pbs)
        R1 = lambda X, Ps, Pbs: Ps[1] + c1(Pbs)
        R0 = lambda X, Ps, Pbs: Ps[0]
        return R2, R1, R0
    else:
        c1 = lambda Pbs: -Pbs[2]
        c2 = lambda Pbs: 0 * Pbs[2]
        c3 = lambda Pbs: 1 / 6 * (-6 * Pbs[4] + Pbs[2])
        c4 = lambda Pbs: 0 * Pbs[2]
        R4 = lambda X, Ps, Pbs: Ps[4] + c1(Pbs) @ X ** 3 / 6 + c2(Pbs
            ) @ X ** 2 / 2 + c3(Pbs) @ X + c4(Pbs)
        R3 = lambda X, Ps, Pbs: Ps[3] + c1(Pbs) @ X ** 2 / 2 + c2(Pbs
            ) @ X + c3(Pbs)
        R2 = lambda X, Ps, Pbs: Ps[2] + c1(Pbs) @ X + c2(Pbs)
        return R4, R3, R2


def get_Ss(X, bHO):
    if not bHO:
        c1 = lambda Pbs: 2 * np.pi
        c2 = lambda Pbs: 0
        S2 = lambda X, Pbs: c1(Pbs) * X + c2(Pbs)
        S1 = lambda X, Pbs: c1(Pbs) * X ** 0
        S0 = lambda X, Pbs: 0 * Pbs[2] * X ** 0
        return S2, S1, S0
    else:
        c1 = lambda Pbs: 0
        c2 = lambda Pbs: 0
        c3 = lambda Pbs: 2 * np.pi
        c4 = lambda Pbs: 0
        S4 = lambda X, Pbs: c1(Pbs) * X ** 3 / 6 + c2(Pbs) * X ** 2 / 2 + c3(
            Pbs) * X + c4(Pbs)
        S3 = lambda X, Pbs: c1(Pbs) * X ** 2 / 2 + c2(Pbs) * X + c3(Pbs)
        S2 = lambda X, Pbs: c1(Pbs) * X + c2(Pbs)
        return S4, S3, S2


def get_Dxs(R2, R1, R0, X, Ps, Pbs):
    R2c = R2(X, Ps, Pbs)
    Dx = np.linalg.lstsq(R2c, R0(X, Ps, Pbs))[0]
    return Dx


<mask token>


def get_exact_x(X, T, c, x0=1 / 4):
    y = np.exp((X - c * T - x0) / np.sqrt(1 - c ** 2))
    yp = y / np.sqrt(1 - c ** 2)
    return 4 * yp / (1 + y ** 2)


def get_exact_t(X, T, c, x0=1 / 4):
    y = np.exp((X - c * T - x0) / np.sqrt(1 - c ** 2))
    yp = -c * y / np.sqrt(1 - c ** 2)
    return 4 * yp / (1 + y ** 2)


def fun(t, uv, X, M2, Dx, S2, S1, S0):
    u = uv[:M2].reshape(1, M2)
    v = uv[M2:]
    uxx = (u - S2) @ Dx + S0
    dvdt = uxx - np.sin(u)
    return np.hstack((v, dvdt.flatten()))


<mask token>

<mask token>


def solve_SG(J=3, c=1 - 1 / 10000.0, x0=1 / 4, fineWidth=3 / 16, bHO=True,
    widthTol=1 / 25, borders=1, a=0.9, bc=[0, 2 * np.pi], bDebug=True):
    tf = (1 - 2 * x0) / c
    M2 = 2 * 2 ** J
    adaptiveGrid = AdaptiveGrid(J, AdaptiveGridType.DERIV_NU_PLUS_BORDERS,
        x0, borders, a=a, hw=fineWidth, onlyMax=True, bc=bc, Nper=int(M2 / 
        2) - borders)
    X, Xg = nonuniform_grid(J, 1)
    u0x = get_exact_x(X, 0, c, x0)
    Xg, X = adaptiveGrid.get_grid(X, u0x)
    u0 = get_exact(X, 0, c, x0)
    v0 = get_exact_t(X, 0, c, x0)
    H_and_P = get_H_and_P(J, Xg, X, bHO)
    Ps, Pbs = H_and_P[0], H_and_P[1]
    R2, R1, R0 = get_Rs(X, bHO)
    S2, S1, S0 = get_Ss(X, bHO)
    Dx = get_Dxs(R2, R1, R0, X, Ps, Pbs)
    Sc2 = S2(X, Pbs)
    Sc1 = S1(X, Pbs)
    Sc0 = S0(X, Pbs)
    cluster = getuser() == 'mart.ratas'
    r = ode(fun).set_integrator('lsoda', nsteps=int(1000000000.0))
    r.set_initial_value(np.hstack((u0.flatten(), v0.flatten())), 0
        ).set_f_params(X, M2, Dx, Sc2, Sc1, Sc0)
    dt = 0.001
    xres = [np.hstack((0, X.flatten(), 1))]
    tres = [0]
    ures = [np.hstack((bc[0], u0.flatten(), bc[1]))]
    ue = [np.hstack((bc[0], u0.flatten(), bc[1]))]
    mdiff = [0]
    xmaxCur = x0
    mids = [x0]
    rmids1 = [get_cur_mid(X, get_exact(X, 0, c, x0))]
    exact_mid = lambda t: c * t + x0
    rmids2 = [exact_mid(0)]
    swaps = []
    while r.t < tf:
        r.integrate(r.t + dt)
        toPrint = 'tf=%g' % r.t
        if cluster:
            print(toPrint)
        else:
            reprint(toPrint)
        if np.any(r.y != r.y):
            print("'nan's in the result! CANNOT have this")
            print(r.y)
            break
        xres.append(np.hstack((0, X.flatten(), 1)))
        tres.append(r.t)
        ucur = r.y[:M2]
        vcur = r.y[M2:]
        ures.append(np.hstack((bc[0], ucur, bc[1])))
        ue.append(np.hstack((bc[0], get_exact(X, r.t, c, x0).flatten(), bc[1]))
            )
        mdiff.append(np.max(np.abs(ue[-1] - ures[-1])))
        xMid = get_cur_mid(X, vcur)
        realMid1 = X[0, np.argmax(get_exact_x(X, r.t, c, x0))]
        realMid2 = exact_mid(r.t)
        mids.append(xMid)
        rmids1.append(realMid1)
        rmids2.append(realMid2)
        if abs(xMid - xmaxCur) > widthTol:
            uxcur = np.diff(ucur)
            Xog, Xo = Xg, X
            Ps_old, Pbs_old = Ps, Pbs
            Xg, X = adaptiveGrid.get_grid((Xo[0, :-1] + Xo[0, 1:]) / 2, uxcur)
            if np.any(Xog != Xg) or np.any(Xo != X):
                H_and_P = get_H_and_P(J, Xg, X, bHO)
                Ps, Pbs = H_and_P[0], H_and_P[1]
                try:
                    ucur = change_grid(J, ucur, Ps_old, Pbs_old, X, Xog, Xo,
                        R2, bHO, bc=bc)
                except ValueError as e:
                    print(xMid)
                    print(Xo, Xog, '\n', X, Xg)
                    print(e)
                    break
                try:
                    vcur = change_grid(J, vcur, Ps_old, Pbs_old, X, Xog, Xo,
                        R2, bHO, bc=[0, 0], bInterpolate=True)
                except ValueError as e:
                    print(xMid)
                    print(Xo, Xog, '\n', X, Xg)
                    print(e)
                    break
                Dx = get_Dxs(R2, R1, R0, X, Ps, Pbs)
                Sc2 = S2(X, Pbs)
                Sc1 = S1(X, Pbs)
                Sc0 = S0(X, Pbs)
                r.set_initial_value(np.hstack((ucur.flatten(), vcur.flatten
                    ())), r.t)
                r.set_f_params(X, M2, Dx, Sc2, Sc1, Sc0)
                xmaxCur = xMid
                swaps.append(r.t)
    print('DONE:', toPrint)
    T = np.array(tres).reshape(len(tres), 1)
    xx, tt = np.meshgrid(np.hstack((0, X.flatten(), 1)), T)
    xx = np.array(xres)
    print('TF=', tres[-1], '(target tf=%g)' % tf)
    U = np.array(ures)
    Ue = np.array(ue)
    print('SHAPES:', U.shape, Ue.shape)
    print('MAX diff:', np.max(mdiff))
    if bDebug:
        from matplotlib import pyplot as plt
        plt.plot(T.flatten(), mids)
        plt.plot(T.flatten(), rmids1)
        plt.plot(T.flatten(), rmids2)
        plt.plot(T.flatten(), mdiff)
        plt.plot(swaps, swaps, 'o')
        plt.legend(('CALC', 'Real1', 'REALREL', 'ERROR', 'swaps'))
        plt.ylim((0, 1))
    return xx, tt, U, Ue


def get_cur_mid(X, v):
    return X[0, np.argmin(v)]


def get_H_and_P(J, Xg, X, bHO):
    H = Hm_nu(Xg)
    P2 = Pn_nu(J, 2, Xg, X.flatten())
    P2b = Pnx_nu(J, 2, 1, Xg)
    P1 = Pn_nu(J, 1, Xg, X.flatten())
    if bHO:
        P4 = Pn_nu(J, 4, Xg, X.flatten())
        P4b = Pnx_nu(J, 4, 1, Xg)
        P3 = Pn_nu(J, 3, Xg, X.flatten())
        return (H, P1, P2, P3, P4), (None, None, P2b, None, P4b)
    else:
        return (H, P1, P2), (None, None, P2b)


def change_grid(J, ucur, Pso, Pbso, X, Xog, Xo, R2, bHO, bInterpolate=True,
    bc=None):
    if not bInterpolate:
        Px2 = np.hstack([Pnx_nu(J, 2, x, Xog) for x in X.flatten()])
        if bHO:
            Px4 = np.hstack([Pnx_nu(J, 4, x, Xog) for x in X.flatten()])
        PsMID = [None, None, Px2] if not bHO else [None, None, Px2, None, Px4]
        return ucur @ np.linalg.lstsq(R2(Xo, Pso, Pbso), R2(X, PsMID, Pbso))[0]
    else:
        filledUcur = np.hstack((bc[0], ucur.flatten(), bc[1]))
        filledX = np.hstack((0, Xo.flatten(), 1))
        interpolant = interp1d(filledX, filledUcur, 3)
        unew = interpolant(X.flatten())
        return unew.reshape(*ucur.shape)


def get_Rs(X, bHO):
    if not bHO:
        c1 = lambda Pbs: 2 * np.pi - Pbs[2]
        c2 = lambda Pbs: 0 * Pbs[2] + 1
        R2 = lambda X, Ps, Pbs: Ps[2] + c1(Pbs) @ X + c2(Pbs)
        R1 = lambda X, Ps, Pbs: Ps[1] + c1(Pbs)
        R0 = lambda X, Ps, Pbs: Ps[0]
        return R2, R1, R0
    else:
        c1 = lambda Pbs: -Pbs[2]
        c2 = lambda Pbs: 0 * Pbs[2]
        c3 = lambda Pbs: 1 / 6 * (-6 * Pbs[4] + Pbs[2])
        c4 = lambda Pbs: 0 * Pbs[2]
        R4 = lambda X, Ps, Pbs: Ps[4] + c1(Pbs) @ X ** 3 / 6 + c2(Pbs
            ) @ X ** 2 / 2 + c3(Pbs) @ X + c4(Pbs)
        R3 = lambda X, Ps, Pbs: Ps[3] + c1(Pbs) @ X ** 2 / 2 + c2(Pbs
            ) @ X + c3(Pbs)
        R2 = lambda X, Ps, Pbs: Ps[2] + c1(Pbs) @ X + c2(Pbs)
        return R4, R3, R2


def get_Ss(X, bHO):
    if not bHO:
        c1 = lambda Pbs: 2 * np.pi
        c2 = lambda Pbs: 0
        S2 = lambda X, Pbs: c1(Pbs) * X + c2(Pbs)
        S1 = lambda X, Pbs: c1(Pbs) * X ** 0
        S0 = lambda X, Pbs: 0 * Pbs[2] * X ** 0
        return S2, S1, S0
    else:
        c1 = lambda Pbs: 0
        c2 = lambda Pbs: 0
        c3 = lambda Pbs: 2 * np.pi
        c4 = lambda Pbs: 0
        S4 = lambda X, Pbs: c1(Pbs) * X ** 3 / 6 + c2(Pbs) * X ** 2 / 2 + c3(
            Pbs) * X + c4(Pbs)
        S3 = lambda X, Pbs: c1(Pbs) * X ** 2 / 2 + c2(Pbs) * X + c3(Pbs)
        S2 = lambda X, Pbs: c1(Pbs) * X + c2(Pbs)
        return S4, S3, S2


def get_Dxs(R2, R1, R0, X, Ps, Pbs):
    R2c = R2(X, Ps, Pbs)
    Dx = np.linalg.lstsq(R2c, R0(X, Ps, Pbs))[0]
    return Dx


def get_exact(X, T, c, x0=1 / 4):
    return 4 * np.arctan(np.exp((X - c * T - x0) / np.sqrt(1 - c ** 2)))


def get_exact_x(X, T, c, x0=1 / 4):
    y = np.exp((X - c * T - x0) / np.sqrt(1 - c ** 2))
    yp = y / np.sqrt(1 - c ** 2)
    return 4 * yp / (1 + y ** 2)


def get_exact_t(X, T, c, x0=1 / 4):
    y = np.exp((X - c * T - x0) / np.sqrt(1 - c ** 2))
    yp = -c * y / np.sqrt(1 - c ** 2)
    return 4 * yp / (1 + y ** 2)


def fun(t, uv, X, M2, Dx, S2, S1, S0):
    u = uv[:M2].reshape(1, M2)
    v = uv[M2:]
    uxx = (u - S2) @ Dx + S0
    dvdt = uxx - np.sin(u)
    return np.hstack((v, dvdt.flatten()))


<mask token>

import numpy as np
import sys
from scipy.integrate import ode
from scipy.interpolate import interp1d
if sys.version_info[0] < 3:
    raise Exception('Must be using Python 3')
sys.path.append('/home/mart/Documents/KybI/2019/python/hw-burgers')
from utils.nonuniform_grid import nonuniform_grid
from hwbasics.HwBasics import Pn_nu, Pnx_nu, Hm_nu
from utils.reprint import reprint
from utils.utils import plot2D, plot_grid, plot3D
from getpass import getuser
from dynnu.adaptive_grid import AdaptiveGrid, AdaptiveGridType


def solve_SG(J=3, c=1 - 1 / 10000.0, x0=1 / 4, fineWidth=3 / 16, bHO=True,
    widthTol=1 / 25, borders=1, a=0.9, bc=[0, 2 * np.pi], bDebug=True):
    tf = (1 - 2 * x0) / c
    M2 = 2 * 2 ** J
    adaptiveGrid = AdaptiveGrid(J, AdaptiveGridType.DERIV_NU_PLUS_BORDERS,
        x0, borders, a=a, hw=fineWidth, onlyMax=True, bc=bc, Nper=int(M2 / 
        2) - borders)
    X, Xg = nonuniform_grid(J, 1)
    u0x = get_exact_x(X, 0, c, x0)
    Xg, X = adaptiveGrid.get_grid(X, u0x)
    u0 = get_exact(X, 0, c, x0)
    v0 = get_exact_t(X, 0, c, x0)
    H_and_P = get_H_and_P(J, Xg, X, bHO)
    Ps, Pbs = H_and_P[0], H_and_P[1]
    R2, R1, R0 = get_Rs(X, bHO)
    S2, S1, S0 = get_Ss(X, bHO)
    Dx = get_Dxs(R2, R1, R0, X, Ps, Pbs)
    Sc2 = S2(X, Pbs)
    Sc1 = S1(X, Pbs)
    Sc0 = S0(X, Pbs)
    cluster = getuser() == 'mart.ratas'
    r = ode(fun).set_integrator('lsoda', nsteps=int(1000000000.0))
    r.set_initial_value(np.hstack((u0.flatten(), v0.flatten())), 0
        ).set_f_params(X, M2, Dx, Sc2, Sc1, Sc0)
    dt = 0.001
    xres = [np.hstack((0, X.flatten(), 1))]
    tres = [0]
    ures = [np.hstack((bc[0], u0.flatten(), bc[1]))]
    ue = [np.hstack((bc[0], u0.flatten(), bc[1]))]
    mdiff = [0]
    xmaxCur = x0
    mids = [x0]
    rmids1 = [get_cur_mid(X, get_exact(X, 0, c, x0))]
    exact_mid = lambda t: c * t + x0
    rmids2 = [exact_mid(0)]
    swaps = []
    while r.t < tf:
        r.integrate(r.t + dt)
        toPrint = 'tf=%g' % r.t
        if cluster:
            print(toPrint)
        else:
            reprint(toPrint)
        if np.any(r.y != r.y):
            print("'nan's in the result! CANNOT have this")
            print(r.y)
            break
        xres.append(np.hstack((0, X.flatten(), 1)))
        tres.append(r.t)
        ucur = r.y[:M2]
        vcur = r.y[M2:]
        ures.append(np.hstack((bc[0], ucur, bc[1])))
        ue.append(np.hstack((bc[0], get_exact(X, r.t, c, x0).flatten(), bc[1]))
            )
        mdiff.append(np.max(np.abs(ue[-1] - ures[-1])))
        xMid = get_cur_mid(X, vcur)
        realMid1 = X[0, np.argmax(get_exact_x(X, r.t, c, x0))]
        realMid2 = exact_mid(r.t)
        mids.append(xMid)
        rmids1.append(realMid1)
        rmids2.append(realMid2)
        if abs(xMid - xmaxCur) > widthTol:
            uxcur = np.diff(ucur)
            Xog, Xo = Xg, X
            Ps_old, Pbs_old = Ps, Pbs
            Xg, X = adaptiveGrid.get_grid((Xo[0, :-1] + Xo[0, 1:]) / 2, uxcur)
            if np.any(Xog != Xg) or np.any(Xo != X):
                H_and_P = get_H_and_P(J, Xg, X, bHO)
                Ps, Pbs = H_and_P[0], H_and_P[1]
                try:
                    ucur = change_grid(J, ucur, Ps_old, Pbs_old, X, Xog, Xo,
                        R2, bHO, bc=bc)
                except ValueError as e:
                    print(xMid)
                    print(Xo, Xog, '\n', X, Xg)
                    print(e)
                    break
                try:
                    vcur = change_grid(J, vcur, Ps_old, Pbs_old, X, Xog, Xo,
                        R2, bHO, bc=[0, 0], bInterpolate=True)
                except ValueError as e:
                    print(xMid)
                    print(Xo, Xog, '\n', X, Xg)
                    print(e)
                    break
                Dx = get_Dxs(R2, R1, R0, X, Ps, Pbs)
                Sc2 = S2(X, Pbs)
                Sc1 = S1(X, Pbs)
                Sc0 = S0(X, Pbs)
                r.set_initial_value(np.hstack((ucur.flatten(), vcur.flatten
                    ())), r.t)
                r.set_f_params(X, M2, Dx, Sc2, Sc1, Sc0)
                xmaxCur = xMid
                swaps.append(r.t)
    print('DONE:', toPrint)
    T = np.array(tres).reshape(len(tres), 1)
    xx, tt = np.meshgrid(np.hstack((0, X.flatten(), 1)), T)
    xx = np.array(xres)
    print('TF=', tres[-1], '(target tf=%g)' % tf)
    U = np.array(ures)
    Ue = np.array(ue)
    print('SHAPES:', U.shape, Ue.shape)
    print('MAX diff:', np.max(mdiff))
    if bDebug:
        from matplotlib import pyplot as plt
        plt.plot(T.flatten(), mids)
        plt.plot(T.flatten(), rmids1)
        plt.plot(T.flatten(), rmids2)
        plt.plot(T.flatten(), mdiff)
        plt.plot(swaps, swaps, 'o')
        plt.legend(('CALC', 'Real1', 'REALREL', 'ERROR', 'swaps'))
        plt.ylim((0, 1))
    return xx, tt, U, Ue


def get_cur_mid(X, v):
    return X[0, np.argmin(v)]


def get_H_and_P(J, Xg, X, bHO):
    H = Hm_nu(Xg)
    P2 = Pn_nu(J, 2, Xg, X.flatten())
    P2b = Pnx_nu(J, 2, 1, Xg)
    P1 = Pn_nu(J, 1, Xg, X.flatten())
    if bHO:
        P4 = Pn_nu(J, 4, Xg, X.flatten())
        P4b = Pnx_nu(J, 4, 1, Xg)
        P3 = Pn_nu(J, 3, Xg, X.flatten())
        return (H, P1, P2, P3, P4), (None, None, P2b, None, P4b)
    else:
        return (H, P1, P2), (None, None, P2b)


def change_grid(J, ucur, Pso, Pbso, X, Xog, Xo, R2, bHO, bInterpolate=True,
    bc=None):
    if not bInterpolate:
        Px2 = np.hstack([Pnx_nu(J, 2, x, Xog) for x in X.flatten()])
        if bHO:
            Px4 = np.hstack([Pnx_nu(J, 4, x, Xog) for x in X.flatten()])
        PsMID = [None, None, Px2] if not bHO else [None, None, Px2, None, Px4]
        return ucur @ np.linalg.lstsq(R2(Xo, Pso, Pbso), R2(X, PsMID, Pbso))[0]
    else:
        filledUcur = np.hstack((bc[0], ucur.flatten(), bc[1]))
        filledX = np.hstack((0, Xo.flatten(), 1))
        interpolant = interp1d(filledX, filledUcur, 3)
        unew = interpolant(X.flatten())
        return unew.reshape(*ucur.shape)


def get_Rs(X, bHO):
    if not bHO:
        c1 = lambda Pbs: 2 * np.pi - Pbs[2]
        c2 = lambda Pbs: 0 * Pbs[2] + 1
        R2 = lambda X, Ps, Pbs: Ps[2] + c1(Pbs) @ X + c2(Pbs)
        R1 = lambda X, Ps, Pbs: Ps[1] + c1(Pbs)
        R0 = lambda X, Ps, Pbs: Ps[0]
        return R2, R1, R0
    else:
        c1 = lambda Pbs: -Pbs[2]
        c2 = lambda Pbs: 0 * Pbs[2]
        c3 = lambda Pbs: 1 / 6 * (-6 * Pbs[4] + Pbs[2])
        c4 = lambda Pbs: 0 * Pbs[2]
        R4 = lambda X, Ps, Pbs: Ps[4] + c1(Pbs) @ X ** 3 / 6 + c2(Pbs
            ) @ X ** 2 / 2 + c3(Pbs) @ X + c4(Pbs)
        R3 = lambda X, Ps, Pbs: Ps[3] + c1(Pbs) @ X ** 2 / 2 + c2(Pbs
            ) @ X + c3(Pbs)
        R2 = lambda X, Ps, Pbs: Ps[2] + c1(Pbs) @ X + c2(Pbs)
        return R4, R3, R2


def get_Ss(X, bHO):
    if not bHO:
        c1 = lambda Pbs: 2 * np.pi
        c2 = lambda Pbs: 0
        S2 = lambda X, Pbs: c1(Pbs) * X + c2(Pbs)
        S1 = lambda X, Pbs: c1(Pbs) * X ** 0
        S0 = lambda X, Pbs: 0 * Pbs[2] * X ** 0
        return S2, S1, S0
    else:
        c1 = lambda Pbs: 0
        c2 = lambda Pbs: 0
        c3 = lambda Pbs: 2 * np.pi
        c4 = lambda Pbs: 0
        S4 = lambda X, Pbs: c1(Pbs) * X ** 3 / 6 + c2(Pbs) * X ** 2 / 2 + c3(
            Pbs) * X + c4(Pbs)
        S3 = lambda X, Pbs: c1(Pbs) * X ** 2 / 2 + c2(Pbs) * X + c3(Pbs)
        S2 = lambda X, Pbs: c1(Pbs) * X + c2(Pbs)
        return S4, S3, S2


def get_Dxs(R2, R1, R0, X, Ps, Pbs):
    R2c = R2(X, Ps, Pbs)
    Dx = np.linalg.lstsq(R2c, R0(X, Ps, Pbs))[0]
    return Dx


def get_exact(X, T, c, x0=1 / 4):
    return 4 * np.arctan(np.exp((X - c * T - x0) / np.sqrt(1 - c ** 2)))


def get_exact_x(X, T, c, x0=1 / 4):
    y = np.exp((X - c * T - x0) / np.sqrt(1 - c ** 2))
    yp = y / np.sqrt(1 - c ** 2)
    return 4 * yp / (1 + y ** 2)


def get_exact_t(X, T, c, x0=1 / 4):
    y = np.exp((X - c * T - x0) / np.sqrt(1 - c ** 2))
    yp = -c * y / np.sqrt(1 - c ** 2)
    return 4 * yp / (1 + y ** 2)


def fun(t, uv, X, M2, Dx, S2, S1, S0):
    u = uv[:M2].reshape(1, M2)
    v = uv[M2:]
    uxx = (u - S2) @ Dx + S0
    dvdt = uxx - np.sin(u)
    return np.hstack((v, dvdt.flatten()))


if __name__ == '__main__':
    x0 = 0.3
    c = 1 - 0.0001
    tf = (1 - 2 * x0) / c
    nrOfBorders = 1
    JRange = [4]
    for J in JRange:
        bests = []
        for fineWidth in np.arange(0.2, 0.5, 0.01):
            for widthTol in np.arange(0.04, 0.15, 0.005):
                if J == 4:
                    aValues = np.arange(0.881, 0.906, 0.001)
                elif J == 5:
                    aValues = np.arange(0.96, 0.98, 0.001)
                elif J == 6:
                    aValues = np.arange(1, 1.0001, 0.001)
                elif J == 7:
                    aValues = [1.0]
                for a in aValues:
                    mStr = 'J=%d, HOHWM, fineWidth = %g, widthTol=%g, a=%g' % (
                        J, fineWidth, widthTol, a)
                    print(mStr)
                    try:
                        X, T, U, Ue = solve_SG(J, c=c, x0=x0, a=a, widthTol
                            =widthTol, fineWidth=fineWidth, borders=nrOfBorders
                            )
                    except ValueError as e:
                        print('Got exception (continuing on next)', e)
                    md = np.max(np.abs(U - Ue))
                    bests.append((fineWidth, widthTol, a, np.max(T), md))
        print('BEST:\n', bests)
        minDiff = 10000000000.0
        minFW = 'N/A'
        minWT = 'N/A'
        minA = 'N/A'
        for fw, wt, a, ctf, md in bests:
            if ctf >= tf and md < minDiff:
                minDiff = md
                minFW = fw
                minWT = wt
                minA = a
        print('Best at fw=', minFW, 'wt=', minWT, 'a=', minA,
            ' with max diff of ', minDiff)

#!/usr/bin/env python3

import numpy as np
import sys
# integration
from scipy.integrate import ode
from scipy.interpolate import interp1d

if sys.version_info[0] < 3:
    raise Exception("Must be using Python 3")

# my "package"
sys.path.append('/home/mart/Documents/KybI/2019/python/hw-burgers')
# my stuff
from utils.nonuniform_grid import nonuniform_grid
from hwbasics.HwBasics import Pn_nu, Pnx_nu, Hm_nu
from utils.reprint import reprint
from utils.utils import plot2D, plot_grid, plot3D

# cluster or not
from getpass import getuser

from dynnu.adaptive_grid import AdaptiveGrid, AdaptiveGridType
    
def solve_SG(J=3, c=1-1/1e4, x0=1/4, fineWidth=3/16, bHO=True, widthTol=1/25, borders=1, a=0.9, bc=[0, 2 * np.pi], bDebug=True):
    # calculate tf
    tf = (1 - 2 * x0)/c

    M2 = 2 * 2 ** J
    adaptiveGrid = AdaptiveGrid(J, AdaptiveGridType.DERIV_NU_PLUS_BORDERS, x0, borders, 
                    a=a, hw=fineWidth, onlyMax=True, bc=bc, Nper=int(M2/2)-borders)
    # adaptiveGrid = AdaptiveGrid(J, AdaptiveGridType.STATIONARY, x0, borders, gap=.11)
    # first, get uniform grid
    X, Xg = nonuniform_grid(J, 1)
    u0x = get_exact_x(X, 0, c, x0)
    # then 
    Xg, X = adaptiveGrid.get_grid(X, u0x)
    # print('MIN grid, coll:', np.min(np.diff(Xg)),np.min(np.diff(X)))
    # plot_grid(Xg, X)

    u0 = get_exact(X, 0, c, x0)
    v0 = get_exact_t(X, 0, c, x0)
    # ax = plot2D(X, u0, bShow=False)
    # plot2D(X, v0, ax=ax, legend=("u0", v0))

    # H, P
    H_and_P = get_H_and_P(J, Xg, X, bHO)
    Ps, Pbs = H_and_P[0], H_and_P[1]
    R2, R1, R0 = get_Rs(X, bHO)
    S2, S1, S0 = get_Ss(X, bHO) 
    Dx = get_Dxs(R2, R1, R0, X, Ps, Pbs)
    Sc2 = S2(X, Pbs)
    Sc1 = S1(X, Pbs)
    Sc0 = S0(X, Pbs)
    # in cluster or not
    cluster = getuser() == "mart.ratas"

    r = ode(fun).set_integrator('lsoda', nsteps=int(1e9))
    r.set_initial_value(np.hstack((u0.flatten(), v0.flatten())), 0).set_f_params(X, M2, Dx, Sc2, Sc1, Sc0)
    dt = 1e-3#tf/1e3 # TODO different? adaptive?
    xres = [np.hstack((0, X.flatten(), 1))]
    tres = [0]
    ures = [np.hstack((bc[0], u0.flatten(), bc[1]))]
    ue = [np.hstack((bc[0], u0.flatten(), bc[1]))]
    mdiff = [0]
    xmaxCur = x0
    mids = [x0,]
    rmids1 = [get_cur_mid(X, get_exact(X, 0, c, x0)),]
    exact_mid = lambda t: c * t + x0
    rmids2 = [exact_mid(0),]
    swaps = []
    while r.t < tf:
        r.integrate(r.t+dt)
        toPrint = "tf=%g" % r.t
        if cluster:
            print(toPrint)
        else:
            reprint(toPrint)
        # print('r.y', r.y)
        if np.any(r.y != r.y):
            print("'nan's in the result! CANNOT have this")
            print(r.y)
            break
        # find new (if needed) 
        xres.append(np.hstack((0, X.flatten(), 1)))
        tres.append(r.t)       # TODO - need to get the X axes!
        ucur = r.y[:M2]
        vcur = r.y[M2:]
        # if len(mids) < 10:
        #     vstart = vcur[0]
        #     vmax = np.max(np.abs(vcur))
        #     vstartRel = vstart/vmax
        #     toPrint = "%s : %g, %g"%(toPrint, vstart, vstartRel)
        #     plot2D(X, vcur, ax=plot2D(X, ucur, bShow=False), legend=("u", "v"), title=toPrint)
        ures.append(np.hstack((bc[0], ucur, bc[1]))) # reshape?
        ue.append(np.hstack((bc[0], get_exact(X, r.t, c, x0).flatten(), bc[1])))
        mdiff.append(np.max(np.abs(ue[-1]-ures[-1])))
        # plot2D(X, r.y, bShow=False)
        xMid = get_cur_mid(X, vcur) # TODO - finding middle is not trivial...
        realMid1 = X[0, np.argmax(get_exact_x(X, r.t, c, x0))] 
        realMid2 = exact_mid(r.t)
        mids.append(xMid)
        rmids1.append(realMid1)
        rmids2.append(realMid2)
        if abs(xMid - xmaxCur) > widthTol:
            # uxcur = r.y.reshape(1, M2) @ (np.linalg.lstsq(R2(X, Ps, Pbs), R1(X, Ps, Pbs))[0])
            # break; # TODO - REMOVE (just not going fowrad here)
            uxcur = np.diff(ucur)
            Xog, Xo = Xg, X
            Ps_old, Pbs_old = Ps, Pbs
            # Xg, X = adaptiveGrid.get_grid(Xo, uxcur)
            Xg, X = adaptiveGrid.get_grid((Xo[0,:-1] + Xo[0,1:])/2, uxcur)
            if np.any(Xog != Xg) or np.any(Xo != X): # if there is an actual change in the grid
                # print("DIFF:\n", Xog != Xg)
                H_and_P = get_H_and_P(J, Xg, X, bHO)
                Ps, Pbs = H_and_P[0], H_and_P[1]
                # plot_grid(Xg, X)
                try:
                    ucur = change_grid(J, ucur, Ps_old, Pbs_old, X, Xog, Xo, R2, bHO, bc=bc)
                except ValueError as e:
                    print(xMid)
                    print(Xo, Xog, "\n", X, Xg)
                    print(e)
                    break
                # if len(swaps) < 2:
                #     ax = plot2D(Xo, ucur, bShow=False)
                #     plot2D(X, ucur, ax=ax, bShow=False)
                #     Xop = np.hstack((0, Xo.flatten(), 1))
                #     plot2D(Xop, ue[-1], ax=ax, legend=("PREV", "NEW", "EXACT"))
                try:
                    vcur = change_grid(J, vcur, Ps_old, Pbs_old, X, Xog, Xo, R2, bHO, bc=[0, 0], bInterpolate=True) # NEED to interpolate
                except ValueError as e:
                    print(xMid)
                    print(Xo, Xog, "\n", X, Xg)
                    print(e)
                    break
                Dx = get_Dxs(R2, R1, R0, X, Ps, Pbs)
                Sc2 = S2(X, Pbs)
                Sc1 = S1(X, Pbs)
                Sc0 = S0(X, Pbs)
                # plot2D(Xo, r.y,bShow=False),plot2D(X, ucur) # TODO - remove
                r.set_initial_value(np.hstack((ucur.flatten(), vcur.flatten())), r.t)
                r.set_f_params(X, M2, Dx, Sc2, Sc1, Sc0)
                xmaxCur = xMid
                swaps.append(r.t)
    print("DONE:", toPrint)
    T = np.array(tres).reshape(len(tres), 1)
    xx, tt = np.meshgrid(np.hstack((0, X.flatten(), 1)), T)
    xx = np.array(xres)
    print ('TF=', tres[-1], "(target tf=%g)"%tf)
    U = np.array(ures)
    Ue = np.array(ue)
    print('SHAPES:', U.shape, Ue.shape)
    print('MAX diff:', np.max(mdiff))
    if bDebug:
        from matplotlib import pyplot as plt
        plt.plot(T.flatten(), mids)
        plt.plot(T.flatten(), rmids1)
        plt.plot(T.flatten(), rmids2)
        plt.plot(T.flatten(), mdiff)
        plt.plot(swaps, swaps, 'o')
        plt.legend(("CALC", "Real1", "REALREL", "ERROR", "swaps"))
        plt.ylim((0, 1))
    return xx, tt, U, Ue

# def get_cur_mid(X, ux, target=0.5):
#     # iclosest = np.argmin(np.abs(ux - target))
#     # x1 = X[0, iclosest]
#     # # print(x1, iclosest)
#     # val1 = ux[iclosest]
#     # if (val1 > target):
#     #     x2 = X[0, iclosest + 1]
#     #     val2 = ux[iclosest + 1]
#     # else:
#     #     x2 = X[0, iclosest - 1]
#     #     val2 = ux[iclosest - 1]
#     # diff1 = abs(val1 - target)
#     # diff2 = abs(val2 - target)
#     # return (diff1 * x2 + diff2 * x1)/(diff1 + diff2)
#     return X[0, np.argmin(ux)]

# def get_cur_mid(X, u):
    # Xc = (X[0,:-1] + X[0,1:])/2
    # ux = np.diff(u)
    # return Xc[np.argmin(ux)]
def get_cur_mid(X, v):
    return X[0, np.argmin(v)]


def get_H_and_P(J, Xg, X, bHO):
    H = Hm_nu(Xg)
    P2 = Pn_nu(J, 2, Xg, X.flatten())
    P2b = Pnx_nu(J, 2, 1, Xg) 
    P1 = Pn_nu(J, 1, Xg, X.flatten())
    if bHO:
        P4 = Pn_nu(J, 4, Xg, X.flatten())
        P4b = Pnx_nu(J, 4, 1, Xg)
        P3 = Pn_nu(J, 3, Xg, X.flatten())
        return (H, P1, P2, P3, P4), (None, None, P2b, None, P4b)
    else:
        return (H, P1, P2)        , (None, None, P2b)

def change_grid(J, ucur, Pso, Pbso, X, Xog, Xo, R2, bHO, bInterpolate=True, bc=None):
    if not bInterpolate:
        Px2 = np.hstack([Pnx_nu(J, 2, x, Xog) for x in X.flatten()])
        if bHO:
            Px4 = np.hstack([Pnx_nu(J, 4, x, Xog) for x in X.flatten()])
        PsMID = [None, None, Px2] if not bHO else [None, None, Px2, None, Px4]
        return ucur @ np.linalg.lstsq(R2(Xo, Pso, Pbso), R2(X, PsMID, Pbso))[0]
    else:
        filledUcur = np.hstack((bc[0], ucur.flatten(), bc[1]))
        filledX = np.hstack((0, Xo.flatten(), 1))
        interpolant = interp1d(filledX, filledUcur, 3)
        unew = interpolant(X.flatten())
        return unew.reshape(*ucur.shape)

def get_Rs(X, bHO):
    if not bHO:
        c1 = lambda Pbs: 2 * np.pi - Pbs[2]
        c2 = lambda Pbs: 0 * Pbs[2] + 1
        R2 = lambda X, Ps, Pbs: Ps[2] + c1(Pbs) @ X      + c2(Pbs)
        R1 = lambda X, Ps, Pbs: Ps[1] + c1(Pbs) 
        R0 = lambda X, Ps, Pbs: Ps[0]
        return R2, R1, R0
    else:
        c1 = lambda Pbs: - Pbs[2]
        c2 = lambda Pbs: 0 * Pbs[2]
        c3 = lambda Pbs: 1/6 * (- 6 * Pbs[4] + Pbs[2]) # 1/16 * (12 * np.pi) = 2 * np.pi -> S
        c4 = lambda Pbs: 0 * Pbs[2]
        R4 = lambda X, Ps, Pbs: Ps[4] + c1(Pbs) @ X**3/6  + c2(Pbs) @ X**2/2 + c3(Pbs) @ X      + c4(Pbs)
        R3 = lambda X, Ps, Pbs: Ps[3] + c1(Pbs) @ X**2/2  + c2(Pbs) @ X      + c3(Pbs)
        R2 = lambda X, Ps, Pbs: Ps[2] + c1(Pbs) @ X       + c2(Pbs)
        return R4, R3, R2


def get_Ss(X, bHO):
    if not bHO:
        c1 = lambda Pbs: 2 * np.pi
        c2 = lambda Pbs: 0
        S2 = lambda X, Pbs: c1(Pbs) * X + c2(Pbs)
        S1 = lambda X, Pbs: c1(Pbs) * X**0
        S0 = lambda X, Pbs: 0 * Pbs[2] * X**0
        return S2, S1, S0
    else:
        c1 = lambda Pbs: 0
        c2 = lambda Pbs: 0
        c3 = lambda Pbs: 2 * np.pi
        c4 = lambda Pbs: 0
        S4 = lambda X, Pbs: c1(Pbs) * X**3/6 + c2(Pbs) * X**2/2 + c3(Pbs) * X + c4(Pbs)
        S3 = lambda X, Pbs: c1(Pbs) * X**2/2 + c2(Pbs) * X      + c3(Pbs)
        S2 = lambda X, Pbs: c1(Pbs) * X      + c2(Pbs) 
        return S4, S3, S2


def get_Dxs(R2, R1, R0, X, Ps, Pbs):
    R2c = R2(X, Ps, Pbs)
    Dx = np.linalg.lstsq(R2c, R0(X, Ps, Pbs))[0]
    return Dx


def get_exact(X, T, c, x0=1/4):
    return 4 * np.arctan(np.exp((X - c * T - x0)/np.sqrt(1 - c**2)))


def get_exact_x(X, T, c, x0=1/4):
    # D(arctan(y),x) = y'/(1 + y^2)
    y = np.exp((X - c * T - x0)/np.sqrt(1 - c**2))
    yp = y/np.sqrt(1 - c**2)
    return 4 * yp/(1 + y**2)

def get_exact_t(X, T, c, x0=1/4):
    # D(arctan(y),t) = y'/(1 + y^2)
    y = np.exp((X - c * T - x0)/np.sqrt(1 - c**2))
    yp = -c * y/np.sqrt(1 - c**2)
    return 4 * yp/(1 + y**2)

def fun(t, uv, X, M2, Dx, S2, S1, S0):
    u = uv[:M2].reshape(1, M2)
    v = uv[M2:] # = dudt
    uxx = (u - S2) @ Dx + S0
    dvdt = uxx - np.sin(u)
    # if t < 4.4e-7:
    #     plot2D(X, v, ax=plot2D(X, u, bShow=False), legend=("u", "v"), title="t=%g"%t, xlims=[0, 0.4], ylims=[-.01,.01])
    return np.hstack((v, dvdt.flatten()))



if __name__ == '__main__':
    x0 = .3
    # widthTol = 1/35
    # fineWidth = .1
    c = 1 - 1e-4
    tf = (1 - 2 * x0)/c
    nrOfBorders = 1
    JRange = [4,]#5,6,7]
    for J in JRange:
        bests = []
        for fineWidth in np.arange(.2, .5, .01):
        # for fineWidth in [0.37]:
        #     for widthTol in [0.095]:
            for widthTol in np.arange(.04, .15, .005):
            # for widthTol in [.055]:#np.arange(.01, .1, .005):
                if J == 4:
                    aValues = np.arange(.881, .906, .001)
                    # aValues = np.arange(.892, .906, .001)
                    # aValues = [.896,]
                elif J == 5:
                    # aValues = np.arange(.963, .976, .001)
                    aValues = np.arange(.960, .980, .001)
                    # aValues = [.974,]
                elif J == 6:
                    aValues = np.arange(1, 1.0001, .001)
                elif J == 7:
                    aValues = [1.,]
                for a in aValues:
                    mStr = "J=%d, HOHWM, fineWidth = %g, widthTol=%g, a=%g"%(J, fineWidth, widthTol, a)
                    print(mStr)
                    # X, T, U, Ue = solve_kdv(J, alpha=alpha, beta=beta, c=c, tf=tf, bHO=True, x0=x0, fineWidth=fineWidth, widthTol=widthTol, borders=nrOfBorders, a=a)
                    try:
                        X, T, U, Ue = solve_SG(J, c=c, x0=x0, a=a, widthTol=widthTol,fineWidth=fineWidth, borders=nrOfBorders)
                    except ValueError as e:
                        print('Got exception (continuing on next)', e)
                    md = np.max(np.abs(U - Ue))
                    bests.append((fineWidth, widthTol, a, np.max(T), md))
                    #mdiffs.append(np.max(np.abs(U - Ue)))
                    # plot3D(X, T, U, bShow=False, title=mStr, zlims=[0, 2*np.pi]),plot3D(X,T,Ue, bShow=False),plot3D(X, T, U-Ue)
                # iMin = np.argmin(mdiffs)
                # print("best result: ", mdiffs[iMin], "at", aValues[iMin])
                    # bests.append((aValues[iMin], mdiffs[iMin]))
        print ("BEST:\n", bests)
        minDiff = 1e10
        minFW = "N/A"
        minWT = "N/A"
        minA = "N/A"
        for fw, wt, a, ctf, md in bests:
            if (ctf >= tf) and (md < minDiff):
                minDiff = md
                minFW = fw
                minWT = wt
                minA = a
        print ("Best at fw=", minFW, "wt=", minWT, "a=", minA, " with max diff of ", minDiff)