import argparse
import os


past_versions_testing = {
    "pytorch": {
        "1.12": {
            "torch": "1.12.1",
            "torchvision": "0.13.1",
            "torchaudio": "0.12.1",
            "python": 3.9,
            "cuda": "cu113",
            "install": (
                "python3 -m pip install --no-cache-dir -U torch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1"
                " --extra-index-url https://download.pytorch.org/whl/cu113"
            ),
        },
        "1.11": {
            "torch": "1.11.0",
            "torchvision": "0.12.0",
            "torchaudio": "0.11.0",
            "python": 3.9,
            "cuda": "cu113",
            "install": (
                "python3 -m pip install --no-cache-dir -U torch==1.11.0 torchvision==0.12.0 torchaudio==0.11.0"
                " --extra-index-url https://download.pytorch.org/whl/cu113"
            ),
        },
        "1.10": {
            "torch": "1.10.2",
            "torchvision": "0.11.3",
            "torchaudio": "0.10.2",
            "python": 3.9,
            "cuda": "cu113",
            "install": (
                "python3 -m pip install --no-cache-dir -U torch==1.10.2 torchvision==0.11.3 torchaudio==0.10.2"
                " --extra-index-url https://download.pytorch.org/whl/cu113"
            ),
        },
        # torchaudio < 0.10 has no CUDA-enabled binary distributions
        "1.9": {
            "torch": "1.9.1",
            "torchvision": "0.10.1",
            "torchaudio": "0.9.1",
            "python": 3.9,
            "cuda": "cu111",
            "install": (
                "python3 -m pip install --no-cache-dir -U torch==1.9.1 torchvision==0.10.1 torchaudio==0.9.1"
                " --extra-index-url https://download.pytorch.org/whl/cu111"
            ),
        },
        "1.8": {
            "torch": "1.8.1",
            "torchvision": "0.9.1",
            "torchaudio": "0.8.1",
            "python": 3.9,
            "cuda": "cu111",
            "install": (
                "python3 -m pip install --no-cache-dir -U torch==1.8.1 torchvision==0.9.1 torchaudio==0.8.1"
                " --extra-index-url https://download.pytorch.org/whl/cu111"
            ),
        },
        "1.7": {
            "torch": "1.7.1",
            "torchvision": "0.8.2",
            "torchaudio": "0.7.2",
            "python": 3.9,
            "cuda": "cu110",
            "install": (
                "python3 -m pip install --no-cache-dir -U torch==1.7.1 torchvision==0.8.2 torchaudio==0.7.2"
                " --extra-index-url https://download.pytorch.org/whl/cu110"
            ),
        },
        "1.6": {
            "torch": "1.6.0",
            "torchvision": "0.7.0",
            "torchaudio": "0.6.0",
            "python": 3.8,
            "cuda": "cu101",
            "install": (
                "python3 -m pip install --no-cache-dir -U torch==1.6.0 torchvision==0.7.0 torchaudio==0.6.0"
                " --extra-index-url https://download.pytorch.org/whl/cu101"
            ),
        },
        "1.5": {
            "torch": "1.5.1",
            "torchvision": "0.6.1",
            "torchaudio": "0.5.1",
            "python": 3.8,
            "cuda": "cu101",
            "install": (
                "python3 -m pip install --no-cache-dir -U torch==1.5.1 torchvision==0.6.1 torchaudio==0.5.1"
                " --extra-index-url https://download.pytorch.org/whl/cu101"
            ),
        },
        "1.4": {
            "torch": "1.4.0",
            "torchvision": "0.5.0",
            "torchaudio": "0.4.0",
            "python": 3.8,
            "cuda": "cu100",
            "install": (
                "python3 -m pip install --no-cache-dir -U torch==1.4.0 torchvision==0.5.0 torchaudio==0.4.0"
                " --extra-index-url https://download.pytorch.org/whl/cu100"
            ),
        },
    },
    "tensorflow": {
        "2.8": {
            "tensorflow": "2.8.2",
            "install": "python3 -m pip install --no-cache-dir -U tensorflow==2.8.2",
        },
        "2.7": {
            "tensorflow": "2.7.3",
            "install": "python3 -m pip install --no-cache-dir -U tensorflow==2.7.3",
        },
        "2.6": {
            "tensorflow": "2.6.5",
            "install": "python3 -m pip install --no-cache-dir -U tensorflow==2.6.5",
        },
        "2.5": {
            "tensorflow": "2.5.3",
            "install": "python3 -m pip install --no-cache-dir -U tensorflow==2.5.3",
        },
        # need another `nvidia:cuda` docker image, otherwise GPU not working
        "2.4": {
            "tensorflow": "2.4.4",
            "install": "python3 -m pip install --no-cache-dir -U tensorflow==2.4.4",
            # This should be specified as a docker build argument.
            # We keep the information here for reference only.
            "base_docker": "nvidia/cuda:11.0.3-cudnn8-devel-ubuntu20.04",
        },
    },
}


if __name__ == "__main__":
    parser = argparse.ArgumentParser("Choose the framework and version to install")
    parser.add_argument(
        "--framework", help="The framework to install. Should be `torch` or `tensorflow`", type=str, required=True
    )
    parser.add_argument("--version", help="The version of the framework to install.", type=str, required=True)
    args = parser.parse_args()

    info = past_versions_testing[args.framework][args.version]

    os.system(f'echo "export INSTALL_CMD=\'{info["install"]}\'" >> ~/.profile')
    print(f'echo "export INSTALL_CMD=\'{info["install"]}\'" >> ~/.profile')

    cuda = ""
    if args.framework == "pytorch":
        cuda = info["cuda"]
    os.system(f"echo \"export CUDA='{cuda}'\" >> ~/.profile")
    print(f"echo \"export CUDA='{cuda}'\" >> ~/.profile")