# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
#
# 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
import sys
__dir__ = os.path.dirname(__file__)
sys.path.append(os.path.join(__dir__, ""))
sys.path.append(os.path.join(__dir__, "deploy"))

import argparse
import shutil
import textwrap
import tarfile
import requests
import warnings
from functools import partial
from difflib import SequenceMatcher

import cv2
import numpy as np
from tqdm import tqdm
from prettytable import PrettyTable

from deploy.python.predict_cls import ClsPredictor
from deploy.utils.get_image_list import get_image_list
from deploy.utils import config

from ppcls.arch.backbone import *

__all__ = ["PaddleClas"]

BASE_DIR = os.path.expanduser("~/.paddleclas/")
BASE_INFERENCE_MODEL_DIR = os.path.join(BASE_DIR, "inference_model")
BASE_IMAGES_DIR = os.path.join(BASE_DIR, "images")
BASE_DOWNLOAD_URL = "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/inference/{}_infer.tar"
MODEL_SERIES = {
    "AlexNet": ["AlexNet"],
    "DarkNet": ["DarkNet53"],
    "DeiT": [
        "DeiT_base_distilled_patch16_224", "DeiT_base_distilled_patch16_384",
        "DeiT_base_patch16_224", "DeiT_base_patch16_384",
        "DeiT_small_distilled_patch16_224", "DeiT_small_patch16_224",
        "DeiT_tiny_distilled_patch16_224", "DeiT_tiny_patch16_224"
    ],
    "DenseNet": [
        "DenseNet121", "DenseNet161", "DenseNet169", "DenseNet201",
        "DenseNet264"
    ],
    "DPN": ["DPN68", "DPN92", "DPN98", "DPN107", "DPN131"],
    "EfficientNet": [
        "EfficientNetB0", "EfficientNetB0_small", "EfficientNetB1",
        "EfficientNetB2", "EfficientNetB3", "EfficientNetB4", "EfficientNetB5",
        "EfficientNetB6", "EfficientNetB7"
    ],
    "GhostNet":
    ["GhostNet_x0_5", "GhostNet_x1_0", "GhostNet_x1_3", "GhostNet_x1_3_ssld"],
    "HRNet": [
        "HRNet_W18_C", "HRNet_W30_C", "HRNet_W32_C", "HRNet_W40_C",
        "HRNet_W44_C", "HRNet_W48_C", "HRNet_W64_C", "HRNet_W18_C_ssld",
        "HRNet_W48_C_ssld"
    ],
    "Inception": ["GoogLeNet", "InceptionV3", "InceptionV4"],
    "MobileNetV1": [
        "MobileNetV1_x0_25", "MobileNetV1_x0_5", "MobileNetV1_x0_75",
        "MobileNetV1", "MobileNetV1_ssld"
    ],
    "MobileNetV2": [
        "MobileNetV2_x0_25", "MobileNetV2_x0_5", "MobileNetV2_x0_75",
        "MobileNetV2", "MobileNetV2_x1_5", "MobileNetV2_x2_0",
        "MobileNetV2_ssld"
    ],
    "MobileNetV3": [
        "MobileNetV3_small_x0_35", "MobileNetV3_small_x0_5",
        "MobileNetV3_small_x0_75", "MobileNetV3_small_x1_0",
        "MobileNetV3_small_x1_25", "MobileNetV3_large_x0_35",
        "MobileNetV3_large_x0_5", "MobileNetV3_large_x0_75",
        "MobileNetV3_large_x1_0", "MobileNetV3_large_x1_25",
        "MobileNetV3_small_x1_0_ssld", "MobileNetV3_large_x1_0_ssld"
    ],
    "RegNet": ["RegNetX_4GF"],
    "Res2Net": [
        "Res2Net50_14w_8s", "Res2Net50_26w_4s", "Res2Net50_vd_26w_4s",
        "Res2Net200_vd_26w_4s", "Res2Net101_vd_26w_4s",
        "Res2Net50_vd_26w_4s_ssld", "Res2Net101_vd_26w_4s_ssld",
        "Res2Net200_vd_26w_4s_ssld"
    ],
    "ResNeSt": ["ResNeSt50", "ResNeSt50_fast_1s1x64d"],
    "ResNet": [
        "ResNet18", "ResNet18_vd", "ResNet34", "ResNet34_vd", "ResNet50",
        "ResNet50_vc", "ResNet50_vd", "ResNet50_vd_v2", "ResNet101",
        "ResNet101_vd", "ResNet152", "ResNet152_vd", "ResNet200_vd",
        "ResNet34_vd_ssld", "ResNet50_vd_ssld", "ResNet50_vd_ssld_v2",
        "ResNet101_vd_ssld", "Fix_ResNet50_vd_ssld_v2", "ResNet50_ACNet_deploy"
    ],
    "ResNeXt": [
        "ResNeXt50_32x4d", "ResNeXt50_vd_32x4d", "ResNeXt50_64x4d",
        "ResNeXt50_vd_64x4d", "ResNeXt101_32x4d", "ResNeXt101_vd_32x4d",
        "ResNeXt101_32x8d_wsl", "ResNeXt101_32x16d_wsl",
        "ResNeXt101_32x32d_wsl", "ResNeXt101_32x48d_wsl",
        "Fix_ResNeXt101_32x48d_wsl", "ResNeXt101_64x4d", "ResNeXt101_vd_64x4d",
        "ResNeXt152_32x4d", "ResNeXt152_vd_32x4d", "ResNeXt152_64x4d",
        "ResNeXt152_vd_64x4d"
    ],
    "SENet": [
        "SENet154_vd", "SE_HRNet_W64_C_ssld", "SE_ResNet18_vd",
        "SE_ResNet34_vd", "SE_ResNet50_vd", "SE_ResNeXt50_32x4d",
        "SE_ResNeXt50_vd_32x4d", "SE_ResNeXt101_32x4d"
    ],
    "ShuffleNetV2": [
        "ShuffleNetV2_swish", "ShuffleNetV2_x0_25", "ShuffleNetV2_x0_33",
        "ShuffleNetV2_x0_5", "ShuffleNetV2_x1_0", "ShuffleNetV2_x1_5",
        "ShuffleNetV2_x2_0"
    ],
    "SqueezeNet": ["SqueezeNet1_0", "SqueezeNet1_1"],
    "SwinTransformer": [
        "SwinTransformer_large_patch4_window7_224_22kto1k",
        "SwinTransformer_large_patch4_window12_384_22kto1k",
        "SwinTransformer_base_patch4_window7_224_22kto1k",
        "SwinTransformer_base_patch4_window12_384_22kto1k",
        "SwinTransformer_base_patch4_window12_384",
        "SwinTransformer_base_patch4_window7_224",
        "SwinTransformer_small_patch4_window7_224",
        "SwinTransformer_tiny_patch4_window7_224"
    ],
    "VGG": ["VGG11", "VGG13", "VGG16", "VGG19"],
    "VisionTransformer": [
        "ViT_base_patch16_224", "ViT_base_patch16_384", "ViT_base_patch32_384",
        "ViT_large_patch16_224", "ViT_large_patch16_384",
        "ViT_large_patch32_384", "ViT_small_patch16_224"
    ],
    "Xception": [
        "Xception41", "Xception41_deeplab", "Xception65", "Xception65_deeplab",
        "Xception71"
    ]
}


