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',
    'Raspberry Pi 4 Model B (CPU, ONNX Runtime, FP32)': 'pi4_ort_latency',
    'Raspberry Pi 5 Model B (CPU, ONNX Runtime, FP32)': 'pi5_ort_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)
    fig.update_layout(autosize=True)
    return fig


def get_comparison_plot(
    dataset_tag,
    metric_tag,
    latency_keys,
    log_axis=False,
    remove_marker=False,
):
    if len(latency_keys) == 0:
        layout_opts = {
            "annotations": [
                {
                    "text": "Please select at least 2 hardware targets to compare!",
                    "showarrow": False,
                    "font": {"size": 28},
                }
            ]
        }
    else:
        layout_opts = {
            'legend': dict(
                title='Hardware targets selected<br>(click to toggle)',
            )
        }

    fig = go.Figure(
        layout=go.Layout(
            title=go.layout.Title(
                text=f'{METRIC_KEYS_TO_NAMES[metric_tag]}-latency Pareto Optimal Models'
            )
        )
    )

    if remove_marker:
        mode = 'lines'
    else:
        mode = 'lines+markers'

    for latency_key in latency_keys:
        metric_key = f'{metric_tag}_{dataset_tag}'
        frontier = pareto_indices[metric_key][latency_key]
        df_pareto = df.iloc[frontier, :]
        model_name = df_pareto['model_name']
        fig.add_trace(
            go.Scatter(
                x=df_pareto[latency_key],
                y=df_pareto[f'{metric_tag}_{dataset_tag}'],
                name=LATENCY_KEYS_TO_NAMES[latency_key],
                mode=mode,
                customdata=model_name,
                hovertemplate='<b>model name:%{customdata}</b><br>latency:%{x:.2f} <br>metric: %{y:.2f} ',
            )
        )

    if log_axis:
        fig.update_xaxes(type='log')
        x_axis_title = f'Inference latency, ms (BS=1, log scale)'
    else:
        x_axis_title = f'Inference latency, ms (BS=1)'

    fig.update_layout(
        height=600,
        plot_bgcolor='white',
        modebar_remove=[
            'lasso',
            'autoscale',
            'zoomin',
            'zoomout',
            'select2d',
            'select',
        ],
        xaxis_title=x_axis_title,
        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,
    )
    fig.update_layout(autosize=True, template="plotly_white")

    return fig


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+markers',
            opacity=0.5,
            line=go.scatter.Line(color='grey'),
            showlegend=False,
            name=metric_key,
            marker=dict(symbol=['circle']),
        )
    )
    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 create_comparison_plot(dataset_name, hardware_list, metric_name, vis_options):
    log_axis = 'Log x-axis' in vis_options
    remove_marker = 'Remove datapoint markers' in vis_options

    latency_keys = []
    for hardware_name in hardware_list:
        latency_keys.append(LATENCY_KEYS[hardware_name])

    fig = get_comparison_plot(
        DATASET_TAGS[dataset_name],
        METRIC_NAME_MAPPING[metric_name],
        latency_keys,
        log_axis,
        remove_marker,
    )
    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