import plotly.express as px
import plotly.graph_objects as go

from utils import DEEPLITE_LIGHT_BLUE_HEX, load_yolobench_data


df, pareto_indices = load_yolobench_data()


METRIC_NAME_MAPPING = {
    'mAP@0.5': 'mAP_0.5',
    'mAP@0.5:0.95': 'mAP_0.5:0.95',
    'Precision': 'precision',
    'Recall': 'recall',
}

METRIC_KEYS_TO_NAMES = {v: k for k, v in METRIC_NAME_MAPPING.items()}


LATENCY_KEYS = {
    'Raspberry Pi 4 Model B (CPU, TFLite, FP32)': 'raspi4_tflite_latency',
    'Jetson Nano (GPU, ONNX Runtime, FP32)': 'nano_gpu_latency',
    'Intel® Core™i7-10875H (CPU, OpenVINO, FP32)': 'openvino_latency',
    'Khadas VIM3 (NPU, INT16)': 'vim3_latency',
    'Orange Pi 5 (NPU, FP16)': 'orange_pi_latency',
    'NVIDIA A40 (GPU, TensorRT, FP32)': 'a40_trt_latency',
}

LATENCY_KEYS_TO_NAMES = {v: k for k, v in LATENCY_KEYS.items()}

DATASET_TAGS = {
    'PASCAL VOC': 'voc',
    'SKU-110K': 'sku',
    'WIDERFACE': 'wider',
    'COCO': 'coco',
}

DATASET_TAGS_TO_NAMES = {v: k for k, v in DATASET_TAGS.items()}


def get_scatter_plot(
        dataset_tag,
        metric_tag,
        latency_key,
        model_family_coloring=True,
        add_pareto_frontier=False,
        plot_pareto_only=False,
        log_axis=False,
    ):
    fig_opts, layout_opts = {'opacity': 0.5, 'color_discrete_sequence': [DEEPLITE_LIGHT_BLUE_HEX]}, {}
    if model_family_coloring:
        fig_opts = {
            'color': 'model_family',
            'opacity': 0.75,
            'color_discrete_sequence': px.colors.qualitative.Plotly,
        }
        layout_opts = {
            'legend': dict(
                title='Model family<br>(click to toggle)',
            )
        }

    frontier = None
    if plot_pareto_only:
        metric_key = f'{metric_tag}_{dataset_tag}'
        frontier = pareto_indices[metric_key][latency_key]

    fig = px.scatter(
        df if frontier is None else df.iloc[frontier, :],
        x=latency_key,
        y=f'{metric_tag}_{dataset_tag}',
        title=f'{METRIC_KEYS_TO_NAMES[metric_tag]}-latency scatter plot',
        hover_data={
            'model_name': True,
            'model_family': False,
            latency_key: ':.2f',
            f'{metric_tag}_{dataset_tag}': ':.2f',
        },
        labels={
            'model_name': 'Model name',
            latency_key: 'Latency',
            f'{metric_tag}_{dataset_tag}': METRIC_KEYS_TO_NAMES[metric_tag],
        },
        template='plotly_white',
        **fig_opts,
    )
    if log_axis:
        fig.update_xaxes(type='log')

    fig.update_layout(
        height=600,
        modebar_remove=['lasso', 'autoscale', 'zoomin', 'zoomout', 'select2d', 'select'],
        xaxis_title=f'{LATENCY_KEYS_TO_NAMES[latency_key]} latency, ms',
        yaxis_title=f"{METRIC_KEYS_TO_NAMES[metric_tag]}",
        xaxis=dict(
            rangeslider=dict(
                visible=True,
                bgcolor=DEEPLITE_LIGHT_BLUE_HEX,
                thickness=0.02,
            ),
        ),
        yaxis=dict(
            fixedrange=False,
        ),
        hoverlabel=dict(
            # bgcolor="white",
            font_size=14,
            font_family='Source Sans Pro'
        ),
        **layout_opts,
    )
    if add_pareto_frontier:
        fig = pareto_frontier_layer(fig, dataset_tag, metric_tag, latency_key)
    return fig


def create_yolobench_plots(
        dataset_name,
        hardware_name,
        metric_name,
        vis_options,
        table_mode,
    ):
    model_family_coloring = 'Model family' in vis_options
    add_pareto_frontier = 'Highlight Pareto' in vis_options
    plot_pareto_only = 'Show Pareto only' in vis_options
    log_axis = 'Log x-axis' in vis_options
    fig = get_scatter_plot(
        DATASET_TAGS[dataset_name],
        METRIC_NAME_MAPPING[metric_name],
        LATENCY_KEYS[hardware_name],
        model_family_coloring,
        add_pareto_frontier,
        plot_pareto_only,
        log_axis,
    )
    pareto_table = get_pareto_table(
        dataset_name, hardware_name, metric_name, expand_table='Show all' in table_mode
    )
    return fig, pareto_table


def pareto_frontier_layer(
        fig,
        dataset_tag,
        metric_tag,
        latency_key,
    ):
    metric_key = f'{metric_tag}_{dataset_tag}'
    frontier = pareto_indices[metric_key][latency_key]
    fig.add_trace(
        go.Scatter(
            x=df.iloc[frontier, :][latency_key],
            y=df.iloc[frontier, :][metric_key],
            mode='lines',
            opacity=0.5,
            line=go.scatter.Line(color='grey'),
            showlegend=False,
            name=metric_key,
        )
    )
    return fig


def get_pareto_table(
    dataset_name, hardware_name, metric_name, expand_table=False,
):
    dataset_tag = DATASET_TAGS[dataset_name]
    metric_tag = METRIC_NAME_MAPPING[metric_name]
    latency_key = LATENCY_KEYS[hardware_name]
    metric_key = f'{metric_tag}_{dataset_tag}'

    latency_key_final = f'{LATENCY_KEYS_TO_NAMES[latency_key]} latency, ms'
    metric_key_final = METRIC_KEYS_TO_NAMES[metric_tag]

    frontier = pareto_indices[metric_key][latency_key]
    table_df = df.iloc[frontier, :][['model_name', metric_key, latency_key]]
    table_df['Input resolution (px)'] = table_df['model_name'].apply(lambda name: name.split('_')[-1])
    table_df['Model name'] = table_df['model_name'].apply(lambda name: name.split('_')[0])
    table_df[metric_key_final] = table_df[metric_key].apply(lambda val: round(val, 3))
    table_df[latency_key_final] = table_df[latency_key].apply(lambda val: round(val, 2))

    def make_clickable(url, name):
        return f'<a href="{url}">{name}</a>'


    if dataset_name == 'COCO':
        table_df['Download link'] = table_df['model_name'].apply(
            lambda name: f'https://download.deeplite.ai/zoo/models/YOLOBench/{name.split("_")[0]}_640.pt'
        )
        table_df['Download link'] = table_df.apply(lambda x: make_clickable(x['Download link'], 'Weights download'), axis=1)
    else:
        table_df['Download link'] = table_df['model_name'].apply(lambda s: 'Coming soon')


    table_df = table_df[['Model name', 'Input resolution (px)',
                         metric_key_final, latency_key_final, 'Download link']].sort_values(by=metric_key_final, ascending=False)
    if not expand_table:
        table_df = table_df.iloc[:10, :]

    table_df = table_df.to_html(
        classes='table',
        escape=False, render_links=True, index=False
    )

    return table_df