class ImageTypeError(Exception):
    """ImageTypeError.
    """

    def __init__(self, message=""):
        super().__init__(message)


class InputModelError(Exception):
    """InputModelError.
    """

    def __init__(self, message=""):
        super().__init__(message)


def args_cfg():
    parser = config.parser()
    other_options = [
        ("infer_imgs", str, None, "The image(s) to be predicted."),
        ("model_name", str, None, "The model name to be used."),
        ("inference_model_dir", str, None, "The directory of model files."),
        ("use_gpu", bool, True, "Whether use GPU. Default by True."), (
            "enable_mkldnn", bool, False,
            "Whether use MKLDNN. Default by False."),
        ("batch_size", int, 1, "Batch size. Default by 1.")
    ]
    for name, opt_type, default, description in other_options:
        parser.add_argument(
            "--" + name, type=opt_type, default=default, help=description)

    args = parser.parse_args()

    for name, opt_type, default, description in other_options:
        val = eval("args." + name)
        full_name = "Global." + name
        args.override.append(
            f"{full_name}={val}") if val is not default else None

    cfg = config.get_config(
        args.config, overrides=args.override, show=args.verbose)

    return cfg


def get_default_confg():
    return {
        "Global": {
            "model_name": "MobileNetV3_small_x0_35",
            "use_gpu": False,
            "use_fp16": False,
            "enable_mkldnn": False,
            "cpu_num_threads": 1,
            "use_tensorrt": False,
            "ir_optim": False,
            "enable_profile": False
        },
        "PreProcess": {
            "transform_ops": [{
                "ResizeImage": {
                    "resize_short": 256
                }
            }, {
                "CropImage": {
                    "size": 224
                }
            }, {
                "NormalizeImage": {
                    "scale": 0.00392157,
                    "mean": [0.485, 0.456, 0.406],
                    "std": [0.229, 0.224, 0.225],
                    "order": ""
                }
            }, {
                "ToCHWImage": None
            }]
        },
        "PostProcess": {
            "name": "Topk",
            "topk": 5,
            "class_id_map_file": "./ppcls/utils/imagenet1k_label_list.txt"
        }
    }


