#!/usr/bin/env python3

#  Copyright (c) 2021 Arm Limited. All rights reserved.
#  SPDX-License-Identifier: Apache-2.0
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

import os, errno
import urllib.request
import subprocess
import fnmatch
import logging
import sys

from argparse import ArgumentParser, ArgumentTypeError
from urllib.error import URLError
from collections import namedtuple


json_uc_res = [{
    "use_case_name": "ad",
    "resources": [{"name": "ad_medium_int8.tflite",
                   "url": "https://github.com/ARM-software/ML-zoo/raw/7c32b097f7d94aae2cd0b98a8ed5a3ba81e66b18/models/anomaly_detection/micronet_medium/tflite_int8/ad_medium_int8.tflite"},
                  {"name": "ifm0.npy",
                   "url": "https://github.com/ARM-software/ML-zoo/raw/7c32b097f7d94aae2cd0b98a8ed5a3ba81e66b18/models/anomaly_detection/micronet_medium/tflite_int8/testing_input/input/0.npy"},
                  {"name": "ofm0.npy",
                   "url": "https://github.com/ARM-software/ML-zoo/raw/7c32b097f7d94aae2cd0b98a8ed5a3ba81e66b18/models/anomaly_detection/micronet_medium/tflite_int8/testing_output/Identity/0.npy"}]
},
    {
        "use_case_name": "asr",
        "resources": [{"name": "wav2letter_pruned_int8.tflite",
                       "url": "https://github.com/ARM-software/ML-zoo/raw/1a92aa08c0de49a7304e0a7f3f59df6f4fd33ac8/models/speech_recognition/wav2letter/tflite_pruned_int8/wav2letter_pruned_int8.tflite"},
                      {"name": "ifm0.npy",
                       "url": "https://github.com/ARM-software/ML-zoo/raw/1a92aa08c0de49a7304e0a7f3f59df6f4fd33ac8/models/speech_recognition/wav2letter/tflite_pruned_int8/testing_input/input_2_int8/0.npy"},
                      {"name": "ofm0.npy",
                       "url": "https://github.com/ARM-software/ML-zoo/raw/1a92aa08c0de49a7304e0a7f3f59df6f4fd33ac8/models/speech_recognition/wav2letter/tflite_pruned_int8/testing_output/Identity_int8/0.npy"}]
    },
    {
        "use_case_name": "img_class",
        "resources": [{"name": "mobilenet_v2_1.0_224_INT8.tflite",
                       "url": "https://github.com/ARM-software/ML-zoo/raw/e0aa361b03c738047b9147d1a50e3f2dcb13dbcb/models/image_classification/mobilenet_v2_1.0_224/tflite_int8/mobilenet_v2_1.0_224_INT8.tflite"},
                      {"name": "ifm0.npy",
                       "url": "https://github.com/ARM-software/ML-zoo/raw/e0aa361b03c738047b9147d1a50e3f2dcb13dbcb/models/image_classification/mobilenet_v2_1.0_224/tflite_int8/testing_input/tfl.quantize/0.npy"},
                      {"name": "ofm0.npy",
                       "url": "https://github.com/ARM-software/ML-zoo/raw/e0aa361b03c738047b9147d1a50e3f2dcb13dbcb/models/image_classification/mobilenet_v2_1.0_224/tflite_int8/testing_output/MobilenetV2/Predictions/Reshape_11/0.npy"}]
    },
    {
        "use_case_name": "kws",
        "resources": [{"name": "ds_cnn_clustered_int8.tflite",
                       "url": "https://github.com/ARM-software/ML-zoo/raw/68b5fbc77ed28e67b2efc915997ea4477c1d9d5b/models/keyword_spotting/ds_cnn_large/tflite_clustered_int8/ds_cnn_clustered_int8.tflite"},
                      {"name": "ifm0.npy",
                       "url": "https://github.com/ARM-software/ML-zoo/raw/68b5fbc77ed28e67b2efc915997ea4477c1d9d5b/models/keyword_spotting/ds_cnn_large/tflite_clustered_int8/testing_input/input_2/0.npy"},
                      {"name": "ofm0.npy",
                       "url": "https://github.com/ARM-software/ML-zoo/raw/68b5fbc77ed28e67b2efc915997ea4477c1d9d5b/models/keyword_spotting/ds_cnn_large/tflite_clustered_int8/testing_output/Identity/0.npy"}]
    },
     {
        "use_case_name": "vww",
        "resources": [{"name": "vww4_128_128_INT8.tflite",
                       "url": "https://github.com/ARM-software/ML-zoo/raw/7dd3b16bb84007daf88be8648983c07f3eb21140/models/visual_wake_words/micronet_vww4/tflite_int8/vww4_128_128_INT8.tflite"},
                      {"name": "ifm0.npy",
                       "url": "https://github.com/ARM-software/ML-zoo/raw/7dd3b16bb84007daf88be8648983c07f3eb21140/models/visual_wake_words/micronet_vww4/tflite_int8/testing_input/input/0.npy"},
                      {"name": "ofm0.npy",
                       "url": "https://github.com/ARM-software/ML-zoo/raw/7dd3b16bb84007daf88be8648983c07f3eb21140/models/visual_wake_words/micronet_vww4/tflite_int8/testing_output/Identity/0.npy"}]
    },
    {
        "use_case_name": "kws_asr",
        "resources": [{"name": "wav2letter_pruned_int8.tflite",
                       "url": "https://github.com/ARM-software/ML-zoo/raw/1a92aa08c0de49a7304e0a7f3f59df6f4fd33ac8/models/speech_recognition/wav2letter/tflite_pruned_int8/wav2letter_pruned_int8.tflite"},
                      {"sub_folder": "asr", "name": "ifm0.npy",
                       "url": "https://github.com/ARM-software/ML-zoo/raw/1a92aa08c0de49a7304e0a7f3f59df6f4fd33ac8/models/speech_recognition/wav2letter/tflite_pruned_int8/testing_input/input_2_int8/0.npy"},
                      {"sub_folder": "asr", "name": "ofm0.npy",
                       "url": "https://github.com/ARM-software/ML-zoo/raw/1a92aa08c0de49a7304e0a7f3f59df6f4fd33ac8/models/speech_recognition/wav2letter/tflite_pruned_int8/testing_output/Identity_int8/0.npy"},
                      {"name": "ds_cnn_clustered_int8.tflite",
                       "url": "https://github.com/ARM-software/ML-zoo/raw/68b5fbc77ed28e67b2efc915997ea4477c1d9d5b/models/keyword_spotting/ds_cnn_large/tflite_clustered_int8/ds_cnn_clustered_int8.tflite"},
                      {"sub_folder": "kws", "name": "ifm0.npy",
                       "url": "https://github.com/ARM-software/ML-zoo/raw/68b5fbc77ed28e67b2efc915997ea4477c1d9d5b/models/keyword_spotting/ds_cnn_large/tflite_clustered_int8/testing_input/input_2/0.npy"},
                      {"sub_folder": "kws", "name": "ofm0.npy",
                       "url": "https://github.com/ARM-software/ML-zoo/raw/68b5fbc77ed28e67b2efc915997ea4477c1d9d5b/models/keyword_spotting/ds_cnn_large/tflite_clustered_int8/testing_output/Identity/0.npy"}]
    },
    {
        "use_case_name": "noise_reduction",
        "resources": [{"name": "rnnoise_INT8.tflite",
                       "url": "https://github.com/ARM-software/ML-zoo/raw/a061600058097a2785d6f1f7785e5a2d2a142955/models/noise_suppression/RNNoise/tflite_int8/rnnoise_INT8.tflite"},
                      {"name": "ifm0.npy",
                       "url": "https://github.com/ARM-software/ML-zoo/raw/a061600058097a2785d6f1f7785e5a2d2a142955/models/noise_suppression/RNNoise/tflite_int8/testing_input/main_input_int8/0.npy"},
                      {"name": "ifm1.npy",
                       "url": "https://github.com/ARM-software/ML-zoo/raw/a061600058097a2785d6f1f7785e5a2d2a142955/models/noise_suppression/RNNoise/tflite_int8/testing_input/vad_gru_prev_state_int8/0.npy"},
                      {"name": "ifm2.npy",
                       "url": "https://github.com/ARM-software/ML-zoo/raw/a061600058097a2785d6f1f7785e5a2d2a142955/models/noise_suppression/RNNoise/tflite_int8/testing_input/noise_gru_prev_state_int8/0.npy"},
                      {"name": "ifm3.npy",
                       "url": "https://github.com/ARM-software/ML-zoo/raw/a061600058097a2785d6f1f7785e5a2d2a142955/models/noise_suppression/RNNoise/tflite_int8/testing_input/denoise_gru_prev_state_int8/0.npy"},
                      {"name": "ofm0.npy",
                       "url": "https://github.com/ARM-software/ML-zoo/raw/a061600058097a2785d6f1f7785e5a2d2a142955/models/noise_suppression/RNNoise/tflite_int8/testing_output/Identity_int8/0.npy"},
                      {"name": "ofm1.npy",
                       "url": "https://github.com/ARM-software/ML-zoo/raw/a061600058097a2785d6f1f7785e5a2d2a142955/models/noise_suppression/RNNoise/tflite_int8/testing_output/Identity_1_int8/0.npy"},
                      {"name": "ofm2.npy",
                       "url": "https://github.com/ARM-software/ML-zoo/raw/a061600058097a2785d6f1f7785e5a2d2a142955/models/noise_suppression/RNNoise/tflite_int8/testing_output/Identity_2_int8/0.npy"},
                      {"name": "ofm3.npy",
                       "url": "https://github.com/ARM-software/ML-zoo/raw/a061600058097a2785d6f1f7785e5a2d2a142955/models/noise_suppression/RNNoise/tflite_int8/testing_output/Identity_3_int8/0.npy"},
                      {"name": "ofm4.npy",
                       "url": "https://github.com/ARM-software/ML-zoo/raw/a061600058097a2785d6f1f7785e5a2d2a142955/models/noise_suppression/RNNoise/tflite_int8/testing_output/Identity_4_int8/0.npy"},
                     ]
    },
    {
        "use_case_name": "inference_runner",
        "resources": [{"name": "dnn_s_quantized.tflite",
                       "url": "https://github.com/ARM-software/ML-zoo/raw/68b5fbc77ed28e67b2efc915997ea4477c1d9d5b/models/keyword_spotting/dnn_small/tflite_int8/dnn_s_quantized.tflite"}
                      ]
    },]

