from distutils.command.clean import clean
import nox
import os
import sys

# Use system installed Python packages
PYT_PATH = (
    "/usr/local/lib/python3.10/dist-packages"
    if not "PYT_PATH" in os.environ
    else os.environ["PYT_PATH"]
)
print(f"Using python path {PYT_PATH}")

# Set the root directory to the directory of the noxfile unless the user wants to
# TOP_DIR
TOP_DIR = (
    os.path.dirname(os.path.realpath(__file__))
    if not "TOP_DIR" in os.environ
    else os.environ["TOP_DIR"]
)
print(f"Test root directory {TOP_DIR}")

# Set the USE_CXX11=1 to use cxx11_abi
USE_CXX11 = 0 if not "USE_CXX11" in os.environ else os.environ["USE_CXX11"]
if USE_CXX11:
    print("Using cxx11 abi")

# Set the USE_HOST_DEPS=1 to use host dependencies for tests
USE_HOST_DEPS = 0 if not "USE_HOST_DEPS" in os.environ else os.environ["USE_HOST_DEPS"]
if USE_HOST_DEPS:
    print("Using dependencies from host python")

# Set epochs to train VGG model for accuracy tests
EPOCHS = 25

SUPPORTED_PYTHON_VERSIONS = ["3.7", "3.8", "3.9", "3.10"]

nox.options.sessions = [
    "l0_api_tests-" + "{}.{}".format(sys.version_info.major, sys.version_info.minor)
]


def install_deps(session):
    print("Installing deps")
    session.install("-r", os.path.join(TOP_DIR, "py", "requirements.txt"))
    session.install("-r", os.path.join(TOP_DIR, "tests", "py", "requirements.txt"))


def download_models(session):
    print("Downloading test models")
    session.install("-r", os.path.join(TOP_DIR, "tests", "modules", "requirements.txt"))
    print(TOP_DIR)
    session.chdir(os.path.join(TOP_DIR, "tests", "modules"))
    if USE_HOST_DEPS:
        session.run_always("python", "hub.py", env={"PYTHONPATH": PYT_PATH})
    else:
        session.run_always("python", "hub.py")


def install_torch_trt(session):
    print("Installing latest torch-tensorrt build")
    session.chdir(os.path.join(TOP_DIR, "py"))
    if USE_CXX11:
        session.run("python", "setup.py", "develop", "--use-cxx11-abi")
    else:
        session.run("python", "setup.py", "develop")


def train_model(session):
    session.chdir(os.path.join(TOP_DIR, "examples/int8/training/vgg16"))
    session.install("-r", "requirements.txt")
    if os.path.exists("vgg16_ckpts/ckpt_epoch25.pth"):
        session.run_always("python", "export_ckpt.py", "vgg16_ckpts/ckpt_epoch25.pth")
        return
    if USE_HOST_DEPS:
        session.run_always(
            "python",
            "main.py",
            "--lr",
            "0.01",
            "--batch-size",
            "128",
            "--drop-ratio",
            "0.15",
            "--ckpt-dir",
            "vgg16_ckpts",
            "--epochs",
            str(EPOCHS),
            env={"PYTHONPATH": PYT_PATH},
        )

        session.run_always(
            "python",
            "export_ckpt.py",
            "vgg16_ckpts/ckpt_epoch" + str(EPOCHS) + ".pth",
            env={"PYTHONPATH": PYT_PATH},
        )
    else:
        session.run_always(
            "python",
            "main.py",
            "--lr",
            "0.01",
            "--batch-size",
            "128",
            "--drop-ratio",
            "0.15",
            "--ckpt-dir",
            "vgg16_ckpts",
            "--epochs",
            str(EPOCHS),
        )

        session.run_always(
            "python", "export_ckpt.py", "vgg16_ckpts/ckpt_epoch" + str(EPOCHS) + ".pth"
        )