def print_info():
    """Print list of supported models in formatted.
    """
    table = PrettyTable(["Series", "Name"])
    try:
        sz = os.get_terminal_size()
        width = sz.columns - 30 if sz.columns > 50 else 10
    except OSError:
        width = 100
    for series in MODEL_SERIES:
        names = textwrap.fill("  ".join(MODEL_SERIES[series]), width=width)
        table.add_row([series, names])
    width = len(str(table).split("\n")[0])
    print("{}".format("-" * width))
    print("Models supported by PaddleClas".center(width))
    print(table)
    print("Powered by PaddlePaddle!".rjust(width))
    print("{}".format("-" * width))


def get_model_names():
    """Get the model names list.
    """
    model_names = []
    for series in MODEL_SERIES:
        model_names += (MODEL_SERIES[series])
    return model_names


def similar_architectures(name="", names=[], thresh=0.1, topk=10):
    """Find the most similar topk model names.
    """
    scores = []
    for idx, n in enumerate(names):
        if n.startswith("__"):
            continue
        score = SequenceMatcher(None, n.lower(), name.lower()).quick_ratio()
        if score > thresh:
            scores.append((idx, score))
    scores.sort(key=lambda x: x[1], reverse=True)
    similar_names = [names[s[0]] for s in scores[:min(topk, len(scores))]]
    return similar_names


def download_with_progressbar(url, save_path):
    """Download from url with progressbar.
    """
    if os.path.isfile(save_path):
        os.remove(save_path)
    response = requests.get(url, stream=True)
    total_size_in_bytes = int(response.headers.get("content-length", 0))
    block_size = 1024  # 1 Kibibyte
    progress_bar = tqdm(total=total_size_in_bytes, unit="iB", unit_scale=True)
    with open(save_path, "wb") as file:
        for data in response.iter_content(block_size):
            progress_bar.update(len(data))
            file.write(data)
    progress_bar.close()
    if total_size_in_bytes == 0 or progress_bar.n != total_size_in_bytes or not os.path.isfile(
            save_path):
        raise Exception(
            f"Something went wrong while downloading model/image from {url}")