# Valid NPU configurations:
valid_npu_config_names = [
        'ethos-u55-32', 'ethos-u55-64',
        'ethos-u55-128', 'ethos-u55-256',
        'ethos-u65-256','ethos-u65-512']

# Default NPU configurations (these are always run when the models are optimised)
default_npu_config_names = [valid_npu_config_names[2], valid_npu_config_names[4]]

# NPU config named tuple
NPUConfig = namedtuple('NPUConfig',['config_name',
                                    'memory_mode',
                                    'system_config',
                                    'ethos_u_npu_id',
                                    'ethos_u_config_id',
                                    'arena_cache_size'])

# The internal SRAM size for Corstone-300 implementation on MPS3 specified by AN552
mps3_max_sram_sz = 2 * 1024 * 1024 # 2 MiB (2 banks of 1 MiB each)


def call_command(command: str) -> str:
    """
    Helpers function that call subprocess and return the output.

    Parameters:
    ----------
    command (string):  Specifies the command to run.
    """
    logging.info(command)
    proc = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, shell=True)
    log = proc.stdout.decode("utf-8")
    if proc.returncode == 0:
        logging.info(log)
    else:
        logging.error(log)
    proc.check_returncode()
    return log


def get_default_npu_config_from_name(config_name: str, arena_cache_size: int = 0) -> NPUConfig:
    """
    Gets the file suffix for the tflite file from the
    `accelerator_config` string.

    Parameters:
    ----------
    config_name (str):      Ethos-U NPU configuration from valid_npu_config_names

    arena_cache_size (int): Specifies arena cache size in bytes. If a value
                            greater than 0 is provided, this will be taken
                            as the cache size. If 0, the default values, as per
                            the NPU config requirements, are used.

    Returns:
    -------
    NPUConfig: An NPU config named tuple populated with defaults for the given
               config name
    """
    if config_name not in valid_npu_config_names:
        raise ValueError(f"""
            Invalid Ethos-U NPU configuration.
            Select one from {valid_npu_config_names}.
            """)

    strings_ids = ["ethos-u55-", "ethos-u65-"]
    processor_ids = ["U55", "U65"]
    prefix_ids = ["H", "Y"]
    memory_modes = ["Shared_Sram", "Dedicated_Sram"]
    system_configs = ["Ethos_U55_High_End_Embedded", "Ethos_U65_High_End"]
    memory_modes_arena = {
        # For shared SRAM memory mode, we use the MPS3 SRAM size by default
        "Shared_Sram"    : mps3_max_sram_sz if arena_cache_size <= 0 else arena_cache_size,
        # For dedicated SRAM memory mode, we do no override the arena size. This is expected to
        # be defined in the vela configuration file instead.
        "Dedicated_Sram" : None if arena_cache_size <= 0 else arena_cache_size
    }

    for i in range(len(strings_ids)):
        if config_name.startswith(strings_ids[i]):
            npu_config_id = config_name.replace(strings_ids[i], prefix_ids[i])
            return  NPUConfig(config_name=config_name,
                              memory_mode=memory_modes[i],
                              system_config=system_configs[i],
                              ethos_u_npu_id=processor_ids[i],
                              ethos_u_config_id=npu_config_id,
                              arena_cache_size=memory_modes_arena[memory_modes[i]])

    return None