def finetune_model(session):
    # Install pytorch-quantization dependency
    session.install(
        "pytorch-quantization", "--extra-index-url", "https://pypi.ngc.nvidia.com"
    )
    session.chdir(os.path.join(TOP_DIR, "examples/int8/training/vgg16"))

    if USE_HOST_DEPS:
        session.run_always(
            "python",
            "finetune_qat.py",
            "--lr",
            "0.01",
            "--batch-size",
            "128",
            "--drop-ratio",
            "0.15",
            "--ckpt-dir",
            "vgg16_ckpts",
            "--start-from",
            str(EPOCHS),
            "--epochs",
            str(EPOCHS + 1),
            env={"PYTHONPATH": PYT_PATH},
        )

        # Export model
        session.run_always(
            "python",
            "export_qat.py",
            "vgg16_ckpts/ckpt_epoch" + str(EPOCHS + 1) + ".pth",
            env={"PYTHONPATH": PYT_PATH},
        )
    else:
        session.run_always(
            "python",
            "finetune_qat.py",
            "--lr",
            "0.01",
            "--batch-size",
            "128",
            "--drop-ratio",
            "0.15",
            "--ckpt-dir",
            "vgg16_ckpts",
            "--start-from",
            str(EPOCHS),
            "--epochs",
            str(EPOCHS + 1),
        )

        # Export model
        session.run_always(
            "python",
            "export_qat.py",
            "vgg16_ckpts/ckpt_epoch" + str(EPOCHS + 1) + ".pth",
        )


def cleanup(session):
    target = [
        "examples/int8/training/vgg16/*.jit.pt",
        "examples/int8/training/vgg16/vgg16_ckpts",
        "examples/int8/training/vgg16/cifar-10-*",
        "examples/int8/training/vgg16/data",
        "tests/modules/*.jit.pt",
        "tests/py/*.jit.pt",
    ]

    target = " ".join(x for x in [os.path.join(TOP_DIR, i) for i in target])
    session.run_always("bash", "-c", str("rm -rf ") + target, external=True)


def run_base_tests(session):
    print("Running basic tests")
    session.chdir(os.path.join(TOP_DIR, "tests/py"))
    tests = [
        "api",
        "integrations/test_to_backend_api.py",
    ]
    for test in tests:
        if USE_HOST_DEPS:
            session.run_always("pytest", test, env={"PYTHONPATH": PYT_PATH})
        else:
            session.run_always("pytest", test)


def run_fx_core_tests(session):
    print("Running FX core tests")
    session.chdir(os.path.join(TOP_DIR, "py/torch_tensorrt/fx/test"))
    tests = [
        "core",
    ]
    for test in tests:
        if USE_HOST_DEPS:
            session.run_always("pytest", test, env={"PYTHONPATH": PYT_PATH})
        else:
            session.run_always("pytest", test)


def run_fx_converter_tests(session):
    print("Running FX converter tests")
    session.chdir(os.path.join(TOP_DIR, "py/torch_tensorrt/fx/test"))
    tests = [
        "converters",
    ]
    # Skipping this test as it fails inside NGC container with the following error.
    # Error Code 4: Internal Error (Could not find any implementation for node conv due to insufficient workspace. See verbose log for requested sizes.)
    skip_tests = "-k not conv3d"
    for test in tests:
        if USE_HOST_DEPS:
            session.run_always("pytest", test, skip_tests, env={"PYTHONPATH": PYT_PATH})
        else:
            session.run_always("pytest", test, skip_tests)


def run_fx_lower_tests(session):
    print("Running FX passes and trt_lower tests")
    session.chdir(os.path.join(TOP_DIR, "py/torch_tensorrt/fx/test"))
    tests = [
        "passes/test_multi_fuse_trt.py",
        # "passes/test_fuse_permute_linear_trt.py",
        "passes/test_remove_duplicate_output_args.py",
        "passes/test_fuse_permute_matmul_trt.py",
        # "passes/test_graph_opts.py"
        "trt_lower",
    ]
    for test in tests:
        if USE_HOST_DEPS:
            session.run_always("pytest", test, env={"PYTHONPATH": PYT_PATH})
        else:
            session.run_always("pytest", test)


def run_fx_quant_tests(session):
    print("Running FX Quant tests")
    session.chdir(os.path.join(TOP_DIR, "py/torch_tensorrt/fx/test"))
    tests = [
        "quant",
    ]
    # Skipping this test as it fails inside NGC container with the following error.
    # ImportError: cannot import name 'ObservationType' from 'torch.ao.quantization.backend_config.observation_type'
    skip_tests = "-k not conv_add_standalone_module"
    for test in tests:
        if USE_HOST_DEPS:
            session.run_always("pytest", test, skip_tests, env={"PYTHONPATH": PYT_PATH})
        else:
            session.run_always("pytest", test, skip_tests)