def check_model_file(model_name):
    """Check the model files exist and download and untar when no exist. 
    """
    storage_directory = partial(os.path.join, BASE_INFERENCE_MODEL_DIR,
                                model_name)
    url = BASE_DOWNLOAD_URL.format(model_name)

    tar_file_name_list = [
        "inference.pdiparams", "inference.pdiparams.info", "inference.pdmodel"
    ]
    model_file_path = storage_directory("inference.pdmodel")
    params_file_path = storage_directory("inference.pdiparams")
    if not os.path.exists(model_file_path) or not os.path.exists(
            params_file_path):
        tmp_path = storage_directory(url.split("/")[-1])
        print(f"download {url} to {tmp_path}")
        os.makedirs(storage_directory(), exist_ok=True)
        download_with_progressbar(url, tmp_path)
        with tarfile.open(tmp_path, "r") as tarObj:
            for member in tarObj.getmembers():
                filename = None
                for tar_file_name in tar_file_name_list:
                    if tar_file_name in member.name:
                        filename = tar_file_name
                if filename is None:
                    continue
                file = tarObj.extractfile(member)
                with open(storage_directory(filename), "wb") as f:
                    f.write(file.read())
        os.remove(tmp_path)
    if not os.path.exists(model_file_path) or not os.path.exists(
            params_file_path):
        raise Exception(
            f"Something went wrong while praparing the model[{model_name}] files!"
        )

    return storage_directory()


def save_prelabel_results(class_id, input_file_path, output_dir):
    """Save the predicted image according to the prediction.
    """
    output_dir = os.path.join(output_dir, str(class_id))
    if not os.path.isdir(output_dir):
        os.makedirs(output_dir)
    shutil.copy(input_file_path, output_dir)