def set_up_resources(run_vela_on_models: bool = False,
                     additional_npu_config_names: list = [],
                     arena_cache_size: int = 0):
    """
    Helpers function that retrieve the output from a command.

    Parameters:
    ----------
    run_vela_on_models (bool):  Specifies if run vela on downloaded models.
    additional_npu_config_names(list):  list of strings of Ethos-U NPU configs.
    arena_cache_size (int): Specifies arena cache size in bytes. If a value
                            greater than 0 is provided, this will be taken
                            as the cache size. If 0, the default values, as per
                            the NPU config requirements, are used.
    """
    current_file_dir = os.path.dirname(os.path.abspath(__file__))
    download_dir = os.path.abspath(os.path.join(current_file_dir, "resources_downloaded"))

    try:
        #   1.1 Does the download dir exist?
        os.mkdir(download_dir)
    except OSError as e:
        if e.errno == errno.EEXIST:
            logging.info("'resources_downloaded' directory exists.")
        else:
            raise

    # 1.2 Does the virtual environment exist?
    env_python = str(os.path.abspath(os.path.join(download_dir, "env", "bin", "python3")))
    env_activate = str(os.path.abspath(os.path.join(download_dir, "env", "bin", "activate")))
    if not os.path.isdir(os.path.join(download_dir, "env")):
        os.chdir(download_dir)
        # Create the virtual environment
        command = "python3 -m venv env"
        call_command(command)
        commands = ["pip install --upgrade pip", "pip install --upgrade setuptools"]
        for c in commands:
            command = f"{env_python} -m {c}"
            call_command(command)
        os.chdir(current_file_dir)
    # 1.3 Make sure to have all the requirement
    requirements = ["ethos-u-vela==3.2.0"]
    command = f"{env_python} -m pip freeze"
    packages = call_command(command)
    for req in requirements:
        if req not in packages:
            command = f"{env_python} -m pip install {req}"
            call_command(command)

    # 2. Download models
    for uc in json_uc_res:
        try:
            #  Does the usecase_name download dir exist?
            os.mkdir(os.path.join(download_dir, uc["use_case_name"]))
        except OSError as e:
            if e.errno != errno.EEXIST:
                logging.error(f"Error creating {uc['use_case_name']} directory.")
                raise

        for res in uc["resources"]:
            res_name = res["name"]
            res_url = res["url"]
            if "sub_folder" in res:
                try:
                    #  Does the usecase_name/sub_folder download dir exist?
                    os.mkdir(os.path.join(download_dir, uc["use_case_name"], res["sub_folder"]))
                except OSError as e:
                    if e.errno != errno.EEXIST:
                        logging.error(f"Error creating {uc['use_case_name']} / {res['sub_folder']} directory.")
                        raise
                res_dst = os.path.join(download_dir,
                                       uc["use_case_name"],
                                       res["sub_folder"],
                                       res_name)
            else:
                res_dst = os.path.join(download_dir,
                                       uc["use_case_name"],
                                       res_name)

            if os.path.isfile(res_dst):
                logging.info(f"File {res_dst} exists, skipping download.")
            else:
                try:
                    g = urllib.request.urlopen(res_url)
                    with open(res_dst, 'b+w') as f:
                        f.write(g.read())
                        logging.info(f"- Downloaded {res_url} to {res_dst}.")
                except URLError:
                    logging.error(f"URLError while downloading {res_url}.")
                    raise

    # 3. Run vela on models in resources_downloaded
    # New models will have same name with '_vela' appended.
    # For example:
    #   original model:  ds_cnn_clustered_int8.tflite
    # after vela model:  ds_cnn_clustered_int8_vela_H128.tflite
    #
    # Note: To avoid to run vela twice on the same model, it's supposed that
    # downloaded model names don't contain the 'vela' word.
    if run_vela_on_models is True:
        config_file = os.path.join(current_file_dir, "scripts", "vela", "default_vela.ini")
        models = [os.path.join(dirpath, f)
                  for dirpath, dirnames, files in os.walk(download_dir)
                  for f in fnmatch.filter(files, '*.tflite') if "vela" not in f]

        # Consolidate all config names while discarding duplicates:
        config_names = list(set(default_npu_config_names + additional_npu_config_names))

        # Get npu config tuple for each config name in a list:
        npu_configs = [get_default_npu_config_from_name(name, arena_cache_size) for name in config_names]

        logging.info(f'All models will be optimised for these configs:')
        for config in npu_configs:
            logging.info(config)

        optimisation_skipped = False

        for model in models:
            output_dir = os.path.dirname(model)
            # model name after compiling with vela is an initial model name + _vela suffix
            vela_optimised_model_path = str(model).replace(".tflite", "_vela.tflite")

            for config in npu_configs:
                vela_command_arena_cache_size = ""

                if config.arena_cache_size:
                    vela_command_arena_cache_size = f"--arena-cache-size={config.arena_cache_size}"

                vela_command = (f". {env_activate} && vela {model} " +
                       f"--accelerator-config={config.config_name} " +
                       "--optimise Performance " +
                       f"--config {config_file} " +
                       f"--memory-mode={config.memory_mode} " +
                       f"--system-config={config.system_config} " +
                       f"--output-dir={output_dir} " +
                       f"{vela_command_arena_cache_size}")

                # we want the name to include the configuration suffix. For example: vela_H128,
                # vela_Y512 etc.
                new_suffix = "_vela_" + config.ethos_u_config_id + '.tflite'
                new_vela_optimised_model_path = (
                    vela_optimised_model_path.replace("_vela.tflite", new_suffix))

                if os.path.isfile(new_vela_optimised_model_path):
                    logging.info(f"File {new_vela_optimised_model_path} exists, skipping optimisation.")
                    optimisation_skipped = True
                    continue

                call_command(vela_command)

                # rename default vela model
                os.rename(vela_optimised_model_path, new_vela_optimised_model_path)
                logging.info(f"Renaming {vela_optimised_model_path} to {new_vela_optimised_model_path}.")


        # If any optimisation was skipped, show how to regenerate:
        if optimisation_skipped:
            logging.warning("One or more optimisations were skipped.")
            logging.warning(f"To optimise all the models, please remove the directory {download_dir}.")


if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument("--skip-vela",
                        help="Do not run Vela optimizer on downloaded models.",
                        action="store_true")
    parser.add_argument("--additional-ethos-u-config-name",
                        help=f"""Additional (non-default) configurations for Vela:
                        {valid_npu_config_names}""",
                        default=[], action="append")
    parser.add_argument("--arena-cache-size",
                        help="Arena cache size in bytes (if overriding the defaults)",
                        type=int,
                        default=0)
    args = parser.parse_args()

    if args.arena_cache_size < 0:
        raise ArgumentTypeError('Arena cache size cannot not be less than 0')

    logging.basicConfig(filename='log_build_default.log', level=logging.DEBUG)
    logging.getLogger().addHandler(logging.StreamHandler(sys.stdout))

    set_up_resources(not args.skip_vela,
                     args.additional_ethos_u_config_name,
                     args.arena_cache_size)