# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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 logging
import pathlib
from pathlib import Path
from typing import Union
import os
from os import environ
import json

import tvm
import tvm.relay as relay
from tvm.contrib import utils, ndk, graph_executor as runtime
from tvm.contrib.download import download_testdata, download

target = "llvm -mtriple=arm64-linux-android"
target_host = None


def del_dir(target: Union[Path, str], only_if_empty: bool = False):
    target = Path(target).expanduser()
    assert target.is_dir()
    for p in sorted(target.glob("**/*"), reverse=True):
        if not p.exists():
            continue
        p.chmod(0o666)
        if p.is_dir():
            p.rmdir()
        else:
            if only_if_empty:
                raise RuntimeError(f"{p.parent} is not empty!")
            p.unlink()
    target.rmdir()


def get_model(model_name, batch_size=1):
    if model_name == "resnet18_v1":
        import mxnet as mx
        from mxnet import gluon
        from mxnet.gluon.model_zoo import vision

        gluon_model = vision.get_model(model_name, pretrained=True)
        img_size = 224
        data_shape = (batch_size, 3, img_size, img_size)
        net, params = relay.frontend.from_mxnet(gluon_model, {"data": data_shape})
        return (net, params)
    elif model_name == "mobilenet_v2":
        import keras
        from keras.applications.mobilenet_v2 import MobileNetV2

        keras.backend.clear_session()  # Destroys the current TF graph and creates a new one.
        weights_url = "".join(
            [
                "https://github.com/JonathanCMitchell/",
                "mobilenet_v2_keras/releases/download/v1.1/",
                "mobilenet_v2_weights_tf_dim_ordering_tf_kernels_0.5_224.h5",
            ]
        )
        weights_file = "mobilenet_v2_weights.h5"
        weights_path = download_testdata(weights_url, weights_file, module="keras")
        keras_mobilenet_v2 = MobileNetV2(
            alpha=0.5, include_top=True, weights=None, input_shape=(224, 224, 3), classes=1000
        )
        keras_mobilenet_v2.load_weights(weights_path)

        img_size = 224
        data_shape = (batch_size, 3, img_size, img_size)
        mod, params = relay.frontend.from_keras(keras_mobilenet_v2, {"input_1": data_shape})
        return (mod, params)


def main(model_str, output_path):
    if output_path.exists():
        del_dir(output_path)
    output_path.mkdir()
    output_path_str = os.fspath(output_path)
    print(model_str)
    print("getting model...")
    net, params = get_model(model_str)
    try:
        os.mkdir(model_str)
    except FileExistsError:
        pass
    print("building...")
    with tvm.transform.PassContext(opt_level=3):
        graph, lib, params = relay.build(net, tvm.target.Target(target, target_host), params=params)
    print("dumping lib...")
    lib.export_library(output_path_str + "/" + "deploy_lib_cpu.so", ndk.create_shared)
    print("dumping graph...")
    with open(output_path_str + "/" + "deploy_graph.json", "w") as f:
        f.write(graph)
    print("dumping params...")
    with open(output_path_str + "/" + "deploy_param.params", "wb") as f:
        f.write(runtime.save_param_dict(params))
    print("dumping labels...")
    synset_url = "".join(
        [
            "https://gist.githubusercontent.com/zhreshold/",
            "4d0b62f3d01426887599d4f7ede23ee5/raw/",
            "596b27d23537e5a1b5751d2b0481ef172f58b539/",
            "imagenet1000_clsid_to_human.txt",
        ]
    )
    synset_path = output_path_str + "/image_net_labels"
    download(synset_url, output_path_str + "/image_net_labels")
    with open(synset_path) as fi:
        synset = eval(fi.read())
        with open(output_path_str + "/image_net_labels.json", "w") as fo:
            json.dump(synset, fo, indent=4)
    os.remove(synset_path)


if __name__ == "__main__":
    if environ.get("TVM_NDK_CC") is None:
        raise RuntimeError("Require environment variable TVM_NDK_CC")
    models_path = Path().absolute().parent.joinpath("app/src/main/assets/models/")
    if not models_path.exists():
        models_path.mkdir()
    models = {
        "mobilenet_v2": models_path.joinpath("mobilenet_v2"),
        "resnet18_v1": models_path.joinpath("resnet18_v1"),
    }
    for model, output_path in models.items():
        main(model, output_path)