import skopt
import numpy as np
import deepxde as dde
from cases.PDECases import Burgers
from solver.PDESolver import PINNSolver
from distutils.version import LooseVersion

class Uniform(PINNSolver):
    def __init__(self, PDECase, method):
        self.method = method
        super().__init__(name=f'Uniform_{method}', PDECase=PDECase)
    
    def gen_data(self):
        if self.method == 'Grid':
            data = dde.data.TimePDE(self.PDECase.geomtime, self.PDECase.pde, [], num_domain=self.PDECase.NumDomain, train_distribution='uniform')
        elif self.method == 'Random':
            data = dde.data.TimePDE(self.PDECase.geomtime, self.PDECase.pde, [], num_domain=self.PDECase.NumDomain, train_distribution='pseudo')
        elif self.method in ['LHS', 'Halton', 'Hammersley', 'Sobol']:
            sample_pts = self.quasirandom(self.PDECase.NumDomain, self.method)
            data = dde.data.TimePDE(self.PDECase.geomtime, self.PDECase.pde, [], num_domain=0, train_distribution='uniform', anchors=sample_pts)
        return data
    
    def quasirandom(self, n_samples, sampler):
        space = [(-1.0, 1.0), (0.0, 1.0)]
        if sampler == "LHS":
            sampler = skopt.sampler.Lhs(
                lhs_type="centered", criterion="maximin", iterations=1000
            )
        elif sampler == "Halton":
            sampler = skopt.sampler.Halton(min_skip=-1, max_skip=-1)
        elif sampler == "Hammersley":
            sampler = skopt.sampler.Hammersly(min_skip=-1, max_skip=-1)
        elif sampler == "Sobol":
            # Remove the first point [0, 0, ...] and the second point [0.5, 0.5, ...], which
            # are too special and may cause some error.
            if LooseVersion(skopt.__version__) < LooseVersion("0.9"):
                sampler = skopt.sampler.Sobol(min_skip=2, max_skip=2, randomize=False)
            else:
                sampler = skopt.sampler.Sobol(skip=0, randomize=False)
                return np.array(
                    sampler.generate(space, n_samples + 2)[2:]
                )
        return np.array(sampler.generate(space, n_samples))
    
if __name__ == '__main__':
    PDECase = Burgers(NumDomain=2000)
    solver = Uniform(PDECase=PDECase, method='Grid')

    #solver = Uniform(PDECase=PDECase, method='Random')
    #solver = Uniform(PDECase=PDECase, method='LHS')
    #solver = Uniform(PDECase=PDECase, method='Halton')
    #solver = Uniform(PDECase=PDECase, method='Hammersley')
    #solver = Uniform(PDECase=PDECase, method='Sobol')

    solver.train()
    solver.save(add_time=True)
    solver.plot_loss_history(train=True, use_time=True)