from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
import h5py
import deepxde as dde
from deepxde.backend import tf


def gen_testdata():
    len_x = 121; DX = 0.01; len_t = 51; DT = 0.01; start_x = -0.6; start_t = 0
    X = np.zeros([len_x*len_t,2])
    for i in range(len_t):
        for j in range(len_x):
            idx = i * len_x + j
            X[idx, 0] = j*DX + start_x
            X[idx, 1] = i*DT + start_t
    return X


def main():
    def pde(x, y):
        dy_x = tf.gradients(y, x)[0]
        dy_x, dy_t = dy_x[:, 0:1], dy_x[:, 1:2]
        # return dy_t + (-2*y +0.5)* dy_x
        return dy_t + y*dy_x

    geom = dde.geometry.Interval(-0.5, 0.5)
    timedomain = dde.geometry.TimeDomain(0, 0.1)
    geomtime = dde.geometry.GeometryXTime(geom, timedomain)

    bc = dde.PeriodicBC(geomtime, 0, lambda _, on_boundary: on_boundary)
    ic = dde.IC(
        # 0.8 + 0.5 * torch.sin((x0f-0.1)*2.*math.pi)
        
        # geomtime, lambda x: 0.8+0.5*np.sin(2 * np.pi * (x[:, 0:1]-0.1)), lambda _, on_initial: on_initial
        geomtime, lambda x: 0.5+np.sin(2 * np.pi * (x[:, 0:1])), lambda _, on_initial: on_initial
    )

    data = dde.data.TimePDE(
        geomtime, pde, [bc, ic], num_domain=2500, num_boundary=100, num_initial=160
    )
    net = dde.maps.FNN([2] + [20] * 3 + [1], "tanh", "Glorot normal")
    model = dde.Model(data, net)

    model.compile("adam", lr=1.0e-3)
    model.train(epochs=10000)
    model.compile("L-BFGS-B")
    model.train()

    X = geomtime.random_points(100000)
    err = 1
    while err > 0.005:
        f = model.predict(X, operator=pde)
        err_eq = np.absolute(f)
        err = np.mean(err_eq)
        print("Mean residual: %.3e" % (err))

        x_id = np.argmax(err_eq)
        print("Adding new point:", X[x_id], "\n")
        data.add_anchors(X[x_id])
        early_stopping = dde.callbacks.EarlyStopping(min_delta=1e-4, patience=2000)
        model.compile("adam", lr=1e-3)
        model.train(
            epochs=10000, disregard_previous_best=True, callbacks=[early_stopping]
        )
        model.compile("L-BFGS-B")
        losshistory, train_state = model.train()
    dde.saveplot(losshistory, train_state, issave=True, isplot=True)

    X = gen_testdata()
    y_pred = model.predict(X)
    hf = h5py.File("deepxde_discont.h5", "w")
    hf.create_dataset('pred', data=y_pred)


if __name__ == "__main__":
    main()
