import numpy as np
import deepxde as dde
from ..cases.PDECases import AllenCahn
from ..solver.PDESolver import PINNSolver
import time

dde.optimizers.config.set_LBFGS_options(maxiter=1000)

class RAR_G(PINNSolver):
    def __init__(self, PDECase):
        super().__init__(name='RAR_G', PDECase=PDECase)

    def closure(self):
        self.train_step()

        for i in range(100):
            X = self.PDECase.geomtime.random_points(100000)
            Y = np.abs(self.model.predict(X, operator=self.PDECase.pde))[:, 0]
            t1 = time.time()
            X_ids = np.argpartition(Y, self.PDECase.NumDomain//100, axis=0)
            print(f'{time.time()-t1} s')
            self.PDECase.data.add_anchors(X[X_ids])

            self.train_step(iterations=1000)

if __name__ == '__main__':
    PDECase = AllenCahn(NumDomain=2000//2)
    solver = RAR_G(PDECase=PDECase)
    solver.train()
    solver.save(add_time=True)