from model.STGNCDE.GCDE import *


def make_model(args):
    vector_field_f = FinalTanh_f(
        input_channels=args["input_dim"],
        hidden_channels=args["hid_dim"],
        hidden_hidden_channels=args["hid_hid_dim"],
        num_hidden_layers=args["num_layers"],
    )
    vector_field_g = VectorField_g(
        input_channels=args["input_dim"],
        hidden_channels=args["hid_dim"],
        hidden_hidden_channels=args["hid_hid_dim"],
        num_hidden_layers=args["num_layers"],
        num_nodes=args["num_nodes"],
        cheb_k=args["cheb_k"],
        embed_dim=args["embed_dim"],
        g_type=args["g_type"],
    )
    model = NeuralGCDE(
        args,
        func_f=vector_field_f,
        func_g=vector_field_g,
        input_channels=args["input_dim"],
        hidden_channels=args["hid_dim"],
        output_channels=args["output_dim"],
        initial=True,
        device=args["device"],
        atol=1e-9,
        rtol=1e-7,
        solver=args["solver"],
    )
    # return model, vector_field_f, vector_field_g
    return model