import os
from functools import lru_cache

import gradio as gr
import numpy as np
import pandas as pd
from huggingface_hub import HfFileSystem, hf_hub_download
from imgutils.generic import classify_predict_score
from natsort import natsorted

hf_fs = HfFileSystem()

_REPOSITORY = 'deepghs/anime_aesthetic'
_DEFAULT_MODEL = 'swinv2pv3_v0_448_ls0.2_x'
_MODELS = natsorted([
    os.path.dirname(os.path.relpath(file, _REPOSITORY))
    for file in hf_fs.glob(f'{_REPOSITORY}/*/model.onnx')
])

LABELS = ["worst", "low", "normal", "good", "great", "best", "masterpiece"]


@lru_cache()
def _get_mark_table(model):
    df = pd.read_csv(hf_hub_download(
        repo_id=_REPOSITORY,
        repo_type='model',
        filename=f'{model}/samples.csv',
    ))
    df = df.sort_values(['score'])
    df['cnt'] = list(range(len(df)))
    df['final_score'] = df['cnt'] / len(df)

    x = np.concatenate([[0.0], df['score'], [6.0]])
    y = np.concatenate([[0.0], df['final_score'], [1.0]])
    return x, y


def _get_percentile(x, y, v):
    idx = np.searchsorted(x, np.clip(v, a_min=0.0, a_max=6.0))
    if idx < x.shape[0] - 1:
        x0, y0 = x[idx], y[idx]
        x1, y1 = x[idx + 1], y[idx + 1]
        return np.clip((v - x0) / (x1 - x0) * (y1 - y0) + y0, a_min=0.0, a_max=1.0)

    else:
        return y[idx]


def _fn_predict(image, model):
    scores = classify_predict_score(
        image=image,
        repo_id=_REPOSITORY,
        model_name=model,
    )
    weighted_mean = sum(i * scores[label] for i, label in enumerate(LABELS))
    x, y = _get_mark_table(model)
    percentile = _get_percentile(x, y, weighted_mean)
    return weighted_mean, percentile, scores


if __name__ == '__main__':
    with gr.Blocks() as demo:
        with gr.Row():
            with gr.Column():
                gr_input_image = gr.Image(type='pil', label='Original Image')
                gr_model = gr.Dropdown(_MODELS, value=_DEFAULT_MODEL, label='Model')
                gr_submit = gr.Button(value='Submit', variant='primary')

            with gr.Column():
                with gr.Row():
                    gr_score = gr.Text(label='Aesthetic Score (0~6)', value='')
                    gr_percentile = gr.Text(label='Percentile (0.0-1.0)', value='')

                with gr.Row():
                    gr_output = gr.Label(label='Aesthetic Classes')

            gr_submit.click(
                _fn_predict,
                inputs=[gr_input_image, gr_model],
                outputs=[gr_score, gr_percentile, gr_output],
            )

    demo.queue(os.cpu_count()).launch()