def run_fx_tracer_tests(session):
    print("Running FX Tracer tests")
    session.chdir(os.path.join(TOP_DIR, "py/torch_tensorrt/fx/test"))
    # skipping a test since it depends on torchdynamo
    # Enable this test once NGC moves to latest pytorch which has dynamo integrated.
    tests = [
        "tracer/test_acc_shape_prop.py",
        "tracer/test_acc_tracer.py",
        # "tracer/test_dispatch_tracer.py"
    ]
    for test in tests:
        if USE_HOST_DEPS:
            session.run_always("pytest", test, env={"PYTHONPATH": PYT_PATH})
        else:
            session.run_always("pytest", test)


def run_fx_tools_tests(session):
    print("Running FX tools tests")
    session.chdir(os.path.join(TOP_DIR, "py/torch_tensorrt/fx/test"))
    tests = [
        "tools",
    ]
    for test in tests:
        if USE_HOST_DEPS:
            session.run_always("pytest", test, env={"PYTHONPATH": PYT_PATH})
        else:
            session.run_always("pytest", test)


def run_model_tests(session):
    print("Running model tests")
    session.chdir(os.path.join(TOP_DIR, "tests/py"))
    tests = [
        "models",
    ]
    for test in tests:
        if USE_HOST_DEPS:
            session.run_always("pytest", test, env={"PYTHONPATH": PYT_PATH})
        else:
            session.run_always("pytest", test)


def run_accuracy_tests(session):
    print("Running accuracy tests")
    session.chdir(os.path.join(TOP_DIR, "tests/py"))
    tests = []
    for test in tests:
        if USE_HOST_DEPS:
            session.run_always("python", test, env={"PYTHONPATH": PYT_PATH})
        else:
            session.run_always("python", test)


def copy_model(session):
    model_files = ["trained_vgg16.jit.pt", "trained_vgg16_qat.jit.pt"]

    for file_name in model_files:
        src_file = os.path.join(
            TOP_DIR, str("examples/int8/training/vgg16/") + file_name
        )
        if os.path.exists(src_file):
            session.run_always(
                "cp",
                "-rpf",
                os.path.join(TOP_DIR, src_file),
                os.path.join(TOP_DIR, str("tests/modules/") + file_name),
                external=True,
            )


def run_int8_accuracy_tests(session):
    print("Running accuracy tests")
    copy_model(session)
    session.chdir(os.path.join(TOP_DIR, "tests/py"))
    tests = [
        "ptq/test_ptq_to_backend.py",
        "ptq/test_ptq_dataloader_calibrator.py",
        "qat/",
    ]
    for test in tests:
        if USE_HOST_DEPS:
            session.run_always("pytest", test, env={"PYTHONPATH": PYT_PATH})
        else:
            session.run_always("pytest", test)


def run_trt_compatibility_tests(session):
    print("Running TensorRT compatibility tests")
    copy_model(session)
    session.chdir(os.path.join(TOP_DIR, "tests/py"))
    tests = [
        "integrations/test_trt_intercompatibility.py",
        # "ptq/test_ptq_trt_calibrator.py",
    ]
    for test in tests:
        if USE_HOST_DEPS:
            session.run_always("pytest", test, env={"PYTHONPATH": PYT_PATH})
        else:
            session.run_always("pytest", test)


def run_dla_tests(session):
    print("Running DLA tests")
    session.chdir(os.path.join(TOP_DIR, "tests/py"))
    tests = [
        "hw/test_api_dla.py",
    ]
    for test in tests:
        if USE_HOST_DEPS:
            session.run_always("pytest", test, env={"PYTHONPATH": PYT_PATH})
        else:
            session.run_always("pytest", test)


def run_multi_gpu_tests(session):
    print("Running multi GPU tests")
    session.chdir(os.path.join(TOP_DIR, "tests/py"))
    tests = [
        "hw/test_multi_gpu.py",
    ]
    for test in tests:
        if USE_HOST_DEPS:
            session.run_always("pytest", test, env={"PYTHONPATH": PYT_PATH})
        else:
            session.run_always("pytest", test)


def run_l0_api_tests(session):
    if not USE_HOST_DEPS:
        install_deps(session)
        install_torch_trt(session)
    download_models(session)
    run_base_tests(session)
    cleanup(session)


def run_l0_fx_tests(session):
    if not USE_HOST_DEPS:
        install_deps(session)
        install_torch_trt(session)
    run_fx_core_tests(session)
    run_fx_converter_tests(session)
    run_fx_lower_tests(session)
    cleanup(session)


def run_l0_fx_core_tests(session):
    if not USE_HOST_DEPS:
        install_deps(session)
        install_torch_trt(session)
    run_fx_core_tests(session)
    cleanup(session)


