from deepxde.backend import backend_name, tf, torch, jax, paddle


def empty_cache():
    # TODO: generated by github copilot, not tested
    if backend_name == "pytorch":
        torch.cuda.empty_cache()
    elif backend_name == "tensorflow":
        tf.compat.v1.reset_default_graph()
        tf.compat.v1.keras.backend.clear_session()
    elif backend_name == "jax":
        jax.local_devices()[0].invalidate()
    elif backend_name == "paddle":
        paddle.device.cuda.empty_cache()
    else:
        raise ValueError(f"Unknown backend: {backend_name=}")


def setup_seed(seed=0):
    import os
    import random
    import numpy as np

    np.random.seed(seed)
    random.seed(seed)

    if backend_name == "tensorflow":
        tf.compat.v1.set_random_seed(seed)
        tf.random.set_seed(seed)
        os.environ["TF_DETERMINISTIC_OPS"] = "1"
    elif backend_name == "tensorflow.compat.v1":
        os.environ["TF_CUDNN_DETERMINISTIC"] = "1"
        session_conf = tf.compat.v1.ConfigProto(
            intra_op_parallelism_threads=1, inter_op_parallelism_threads=1
        )
        sess = tf.compat.v1.Session(
            graph=tf.compat.v1.get_default_graph(), config=session_conf
        )
        tf.compat.v1.keras.backend.set_session(sess)
    elif backend_name == "jax":
        jax.random.PRNGKey(seed)
    elif backend_name == "paddle":
        paddle.seed(seed)
    elif backend_name == "pytorch":
        torch.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        torch.backends.cudnn.deterministic = True
        os.environ["PYTHONHASHSEED"] = str(seed)
    else:
        raise ValueError(f"Unknown backend: {backend_name=}")
