import altair as alt
from typing import Union, List
import pandas as pd


class AltairChartBuilder:

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

        y_axes = self._normalize(y_axes)

        base = alt.Chart(df).properties(title=title)

        charts = [
            base.mark_bar().encode(
                x=alt.X(x_axis, title=x_axis_title or x_axis),
                y=alt.Y(y, title=y_axis_title or y)
            )
            for y in y_axes
        ]

        return alt.layer(*charts)

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

        y_axes = self._normalize(y_axes)

        base = alt.Chart(df).properties(title=title)

        charts = [
            base.mark_line().encode(
                x=alt.X(x_axis, title=x_axis_title or x_axis),
                y=alt.Y(y, title=y_axis_title or y)
            )
            for y in y_axes
        ]

        return alt.layer(*charts)

    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)

        base = alt.Chart(df).properties(title=title)

        charts = [
            base.mark_area().encode(
                x=alt.X(x_axis, title=x_axis_title or x_axis),
                y=alt.Y(y, title=y_axis_title or y)
            )
            for y in y_axes
        ]

        return alt.layer(*charts)

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

        return alt.Chart(df).mark_arc().encode(
            theta=values,
            color=names
        ).properties(title=title)

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

        return alt.Chart(df).mark_bar().encode(
            alt.X(column, bin=alt.Bin(maxbins=bins),
                  title=x_axis_title or column),
            alt.Y("count()", title=y_axis_title or "Count")
        ).properties(title=title)

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

        bar = alt.Chart(df).mark_bar().encode(
            x=alt.X(x_axis, title=x_axis),
            y=alt.Y(y1, title=label1 or y1)
        )

        line = alt.Chart(df).mark_line(color="red").encode(
            x=x_axis,
            y=alt.Y(y2, title=label2 or y2)
        )

        return alt.layer(bar, line).properties(title=title)

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

        df = pd.DataFrame({"value": [value]})

        return alt.Chart(df).mark_text(
            size=40
        ).encode(
            text="value"
        ).properties(title=title)
