#!/usr/bin/env python3
# Copyright (C) 2020 Intel Corporation
# SPDX-License-Identifier: Apache-2.0
#


""" Script to acquire model IRs for stress tests.
Usage: ./scrips/get_testdata.py
"""
# pylint:disable=line-too-long

import argparse
import logging as log
import multiprocessing
import os
import shutil
import subprocess
import sys
from inspect import getsourcefile
from pathlib import Path
from xml.etree import ElementTree as ET

log.basicConfig(format="{file}: [ %(levelname)s ] %(message)s".format(file=os.path.basename(__file__)),
                level=log.INFO, stream=sys.stdout)

# Parameters
OMZ_NUM_ATTEMPTS = 6
DOWNLOADER_JOBS_NUM = 4


def abs_path(relative_path):
    """Return absolute path given path relative to the current file.
    """
    return os.path.realpath(
        os.path.join(os.path.dirname(getsourcefile(lambda: 0)), relative_path))


class VirtualEnv:
    """Class implemented creation and use of virtual environment."""
    is_created = False

    def __init__(self, venv_dir):
        self.venv_dir = Path(abs_path('..')) / venv_dir
        if sys.platform.startswith('linux') or sys.platform == 'darwin':
            self.venv_executable = self.venv_dir / "bin" / "python3"
        else:
            self.venv_executable = self.venv_dir / "Scripts" / "python3.exe"

    def get_venv_executable(self):
        """Returns path to executable from virtual environment."""
        return str(self.venv_executable)

    def get_venv_dir(self):
        """Returns path to virtual environment root directory."""
        return str(self.venv_dir)

    def create(self):
        """Creates virtual environment."""
        cmd = '{executable} -m venv {venv}'.format(executable=sys.executable,
                                                   venv=self.get_venv_dir())
        run_in_subprocess(cmd)
        self.is_created = True

    def install_requirements(self, *requirements):
        """Installs provided requirements. Creates virtual environment if it hasn't been created."""
        if not self.is_created:
            self.create()
        cmd = '{executable} -m pip install --upgrade pip'.format(executable=self.get_venv_executable())
        for req in requirements:
            # Don't install requirements via one `pip install` call to prevent "ERROR: Double requirement given"
            cmd += ' && {executable} -m pip install -r {req}'.format(executable=self.get_venv_executable(), req=req)
        run_in_subprocess(cmd)

    def create_n_install_requirements(self, *requirements):
        """Creates virtual environment and installs provided requirements in it."""
        self.create()
        self.install_requirements(*requirements)


def run_in_subprocess(cmd, check_call=True):
    """Runs provided command in attached subprocess."""
    log.info(cmd)
    if check_call:
        subprocess.check_call(cmd, shell=True)
    else:
        subprocess.call(cmd, shell=True)