class PaddleClas(object):
    """PaddleClas.
    """

    print_info()

    def __init__(self,
                 config: dict=None,
                 model_name: str=None,
                 inference_model_dir: str=None,
                 use_gpu: bool=None,
                 batch_size: int=None):
        """Init PaddleClas with config.

        Args:
            config: The config of PaddleClas's predictor, default by None. If default, the default configuration is used. Please refer doc for more information.
            model_name: The model name supported by PaddleClas, default by None. If specified, override config.
            inference_model_dir: The directory that contained model file and params file to be used, default by None. If specified, override config.
            use_gpu: Wheather use GPU, default by None. If specified, override config.
            batch_size: The batch size to pridict, default by None. If specified, override config.
        """
        super().__init__()
        self._config = config
        self._check_config(model_name, inference_model_dir, use_gpu,
                           batch_size)
        self._check_input_model()
        self.cls_predictor = ClsPredictor(self._config)

    def get_config(self):
        """Get the config.
        """
        return self._config

    def _check_config(self,
                      model_name=None,
                      inference_model_dir=None,
                      use_gpu=None,
                      batch_size=None):
        if self._config is None:
            self._config = get_default_confg()
            warnings.warn("config is not provided, use default!")
        self._config = config.AttrDict(self._config)
        config.create_attr_dict(self._config)

        if model_name is not None:
            self._config.Global["model_name"] = model_name
        if inference_model_dir is not None:
            self._config.Global["inference_model_dir"] = inference_model_dir
        if use_gpu is not None:
            self._config.Global["use_gpu"] = use_gpu
        if batch_size is not None:
            self._config.Global["batch_size"] = batch_size

    def _check_input_model(self):
        """Check input model name or model files.
        """
        candidate_model_names = get_model_names()
        input_model_name = self._config.Global.get("model_name", None)
        inference_model_dir = self._config.Global.get("inference_model_dir",
                                                      None)
        if input_model_name is not None:
            similar_names = similar_architectures(input_model_name,
                                                  candidate_model_names)
            similar_names_str = ", ".join(similar_names)
            if input_model_name not in similar_names_str:
                err = f"{input_model_name} is not exist! Maybe you want: [{similar_names_str}]"
                raise InputModelError(err)
            if input_model_name not in candidate_model_names:
                err = f"{input_model_name} is not provided by PaddleClas. If you want to use your own model, please input model_file as model path!"
                raise InputModelError(err)
            self._config.Global.inference_model_dir = check_model_file(
                input_model_name)
            return
        elif inference_model_dir is not None:
            model_file_path = os.path.join(inference_model_dir,
                                           "inference.pdmodel")
            params_file_path = os.path.join(inference_model_dir,
                                            "inference.pdiparams")
            if not os.path.isfile(model_file_path) or not os.path.isfile(
                    params_file_path):
                err = f"There is no model file or params file in this directory: {inference_model_dir}"
                raise InputModelError(err)
            return
        else:
            err = f"Please specify the model name supported by PaddleClas or directory contained model file and params file."
            raise InputModelError(err)
        return

    def predict(self, input_data, print_pred=True):
        """Predict label of img with paddleclas.
        Args:
            input_data(str, NumPy.ndarray): 
                image to be classified, support: str(local path of image file, internet URL, directory containing series of images) and NumPy.ndarray(preprocessed image data that has 3 channels and accords with [C, H, W], or raw image data that has 3 channels and accords with [H, W, C]).
        Returns:
            dict: {image_name: "", class_id: [], scores: [], label_names: []}，if label name path == None，label_names will be empty.
        """
        if isinstance(input_data, np.ndarray):
            return self.cls_predictor.predict(input_data)
        elif isinstance(input_data, str):
            if input_data.startswith("http"):
                image_storage_dir = partial(os.path.join, BASE_IMAGES_DIR)
                if not os.path.exists(image_storage_dir()):
                    os.makedirs(image_storage_dir())
                image_save_path = image_storage_dir("tmp.jpg")
                download_with_progressbar(input_data, image_save_path)
                input_data = image_save_path
                warnings.warn(
                    f"Image to be predicted from Internet: {input_data}, has been saved to: {image_save_path}"
                )
            image_list = get_image_list(input_data)

            batch_size = self._config.Global.get("batch_size", 1)
            pre_label_out_idr = self._config.Global.get("pre_label_out_idr",
                                                        False)

            img_list = []
            img_path_list = []
            output_list = []
            cnt = 0
            for idx, img_path in enumerate(image_list):
                img = cv2.imread(img_path)
                if img is None:
                    warnings.warn(
                        f"Image file failed to read and has been skipped. The path: {img_path}"
                    )
                    continue
                img_list.append(img)
                img_path_list.append(img_path)
                cnt += 1

                if cnt % batch_size == 0 or (idx + 1) == len(image_list):
                    outputs = self.cls_predictor.predict(img_list)
                    output_list.append(outputs[0])
                    preds = self.cls_predictor.postprocess(outputs)
                    for nu, pred in enumerate(preds):
                        if pre_label_out_idr:
                            save_prelabel_results(pred["class_ids"][0],
                                                  img_path_list[nu],
                                                  pre_label_out_idr)
                        if print_pred:
                            pred_str_list = [
                                f"filename: {img_path_list[nu]}",
                                f"top-{self._config.PostProcess.get('topk', 1)}"
                            ]
                            for k in pred:
                                pred_str_list.append(f"{k}: {pred[k]}")
                            print(", ".join(pred_str_list))
                    img_list = []
                    img_path_list = []
            return output_list
        else:
            err = "Please input legal image! The type of image supported by PaddleClas are: NumPy.ndarray and string of local path or Ineternet URL"
            raise ImageTypeError(err)
        return


# for CLI
def main():
    """Function API used for commad line.
    """
    cfg = args_cfg()
    clas_engine = PaddleClas(cfg)
    clas_engine.predict(cfg["Global"]["infer_imgs"], print_pred=True)
    return


if __name__ == "__main__":
    main()