import os
import time
import math
import argparse
import numpy as np
from PIL import Image, ImageDraw

import mlx.core as mx
from mlxDeepDanBooru import *

parser = argparse.ArgumentParser(description='Thanks for choosing mlxDeepDanBooruModel.')
parser.add_argument("path", type=str, help='image path or image-dir path')
args = parser.parse_args()

IMAGEDIR = f'example'

in_path = args.path
user_file = ""
if os.path.isdir(in_path):
    IMAGEDIR = in_path
    user_file = ""

if os.path.isfile(in_path):
    user_file = in_path
    IMAGEDIR = ""

print(f"user_file: {user_file}")
print(f"IMAGEDIR: {IMAGEDIR}")
#
mlx_dan = mlxDeepDanBooruModel()
# use float32
mlx_dan.load_weights(f"models/model-resnet_custom_v3_fp32_mlx.safetensors")
# or use float16
#mlx_dan.load_weights(f"models/model-resnet_custom_v3_fp16_mlx.safetensors")

mx.eval(mlx_dan.parameters())

# load tag list
model_tags = np.load(f'models/tags-resnet_custom_v3_mlx.npy')
#print(f'total tags: {len(model_tags)}')


def danbooru_tags(fpath):
    results = {}
    tags = []

    try:
        pic = img_resize(fpath)
        a = np.expand_dims(np.array(pic, dtype=np.float32), 0) / 255
        # if using fp16 model,
        # "dtype=np.float16" is better
        # a = np.expand_dims(np.array(pic, dtype=np.float16), 0) / 255

        x = mx.array(a)
        y = mlx_dan(x)[0]
        ylst = y.tolist()
        ylen = len(ylst)
        for i in range(ylen):
            if ylst[i] >= 0.72:
                tags.append(model_tags[i].item())

    except Exception as err:
        print(err)
        tags = []

    results[fpath] = tags
    return results

def image_infer(fpath):
    tags = danbooru_tags(fpath)
    return tags


if __name__ == "__main__":
    image_count = 0
    t1 = time.time()

    if user_file != "":
        tags = image_infer(user_file)
        image_count += 1
        print(tags)

    if IMAGEDIR != "":
        for root, _, files in os.walk(IMAGEDIR,True):
            for f in files:
                if not f[-4:].lower() in (".png",".jpg","jpeg",".bmp"):
                    continue
                fpath = os.path.join(root,f).replace("\\","/")
                tags = image_infer(fpath)
                image_count += 1
                print(tags)

    t2 = time.time()

    if image_count > 0:
        print(f'\n\n{image_count} images: infer speed(with mlx): {(t2 - t1)/image_count} seconds per image')