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


class SeabornChartBuilder:

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

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

        sns.set_theme(style=theme or "whitegrid")

        y_axes = self._normalize(y_axes)
        df_melt = df.melt(id_vars=x_axis, value_vars=y_axes)

        fig, ax = plt.subplots()
        sns.barplot(data=df_melt, x=x_axis, y="value", hue="variable", ax=ax)

        ax.set_title(title)
        ax.set_xlabel(x_axis_title or x_axis)
        ax.set_ylabel(y_axis_title or "Value")

        return fig

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

        sns.set_theme(style=theme or "whitegrid")

        y_axes = self._normalize(y_axes)
        df_melt = df.melt(id_vars=x_axis, value_vars=y_axes)

        fig, ax = plt.subplots()
        sns.lineplot(data=df_melt, x=x_axis, y="value", hue="variable", ax=ax)

        ax.set_title(title)
        ax.set_xlabel(x_axis_title or x_axis)
        ax.set_ylabel(y_axis_title or "Value")

        return fig

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

        # Seaborn doesn't directly support stacked area well,
        # so fallback to matplotlib stackplot.

        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()
        sns.histplot(df[column], bins=bins, ax=ax)

        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()

        sns.barplot(x=df[x_axis], y=df[y1], ax=ax1)
        sns.lineplot(x=df[x_axis], y=df[y2], ax=ax2)

        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):
        import seaborn as sns
        import matplotlib.pyplot as plt

        sns.set_theme(style="white")

        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
