import utils


def advection_periodic_list_beta():
    list_beta = [1.e-02, 1.e-01, 1.e+00, 1.e+01, 4.e+01, 1.e+02, 5.e+02, 1.e+03, 5.e+03, 1.e+04] # follow the setting of SWIM-ODE
    list_n = [16] * len(list_beta)
    list_k = [8] * len(list_beta)

    config = {
        "pde": "advection_periodic",
        "beta": list_beta,
        "n": list_n,
        "k": list_k,
        "output_name": f"advection_periodic_list_beta",
        "plot_type": "beta"
    }
    utils.save_json(f"config/config_advection_periodic_list_beta.json", config)

def advection_periodic_same_beta(beta: float):
    list_n_fun = [2, 5, 8, 11, 14, 17, 20, 23, 26, 30] # follow the setting of SWIM-ODE
    list_n = [i+1 for i in list_n_fun]
    list_k = [i for i in range(1, len(list_n)+1)]
    list_beta = [beta] * len(list_n)

    config = {
        "pde": "advection_periodic",
        "beta": list_beta,
        "n": list_n,
        "k": list_k,
        "output_name": f"advection_periodic_beta_{beta}"
    }
    utils.save_json(f"config/config_advection_periodic_beta_{beta}.json", config)

def advection_periodic_single(beta: float, n: int, k: int):
    config = {
        "pde": "advection_periodic",
        "beta": [beta],
        "n": [n],
        "k": [k],
        "output_name": "advection_periodic_single"
    }
    utils.save_json(f"config/config_advection_periodic_single.json", config)

def euler_bernoulli():
    num_start = 2
    num_end = 10
    list_n = [i*3 for i in range(num_start, num_end)]
    list_k = [i for i in range(num_start, num_end)]

    config = {
        "pde": "euler_bernoulli",
        "n": list_n,
        "k": list_k,
        "output_name": "euler_bernoulli"
    }
    utils.save_json(f"config/config_euler_bernoulli.json", config)

def euler_bernoulli_single(n: int, k: int):
    config = {
        "pde": "euler_bernoulli",
        "n": [n],
        "k": [k],
        "output_name": "euler_bernoulli_single"
    }
    utils.save_json(f"config/config_euler_bernoulli_single.json", config)

def burger_dirichlet():
    num_start = 1
    num_end = 9
    list_n = [i*50-i%2 for i in range(num_start, num_end)]
    list_k = [i for i in range(num_start, num_end)]

    config = {
        "pde": "burger_dirichlet",
        "n": list_n,
        "k": list_k,
        "output_name": "burger_dirichlet"
    }
    utils.save_json(f"config/config_burger_dirichlet.json", config)

def burger_dirichlet_single(n: int, k: int):
    config = {
        "pde": "burger_dirichlet",
        "n": [n],
        "k": [k],
        "output_name": "burger_dirichlet_single"
    }
    utils.save_json(f"config/config_burger_dirichlet_single.json", config)



if __name__ == "__main__":
    advection_periodic_same_beta(beta=40) # convergence test for advection equation with fixed beta
    euler_bernoulli() # convergence test for euler bernoulli beam
    burger_dirichlet() # convergence test for burger
    advection_periodic_list_beta() # test for advection with different betas
    burger_dirichlet_single(n=400, k=8) # a single test for burger
    euler_bernoulli_single(n=27, k=9) # a single test for euler bernoulli
    advection_periodic_single(beta=40, n=16, k=9) # a single test for advection
