import matplotlib.pyplot as plt
from typing import Union, List
import pandas as pd


class MatplotlibChartBuilder:

    def _normalize(self, y_axes: Union[str, List[str]]):
        return [y_axes] if isinstance(y_axes, str) else y_axes

    def bar(self, df: pd.DataFrame, x_axis: str, y_axes,
            title=None, x_axis_title=None, y_axis_title=None, theme=None):

        if theme:
            plt.style.use(theme)

        y_axes = self._normalize(y_axes)

        fig, ax = plt.subplots()

        for y in y_axes:
            ax.bar(df[x_axis], df[y], label=y)

        ax.set_title(title)
        ax.set_xlabel(x_axis_title or x_axis)
        ax.set_ylabel(y_axis_title or ", ".join(y_axes))
        ax.legend()

        return fig

    def line(self, df, x_axis, y_axes,
             title=None, x_axis_title=None, y_axis_title=None, theme=None):

        if theme:
            plt.style.use(theme)

        y_axes = self._normalize(y_axes)

        fig, ax = plt.subplots()

        for y in y_axes:
            ax.plot(df[x_axis], df[y], label=y)

        ax.set_title(title)
        ax.set_xlabel(x_axis_title or x_axis)
        ax.set_ylabel(y_axis_title or ", ".join(y_axes))
        ax.legend()

        return fig

    def area(self, df, x_axis, y_axes,
             title=None, x_axis_title=None, y_axis_title=None, theme=None):

        y_axes = self._normalize(y_axes)

        fig, ax = plt.subplots()
        ax.stackplot(df[x_axis], *(df[y] for y in y_axes), labels=y_axes)

        ax.set_title(title)
        ax.set_xlabel(x_axis_title or x_axis)
        ax.set_ylabel(y_axis_title or ", ".join(y_axes))
        ax.legend()

        return fig

    def pie(self, df, names, values,
            title=None, x_axis_title=None, y_axis_title=None, theme=None):

        fig, ax = plt.subplots()
        ax.pie(df[values], labels=df[names], autopct="%1.1f%%")
        ax.set_title(title)
        return fig

    def histogram(self, df, column, bins=20,
                  title=None, x_axis_title=None, y_axis_title=None, theme=None):

        fig, ax = plt.subplots()
        ax.hist(df[column], bins=bins)

        ax.set_title(title)
        ax.set_xlabel(x_axis_title or column)
        ax.set_ylabel(y_axis_title or "Count")

        return fig

    def combo(self, df, x_axis, y1, y2,
              label1=None, label2=None,
              title=None, theme=None):

        fig, ax1 = plt.subplots()
        ax2 = ax1.twinx()

        ax1.bar(df[x_axis], df[y1], alpha=0.6, label=label1 or y1)
        ax2.plot(df[x_axis], df[y2], label=label2 or y2)

        ax1.set_xlabel(x_axis)
        ax1.set_ylabel(label1 or y1)
        ax2.set_ylabel(label2 or y2)

        fig.suptitle(title)
        return fig

    def kpi(self, value, title=None, theme=None):

        fig, ax = plt.subplots()
        ax.axis("off")

        ax.text(0.5, 0.6, f"{value:,.0f}",
                fontsize=40, ha="center")

        ax.text(0.5, 0.4, title,
                fontsize=14, ha="center")

        return fig


    def kpi_compare(self, current, previous, change_pct,
                    title=None, theme=None):

        fig, ax = plt.subplots()
        ax.axis("off")

        arrow = "▲" if change_pct and change_pct > 0 else "▼"

        ax.text(0.5, 0.65, f"{current:,.0f}",
                fontsize=40, ha="center")

        ax.text(0.5, 0.45,
                f"{arrow} {change_pct:.2f}% vs prev",
                fontsize=14,
                ha="center")

        ax.text(0.5, 0.3, title,
                fontsize=12, ha="center")

        return fig
