import huggingface_hub
from PIL import Image
from pathlib import Path
import csv
import spaces

import onnxruntime as rt
e621_model_path = Path(huggingface_hub.snapshot_download('toynya/Z3D-E621-Convnext'))
e621_model_session = rt.InferenceSession(e621_model_path / 'model.onnx', providers=["CUDAExecutionProvider", "CPUExecutionProvider"])
with open(e621_model_path / 'tags-selected.csv', mode='r', encoding='utf-8') as file:
	csv_reader = csv.DictReader(file)
	e621_model_tags = [row['name'].strip() for row in csv_reader]


def prepare_image_e621(image: Image.Image, target_size: int):
	import numpy as np
	# Pad image to square
	image_shape = image.size
	max_dim = max(image_shape)
	pad_left = (max_dim - image_shape[0]) // 2
	pad_top = (max_dim - image_shape[1]) // 2

	padded_image = Image.new("RGB", (max_dim, max_dim), (255, 255, 255))
	padded_image.paste(image, (pad_left, pad_top))

	# Resize
	if max_dim != target_size:
		padded_image = padded_image.resize((target_size, target_size), Image.BICUBIC)
	
	# Convert to numpy array
	# Based on the ONNX graph, the model appears to expect inputs in the range of 0-255
	image_array = np.asarray(padded_image, dtype=np.float32)

	# Convert PIL-native RGB to BGR
	image_array = image_array[:, :, ::-1]

	return np.expand_dims(image_array, axis=0)


@spaces.GPU
def predict_e621(image: Image.Image, threshold: float = 0.3):
	image_array = prepare_image_e621(image, 448)

	image_array = prepare_image_e621(image, 448)
	input_name = 'input_1:0'
	output_name = 'predictions_sigmoid'

	result = e621_model_session.run([output_name], {input_name: image_array})
	result = result[0][0]

	scores = {e621_model_tags[i]: result[i] for i in range(len(result))}
	predicted_tags = [tag for tag, score in scores.items() if score > threshold]
	tag_string = ', '.join(predicted_tags).replace("_", " ")

	return tag_string


def predict_tags_e621(image: Image.Image, input_tags: str, algo: list[str], threshold: float = 0.3):
    def to_list(s):
        return [x.strip() for x in s.split(",") if not s == ""]
    
    def list_uniq(l):
        return sorted(set(l), key=l.index)
	
    if not "Use Z3D-E621-Convnext" in algo:
        return input_tags
    tag_list = list_uniq(to_list(input_tags) + to_list(predict_e621(image)))
    return ", ".join(tag_list)