def run_l0_fx_converter_tests(session):
    if not USE_HOST_DEPS:
        install_deps(session)
        install_torch_trt(session)
    run_fx_converter_tests(session)
    cleanup(session)


def run_l0_fx_lower_tests(session):
    if not USE_HOST_DEPS:
        install_deps(session)
        install_torch_trt(session)
    run_fx_lower_tests(session)
    cleanup(session)


def run_l0_dla_tests(session):
    if not USE_HOST_DEPS:
        install_deps(session)
        install_torch_trt(session)
    download_models(session)
    run_base_tests(session)
    cleanup(session)


def run_l1_model_tests(session):
    if not USE_HOST_DEPS:
        install_deps(session)
        install_torch_trt(session)
    download_models(session)
    run_model_tests(session)
    cleanup(session)


def run_l1_int8_accuracy_tests(session):
    if not USE_HOST_DEPS:
        install_deps(session)
        install_torch_trt(session)
    train_model(session)
    finetune_model(session)
    run_int8_accuracy_tests(session)
    cleanup(session)


def run_l1_fx_tests(session):
    if not USE_HOST_DEPS:
        install_deps(session)
        install_torch_trt(session)
    run_fx_quant_tests(session)
    run_fx_tracer_tests(session)
    run_fx_tools_tests(session)
    cleanup(session)


def run_l2_trt_compatibility_tests(session):
    if not USE_HOST_DEPS:
        install_deps(session)
        install_torch_trt(session)
    run_trt_compatibility_tests(session)
    cleanup(session)


def run_l2_multi_gpu_tests(session):
    if not USE_HOST_DEPS:
        install_deps(session)
        install_torch_trt(session)
    download_models(session)
    run_multi_gpu_tests(session)
    cleanup(session)


@nox.session(python=SUPPORTED_PYTHON_VERSIONS, reuse_venv=True)
def l0_api_tests(session):
    """When a developer needs to check correctness for a PR or something"""
    run_l0_api_tests(session)


@nox.session(python=SUPPORTED_PYTHON_VERSIONS, reuse_venv=True)
def l0_fx_tests(session):
    """When a developer needs to check correctness for a PR or something"""
    run_l0_fx_tests(session)


@nox.session(python=SUPPORTED_PYTHON_VERSIONS, reuse_venv=True)
def l0_fx_core_tests(session):
    """When a developer needs to check correctness for a PR or something"""
    run_l0_fx_core_tests(session)


@nox.session(python=SUPPORTED_PYTHON_VERSIONS, reuse_venv=True)
def l0_fx_converter_tests(session):
    """When a developer needs to check correctness for a PR or something"""
    run_l0_fx_converter_tests(session)


@nox.session(python=SUPPORTED_PYTHON_VERSIONS, reuse_venv=True)
def l0_fx_lower_tests(session):
    """When a developer needs to check correctness for a PR or something"""
    run_l0_fx_lower_tests(session)


@nox.session(python=SUPPORTED_PYTHON_VERSIONS, reuse_venv=True)
def l0_dla_tests(session):
    """When a developer needs to check basic api functionality using host dependencies"""
    run_l0_dla_tests(session)


@nox.session(python=SUPPORTED_PYTHON_VERSIONS, reuse_venv=True)
def l1_model_tests(session):
    """When a user needs to test the functionality of standard models compilation and results"""
    run_l1_model_tests(session)


@nox.session(python=SUPPORTED_PYTHON_VERSIONS, reuse_venv=True)
def l1_fx_tests(session):
    """When a user needs to test the functionality of standard models compilation and results"""
    run_l1_fx_tests(session)


@nox.session(python=SUPPORTED_PYTHON_VERSIONS, reuse_venv=True)
def l1_int8_accuracy_tests(session):
    """Checking accuracy performance on various usecases"""
    run_l1_int8_accuracy_tests(session)


@nox.session(python=SUPPORTED_PYTHON_VERSIONS, reuse_venv=True)
def l2_trt_compatibility_tests(session):
    """Makes sure that TensorRT Python and Torch-TensorRT can work together"""
    run_l2_trt_compatibility_tests(session)


@nox.session(python=SUPPORTED_PYTHON_VERSIONS, reuse_venv=True)
def l2_multi_gpu_tests(session):
    """Makes sure that Torch-TensorRT can operate on multi-gpu systems"""
    run_l2_multi_gpu_tests(session)