def main():
    """Main entry point.
    """
    parser = argparse.ArgumentParser(
        description='Acquire test data',
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    parser.add_argument('--test_conf', required=True, type=Path,
                        help='Path to a test config .xml file containing models '
                             'which will be downloaded and converted to IRs via OMZ.')
    parser.add_argument('--omz_repo', required=False,
                        help='Path to Open Model Zoo (OMZ) repository. It will be used to skip cloning step.')
    parser.add_argument('--mo_tool', type=Path,
                        default=Path(abs_path('../../../model-optimizer/mo.py')).resolve(),
                        help='Path to Model Optimizer (MO) runner. Required for OMZ converter.py only.')
    parser.add_argument('--omz_models_out_dir', type=Path,
                        default=abs_path('../_omz_out/models'),
                        help='Directory to put test data into. Required for OMZ downloader.py and converter.py')
    parser.add_argument('--omz_irs_out_dir', type=Path,
                        default=abs_path('../_omz_out/irs'),
                        help='Directory to put test data into. Required for OMZ converter.py only.')
    parser.add_argument('--omz_cache_dir', type=Path,
                        default=abs_path('../_omz_out/cache'),
                        help='Directory with test data cache. Required for OMZ downloader.py only.')
    parser.add_argument('--no_venv', action="store_true",
                        help='Skip preparation and use of virtual environment to convert models via OMZ converter.py.')
    parser.add_argument('--skip_omz_errors', action="store_true",
                        help='Skip errors caused by OMZ while downloading and converting.')
    args = parser.parse_args()

    # parse models from test config
    tree = ET.parse(str(args.test_conf))
    root = tree.getroot()
    models_from_cfg = []
    for attributes in root:
        if attributes.tag == "models":
            models = [child.text for child in attributes]
            models_from_cfg = [Path(model).stem for model in models]
            break

    # prepare Open Model Zoo
    if args.omz_repo:
        omz_path = Path(args.omz_repo).resolve()
    else:
        omz_path = Path(abs_path('..')) / "_open_model_zoo"
        # clone Open Model Zoo into temporary path
        if os.path.exists(str(omz_path)):
            shutil.rmtree(str(omz_path))
        cmd = 'git clone --single-branch --branch develop' \
              ' https://github.com/opencv/open_model_zoo {omz_path}'.format(omz_path=omz_path)
        run_in_subprocess(cmd)

    # prepare models
    downloader_path = omz_path / "tools" / "downloader" / "downloader.py"
    models_list_path = args.omz_models_out_dir / "models_list.txt"
    cmd = '{downloader_path} --list {models_list_path}' \
          ' --num_attempts {num_attempts}' \
          ' --output_dir {models_dir}' \
          ' --cache_dir {cache_dir}' \
          ' --jobs {jobs_num}'.format(downloader_path=downloader_path, models_list_path=models_list_path,
                                      num_attempts=OMZ_NUM_ATTEMPTS,
                                      models_dir=args.omz_models_out_dir,
                                      cache_dir=args.omz_cache_dir,
                                      jobs_num=DOWNLOADER_JOBS_NUM)
    models_available = subprocess.check_output(cmd + " --print_all", shell=True, universal_newlines=True).split("\n")
    models_to_run = set(models_from_cfg).intersection(models_available)

    os.makedirs(str(args.omz_models_out_dir), exist_ok=True)
    log.info("List of models from {models_list_path} used for downloader.py and converter.py: "
             "{models_to_run}".format(models_list_path=models_list_path, models_to_run=",".join(models_to_run)))
    with open(str(models_list_path), "w") as file:
        file.writelines([name + "\n" for name in models_to_run])
    if set(models_from_cfg) - models_to_run:
        log.warning("List of models defined in config but not available in OMZ: {}"
                    .format(",".join(set(models_from_cfg) - models_to_run)))

    run_in_subprocess(cmd, check_call=not args.skip_omz_errors)

    # prepare virtual environment and install requirements
    python_executable = sys.executable
    if not args.no_venv:
        Venv = VirtualEnv("./.stress_venv")
        requirements = [
            omz_path / "tools" / "downloader" / "requirements.in",
            args.mo_tool.parent / "requirements.txt",
            args.mo_tool.parent / "requirements_dev.txt",
            omz_path / "tools" / "downloader" / "requirements-caffe2.in",
            omz_path / "tools" / "downloader" / "requirements-pytorch.in"
        ]
        Venv.create_n_install_requirements(*requirements)
        python_executable = Venv.get_venv_executable()

    # convert models to IRs
    converter_path = omz_path / "tools" / "downloader" / "converter.py"
    # NOTE: remove --precision if both precisions (FP32 & FP16) required
    cmd = '{executable} {converter_path} --list {models_list_path}' \
          ' -p {executable}' \
          ' --precision=FP32' \
          ' --output_dir {irs_dir}' \
          ' --download_dir {models_dir}' \
          ' --mo {mo_tool} --jobs {workers_num}'.format(executable=python_executable, converter_path=converter_path,
                                                        models_list_path=models_list_path, irs_dir=args.omz_irs_out_dir,
                                                        models_dir=args.omz_models_out_dir, mo_tool=args.mo_tool,
                                                        workers_num=multiprocessing.cpu_count())
    run_in_subprocess(cmd, check_call=not args.skip_omz_errors)


if __name__ == "__main__":
    main()
