from collections import OrderedDict

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib import cm, colors
from mpl_toolkits.axes_grid1.inset_locator import inset_axes

from .color import plot_colorbar
from .utilities import tight_hue_range


def sunbrust(
    pie_data,
    ax,
    hue=None,
    hue_portion=0.5,
    cmap="viridis",
    colorbar=True,
    colorbar_kws=None,
    inner_radius=0.25,
    outer_radius=1,
    anno_col=None,
    text_anno="text",
    anno_layer_size=0.05,
    col_color_dict=None,
    startangle=0,
    anno_ang_min=5,
    anno_border=1.2,
    text_expend=1.05,
    uniform_section=False,
    order_dict=None,
):
    """Plot a sunburst diagram."""
    return_axes = [ax]
    if order_dict is None:
        order_dict = {}

    # prepare hue colormap
    if hue is not None:
        if isinstance(hue, str):
            hue_data = pie_data.pop(hue)
        elif isinstance(hue, pd.Series):
            hue_data = hue
        else:
            hue_data = pd.Series(hue, index=pie_data.index)
        if not isinstance(cmap, colors.Colormap):
            cmap = plt.get_cmap(cmap)
        cmap.set_bad("gray")

        vmin, vmax = tight_hue_range(hue_data, hue_portion)
        norm = colors.Normalize(vmin=vmin, vmax=vmax, clip=True)
        mapper = cm.ScalarMappable(norm=norm, cmap=cmap)
        col_color_dict = {}
        for col_name, col in pie_data.iteritems():
            reduce_col_data = hue_data.groupby(col).mean()
            col_color_dict[col_name] = {k: mapper.to_rgba(v) for k, v in reduce_col_data.iteritems()}

        # make colorbar
        if colorbar:
            _colorbar_kws = {
                "labelsize": 4,
                "linewidth": 0.5,
                "orientation": "vertical",
            }
            if colorbar_kws is None:
                colorbar_kws = {}
            _colorbar_kws.update(colorbar_kws)
            bbox_to_anchor = _colorbar_kws.pop("bbox_to_anchor")

            cax = inset_axes(
                ax,
                width="3%",
                height="25%",
                loc="lower right",
                borderpad=0,
                bbox_to_anchor=bbox_to_anchor,
            )
            cax = plot_colorbar(
                cax,
                cmap=cmap,
                cnorm=norm,
                hue_norm=(vmin, vmax),
                label=hue_data.name,
                **_colorbar_kws,
            )
            return_axes.append(cax)

    # prepare data
    if uniform_section:
        dedup_groups = pie_data.reset_index().set_index(pie_data.columns.tolist()).index.drop_duplicates()
        dedup_groups = pd.DataFrame(dedup_groups.tolist(), columns=pie_data.columns)
        pie_data = dedup_groups

    # prepare color
    _col_color_dict = {col: None for col in pie_data.columns}
    if col_color_dict is not None:
        _col_color_dict.update(col_color_dict)

    # prepare plot parameters
    ncols = pie_data.columns.size
    if anno_col is None:
        anno_layer_size = 0
    outer_radius = outer_radius - anno_layer_size
    layer_size = (outer_radius - inner_radius) / ncols

    # plot multiple donuts
    previous_order = pd.Series([])
    anno_wedges = []
    anno_names = []
    for col, col_name in enumerate(pie_data.columns):
        cur_radius = inner_radius + (col + 1) * layer_size
        col_pie_data = pie_data[col_name].value_counts()

        # manage order
        if col_name in order_dict:
            _ordered_data = col_pie_data.reindex(pd.Index(order_dict[col_name]))
            previous_order = _ordered_data
        else:
            if col == 0:
                _ordered_data = col_pie_data
                previous_order = _ordered_data
            else:
                records = []
                for section in previous_order.index:
                    section_subs = pie_data[pie_data.iloc[:, col - 1] == section][col_name].unique()
                    records.append(col_pie_data.reindex(pd.Index(section_subs)).sort_values(ascending=False))
                _ordered_data = pd.concat(records)
                previous_order = _ordered_data

        # plot the real pie charts
        pie_color = _col_color_dict[col_name]
        if isinstance(pie_color, dict):
            pie_color = [pie_color[i] for i in _ordered_data.index]
        ax.pie(
            _ordered_data,
            radius=cur_radius,
            colors=pie_color,
            startangle=startangle,
            wedgeprops={"width": layer_size, "edgecolor": "w"},
        )

        # plot an additional thin layer to anchor anno label
        if anno_col == col:
            wedges, texts = ax.pie(
                _ordered_data,
                radius=outer_radius + anno_layer_size,
                colors=pie_color,
                startangle=startangle,
                wedgeprops={"width": anno_layer_size, "edgecolor": "w"},
            )
            if text_anno:
                anno_wedges = wedges
                anno_names = _ordered_data.index.tolist()

    # annotate wedges
    for i, p in enumerate(anno_wedges):
        delta_ang = p.theta2 - p.theta1
        if delta_ang < anno_ang_min:
            continue
        ang = (p.theta2 - p.theta1) / 2.0 + p.theta1
        y = np.sin(np.deg2rad(ang))
        x = np.cos(np.deg2rad(ang))

        if text_anno == "anno_box":
            # wedges annotation
            bbox_props = {"boxstyle": "round,pad=0.2", "fc": "#FFFFFF88", "ec": "#00000022", "lw": 0.72}
            kw = {
                "xycoords": "data",
                "textcoords": "data",
                "arrowprops": {"arrowstyle": "-", "color": "#00000055"},
                "bbox": bbox_props,
                "zorder": 0,
                "va": "center",
            }
            # separate all y
            y_niche = np.arange(-anno_border, anno_border, anno_border / 10)
            allow_y_niche = OrderedDict(enumerate(y_niche))

            horizontalalignment = {-1: "right", 1: "left"}[int(np.sign(x))]
            connectionstyle = f"angle,angleA=0,angleB={ang},rad=5"
            kw["arrowprops"].update({"connectionstyle": connectionstyle})

            suitable_niche = np.abs(np.array(list(allow_y_niche.values())) - anno_border * y).argmin()
            suitable_niche = list(allow_y_niche.keys())[suitable_niche]
            y_niche = allow_y_niche.pop(suitable_niche)
            ax.annotate(
                anno_names[i],
                xy=(x, y),
                xytext=(anno_border * np.sign(x), y_niche),
                horizontalalignment=horizontalalignment,
                **kw,
            )
        elif text_anno == "text":
            ha = "left"
            if x < 0:
                # left side label
                ang += 180
                ha = "right"
            if ang > 180:
                # so label will not be up-side-down
                ang = ang - 360
            ax.text(
                x * text_expend,
                y * text_expend,
                anno_names[i],
                fontdict=None,
                withdash=False,
                rotation=ang,
                va="center",
                ha=ha,
                rotation_mode="anchor",
            )
        elif text_anno is None:
            pass
        else:
            raise ValueError(f'text_anno can only be "text", "anno_box" or None, got {text_anno}')

    if len(return_axes) == 1:
        return ax
    else:
        return tuple(return_axes)
