# Author: Hong Xu. This file is under CC0.

def pip_install_with_confirmation(packages):
    """Run ``pip install`` with a confirmation only if some packages are missing. This is useful to put in the beginning of
    a Jupyter notebook.

    Args:

        packages (dict): Each key is the name of a package to be installed. Each value is a sequence of size 3. The
        first two elements are the ``package`` and ``name`` parameter of ``importlib.import_module()``, respectively.
        The last element is the imported name. Roughly speaking, the three elements correspond to the following three
        spots in an import statement::

            from {param 0} import {param 1} as {param 2}

        Example::

            pip_install_with_confirmation({
                'requests': [None, 'requests', 'requests'],
                'numpy': [None, 'numpy', 'np'],
                'matplotlib': ['matplotlib', '.pyplot', 'plt']
            })

    """
    def import_all_modules():
        from importlib import import_module, invalidate_caches
        for package, info in packages.items():
            globals()[info[2]] = import_module(package=info[0], name=info[1])
            invalidate_caches()

    try:
        ModuleNotFoundError()
    except NameError:
        ModuleNotFoundError = ImportError

    try:
        import_all_modules()
    except ModuleNotFoundError:
        import sys
        pip_install_command = [sys.executable, '-m', 'pip', '-q', 'install'] + list(packages)
        answer = input('Will run "{}". Confirm? [y/n]'.format(' '.join(pip_install_command)))
        if answer == 'y':
            import subprocess
            subprocess.check_call(pip_install_command)
            print("Package installation succeeded!")
            import_all_modules()