from typing import Optional, Union, TYPE_CHECKING

from copy import deepcopy

import pandas as pd

from shimoku.plt.utils import ShimokuPalette
from shimoku.plt.EChart_definitions.default_echart_options import (
    get_common_echart_options,
    get_common_series_options,
)

if TYPE_CHECKING:
    from shimoku.plt.plt_layer import PlotLayer

import logging
from shimoku.execution_logger import log_error

logger = logging.getLogger(__name__)


async def line_chart(
    self: "PlotLayer",
    data: Union[str, pd.DataFrame, list[dict]],
    order: int,
    x: str,
    y: Union[list[str], str, None] = None,
    rows_size: Optional[int] = None,
    cols_size: Optional[int] = None,
    padding: Optional[str] = None,
    title: Optional[str] = None,
    x_axis_name: Optional[str] = None,
    y_axis_name: Optional[str] = None,
    show_values: Union[list[str], str, None] = None,
    option_modifications: Optional[dict] = None,
    variant: Optional[str] = None,
):
    """Create a line chart"""
    common_options = get_common_echart_options()
    series_options = get_common_series_options()
    series_options["type"] = "line"
    common_options["yAxis"][0]["scale"] = True
    common_options["xAxis"][0]["boundaryGap"] = False
    if variant == "minimal":
        series_options["itemStyle"] = {"opacity": 0}
        series_options["lineStyle"] = {"width": 3}

    await self._create_trend_chart(
        option_modifications=option_modifications,
        echart_options=common_options,
        series_options=series_options,
        variant=variant,
        axes=x,
        values=y,
        data_mapping_to_tuples=await self._choose_data(order, data),
        order=order,
        x_axis_names=x_axis_name,
        y_axis_names=y_axis_name,
        title=title,
        rows_size=rows_size,
        cols_size=cols_size,
        padding=padding,
        show_values=show_values,
    )


async def area_chart(
    self: "PlotLayer",
    data: Union[str, pd.DataFrame, list[dict]],
    order: int,
    x: str,
    y: Union[list[str], str, None] = None,
    rows_size: Optional[int] = None,
    cols_size: Optional[int] = None,
    padding: Optional[str] = None,
    title: Optional[str] = None,
    x_axis_name: Optional[str] = None,
    y_axis_name: Optional[str] = None,
    show_values: Union[list[str], str, None] = None,
    option_modifications: Optional[dict] = None,
    variant: Optional[str] = None,
):
    """Create an area chart"""
    common_options = get_common_echart_options()
    series_options = get_common_series_options()
    series_options.update(
        {
            "type": "line",
            "areaStyle": {"opacity": 0.5},
        }
    )
    if variant == "minimal":
        series_options["itemStyle"] = {"opacity": 0}
        series_options["areaStyle"]["opacity"] = 0.2
        series_options["lineStyle"] = {"width": 3}

    common_options["xAxis"][0]["boundaryGap"] = False
    await self._create_trend_chart(
        option_modifications=option_modifications,
        echart_options=common_options,
        series_options=series_options,
        axes=x,
        values=y,
        data_mapping_to_tuples=await self._choose_data(order, data),
        order=order,
        x_axis_names=x_axis_name,
        y_axis_names=y_axis_name,
        title=title,
        rows_size=rows_size,
        cols_size=cols_size,
        padding=padding,
        show_values=show_values,
        variant=variant,
    )


async def stacked_area_chart(
    self: "PlotLayer",
    data: Union[str, pd.DataFrame, list[dict]],
    order: int,
    x: str,
    y: Union[list[str], str, None] = None,
    rows_size: Optional[int] = None,
    cols_size: Optional[int] = None,
    padding: Optional[str] = None,
    title: Optional[str] = None,
    x_axis_name: Optional[str] = None,
    y_axis_name: Optional[str] = None,
    show_values: Union[list[str], str, None] = None,
    option_modifications: Optional[dict] = None,
    variant: Optional[str] = None,
):
    """Create a stacked area chart"""
    common_options = get_common_echart_options()
    series_options = get_common_series_options()
    series_options.update(
        {
            "type": "line",
            "areaStyle": {},
            "stack": "stack",
            "itemStyle": {"borderRadius": [0, 0, 0, 0]},
        }
    )
    common_options["xAxis"][0]["boundaryGap"] = False
    data_mapping_to_tuples = await self._choose_data(order, data)
    len_y = len(y) if y else len(data_mapping_to_tuples.keys()) - 1
    series_options = [series_options] * (len_y - 1)
    series_options.append(deepcopy(series_options[-1]))
    series_options[-1]["itemStyle"]["borderRadius"] = [9, 9, 0, 0]
    await self._create_trend_chart(
        option_modifications=option_modifications,
        echart_options=common_options,
        series_options=series_options,
        axes=x,
        values=y,
        data_mapping_to_tuples=data_mapping_to_tuples,
        order=order,
        x_axis_names=x_axis_name,
        y_axis_names=y_axis_name,
        title=title,
        rows_size=rows_size,
        cols_size=cols_size,
        padding=padding,
        show_values=show_values,
        variant=variant,
    )


async def predictive_line_chart(
    self: "PlotLayer",
    data: Union[str, pd.DataFrame, list[dict]],
    order: int,
    x: str,
    y: Union[list[str], str, None] = None,
    min_value_mark: Optional[Union[str, int]] = None,
    max_value_mark: Optional[Union[str, int]] = None,
    color_mark: str = "rgba(255, 173, 177, 0.4)",
    x_axis_name: Optional[str] = None,
    y_axis_name: Optional[str] = None,
    title: Optional[str] = None,
    rows_size: Optional[int] = None,
    cols_size: Optional[int] = None,
    padding: Optional[str] = None,
    show_values: Union[list[str], str, None] = None,
    option_modifications: Optional[dict] = None,
    variant: Optional[str] = None,
):
    """Create a predictive line chart"""
    common_options = get_common_echart_options()
    series_options = get_common_series_options()
    series_options["type"] = "line"
    common_options["yAxis"][0]["scale"] = True
    common_options["xAxis"][0]["boundaryGap"] = False
    series_options["markArea"] = {
        "itemStyle": {"color": color_mark},
        "data": [
            [
                {"name": "Prediction", "xAxis": min_value_mark},
                {"xAxis": max_value_mark},
            ]
        ],
    }
    await self._create_trend_chart(
        echart_options=common_options,
        series_options=series_options,
        option_modifications=option_modifications,
        axes=x,
        values=y,
        data_mapping_to_tuples=await self._choose_data(order, data),
        order=order,
        x_axis_names=x_axis_name,
        y_axis_names=y_axis_name,
        title=title,
        rows_size=rows_size,
        cols_size=cols_size,
        padding=padding,
        show_values=show_values,
        variant=variant,
    )


async def marked_line_chart(
    self: "PlotLayer",
    data: Union[str, pd.DataFrame, list[dict]],
    order: int,
    x: str,
    marks: list,
    y: Union[list[str], str, None] = None,
    color_mark: str = "rgba(255, 173, 177, 0.4)",
    x_axis_name: Optional[str] = None,
    y_axis_name: Optional[str] = None,
    title: Optional[str] = None,
    rows_size: Optional[int] = None,
    cols_size: Optional[int] = None,
    padding: Optional[str] = None,
    show_values: Union[list[str], str, None] = None,
    option_modifications: Optional[dict] = None,
    variant: Optional[str] = None,
):
    """Create a predictive line chart"""
    common_options = get_common_echart_options()
    series_options = get_common_series_options()
    series_options["type"] = "line"
    common_options["yAxis"][0]["scale"] = True
    common_options["xAxis"][0]["boundaryGap"] = False
    series_options["markArea"] = {
        "itemStyle": {"color": color_mark},
        "data": [
            [
                {"name": name, "xAxis": min_value_mark},
                {"xAxis": max_value_mark},
            ]
            for name, min_value_mark, max_value_mark in marks
        ],
    }
    await self._create_trend_chart(
        echart_options=common_options,
        series_options=series_options,
        option_modifications=option_modifications,
        axes=x,
        values=y,
        data_mapping_to_tuples=await self._choose_data(order, data),
        order=order,
        x_axis_names=x_axis_name,
        y_axis_names=y_axis_name,
        title=title,
        rows_size=rows_size,
        cols_size=cols_size,
        padding=padding,
        show_values=show_values,
        variant=variant,
    )


def calculate_segments(df: pd.DataFrame, threshold: float, r: int, g: int, b: int):
    color_range = df["y"].max() - threshold
    df_as_dict = df.to_dict("records")
    segment = [None, None, None]
    segments = []
    for i in range(len(df_as_dict)):
        if df_as_dict[i]["y"] >= threshold and segment[0] is None:
            segment[0] = i - 1 if i > 0 else 0
        elif (df_as_dict[i]["y"] < threshold or i == len(df_as_dict) - 1) and segment[
            0
        ] is not None:
            if i < len(df_as_dict) - 1 and df_as_dict[i + 1]["y"] >= threshold:
                continue
            segment[1] = i
            max_from_segment = max(
                [val["y"] for val in df_as_dict[segment[0] : segment[1]]]
            )
            color_alpha = 0.3 + (
                0.7 * (max_from_segment - threshold) / color_range
                if color_range > 0
                else 0
            )
            segment[2] = f"rgba({r}, {g}, {b}, {color_alpha})"
            segments.append(tuple(segment))
            segment = [None, None, None]

    return segments


async def segmented_area_chart(
    self: "PlotLayer",
    data: Union[str, pd.DataFrame, list[dict]],
    order: int,
    x: str,
    y: str,
    segments: Optional[list] = None,
    x_axis_name: Optional[str] = None,
    y_axis_name: Optional[str] = None,
    default_color: tuple[int, int, int] = (255, 0, 0),
    title: Optional[str] = None,
    rows_size: Optional[int] = None,
    cols_size: Optional[int] = None,
    padding: Optional[str] = None,
    show_values: Union[list[str], str, None] = None,
    option_modifications: Optional[dict] = None,
    variant: Optional[str] = None,
    top_area: bool = False,
    threshold: Optional[float] = None,
    labels: Optional[list[str]] = None,
):
    """Create a segmented area chart"""
    if not isinstance(y, str):
        raise TypeError("y must be string")
    if threshold is not None and segments is not None:
        raise ValueError("You cannot specify both threshold and segments")
    if threshold is None and segments is None:
        raise ValueError("You must specify either threshold or segments")
    data_mappings_to_tuples = await self._choose_data(order, data)
    if isinstance(data, str):
        data = self._shared_data[data]
    df = pd.DataFrame(data)
    if threshold is not None:
        segments = calculate_segments(df, threshold, *default_color)

    common_options = get_common_echart_options()
    series_options = get_common_series_options()
    common_options["xAxis"][0]["boundaryGap"] = False
    common_options["legend"] = {"show": False}
    del common_options["toolbox"]["feature"]["magicType"]
    del common_options["toolbox"]["bottom"]
    common_options["toolbox"]["top"] = "0px"
    common_options["dataZoom"] = (
        [{"type": "inside"}, {"type": "slider"}] if variant is None else []
    )
    pieces = []
    for segment in segments:
        if len(segment) == 2:
            r, g, b = default_color
            pieces.append(
                {
                    "gte": segment[0],
                    "lte": segment[1],
                    "color": f"rgba({r}, {g}, {b}, 0.75)",
                }
            )
        elif len(segment) == 3:
            color = segment[2]
            if not isinstance(color, str):
                if not isinstance(color, tuple) or len(color) != 3:
                    log_error(
                        logger,
                        f"Color must be a string or a tuple of 3 elements, got {color}",
                        ValueError,
                    )
                color = f"rgba({color[0]}, {color[1]}, {color[2]}, 0.75)"
            pieces.append({"gte": segment[0], "lte": segment[1], "color": color})
        else:
            raise ValueError("Segment must be a tuple of 2 or 3 elements")
    common_options["visualMap"] = {
        "type": "piecewise",
        "show": False,
        "dimension": 0,
        "seriesIndex": 0,
        "pieces": pieces,
    }
    series_options.update(
        {
            "type": "line",
            "emphasis": {"disabled": True},
            "areaStyle": {"origin": "end" if top_area else "start"},
            "itemStyle": {"opacity": 0},
        }
    )

    series_options = [series_options, deepcopy(series_options)]
    del series_options[1]["areaStyle"]

    if labels is None:
        labels = []
    if len(labels) < len(segments):
        labels += [""] * (len(segments) - len(labels))
    if len(labels) > len(segments):
        labels = labels[: len(segments)]

    series_options[1]["markArea"] = {
        "itemStyle": {"color": "rgba(0, 0, 0, 0)"},
        "emphasis": {"disabled": True},
        "label": {
            "show": variant != "minimal",
            "emphasis": {"disabled": True},
            "borderColor": f"rgba({default_color[0]}, {default_color[1]}, {default_color[2]}, 1)",
            "borderWidth": 1,
            "borderRadius": 5,
            "padding": [5, 5, 5, 5],
            "backgroundColor": "rgba(0,0,0,0)",
        },
        "data": [
            [
                {"name": label, "xAxis": marks_tuple[0]},
                {"xAxis": marks_tuple[1]},
            ]
            for label, marks_tuple in zip(labels, segments)
            if label
        ],
    }

    await self._create_trend_chart(
        echart_options=common_options,
        series_options=series_options,
        option_modifications=option_modifications,
        axes=x,
        values=[y, y],
        data_mapping_to_tuples=data_mappings_to_tuples,
        order=order,
        x_axis_names=x_axis_name,
        y_axis_names=y_axis_name,
        title=title,
        rows_size=rows_size,
        cols_size=cols_size,
        padding=padding,
        show_values=show_values,
        variant=variant,
    )


async def line_with_confidence_area_chart(
    self: "PlotLayer",
    data: Union[str, pd.DataFrame, list[dict]],
    order: int,
    x: str,
    lower: str,
    y: str,
    upper: str,
    rows_size: Optional[int] = None,
    cols_size: Optional[int] = None,
    padding: Optional[str] = None,
    title: Optional[str] = None,
    x_axis_name: Optional[str] = None,
    y_axis_name: Optional[str] = None,
    percentages: bool = False,
    option_modifications: Optional[dict] = None,
    variant: Optional[str] = None,
):
    common_options = get_common_echart_options()
    del common_options["toolbox"]["feature"]["magicType"]
    common_options.update(
        {
            "tooltip": {
                "trigger": "axis",
                "axisPointer": {
                    "type": "cross",
                    "animation": False,
                    "label": {
                        "backgroundColor": "#ccc",
                        "borderColor": "#aaa",
                        "borderWidth": 1,
                        "shadowBlur": 0,
                        "shadowOffsetX": 0,
                        "shadowOffsetY": 0,
                        "color": "#222",
                    },
                },
            },
            "xAxis": [
                {
                    "type": "category",
                    "fontFamily": "Rubik",
                    "nameLocation": "middle",
                    "boundaryGap": True,
                    "nameGap": 35,
                    "data": "#set_data#",
                }
            ],
            "yAxis": [
                {
                    "type": "value",
                    "fontFamily": "Rubik",
                    "splitNumber": 3,
                    "boundaryGap": True,
                    "nameLocation": "middle",
                    "nameGap": 60,
                    "axisLabel": {
                        "formatter": "{value}%" if percentages else "{value}"
                    },
                    "axisPointer": {
                        "label": {"formatter": "{value}%" if percentages else "{value}"}
                    },
                },
            ],
        }
    )
    series_options = [
        {
            "name": lower,
            "data": "#set_data#",
            "type": "line",
            "lineStyle": {"color": "#000", "opacity": 0.1},
            "areaStyle": {"color": "#000", "opacity": 0.1, "origin": "end"},
            "symbol": "none",
        },
        {
            "name": upper,
            "data": "#set_data#",
            "silent": True,
            "type": "line",
            "lineStyle": {"color": "#000", "opacity": 0.1},
            "areaStyle": {"color": "#000", "opacity": 0.1, "origin": "start"},
            "symbol": "none",
        },
        {
            "name": y,
            "data": "#set_data#",
            "type": "line",
            "itemStyle": {"color": ShimokuPalette.CHART_C1.value},
            "emphasis": {"lineStyle": {"color": ShimokuPalette.CHART_C1.value}},
            "showSymbol": False,
        },
    ]

    await self._create_trend_chart(
        data_mapping_to_tuples=await self._choose_data(order, data),
        order=order,
        axes=x,
        values=[lower, upper, y],
        rows_size=rows_size,
        cols_size=cols_size,
        padding=padding,
        title=title,
        x_axis_names=x_axis_name,
        y_axis_names=y_axis_name,
        echart_options=common_options,
        series_options=series_options,
        option_modifications=option_modifications,
        variant=variant,
    )


async def segmented_line_chart(
    self: "PlotLayer",
    order: int,
    x: str,
    y: Union[str, list[str]],
    data: Optional[Union[str, pd.DataFrame, list[dict]]],
    x_axis_name: Optional[str] = None,
    y_axis_name: Optional[str] = None,
    rows_size: Optional[int] = None,
    cols_size: Optional[int] = None,
    padding: Optional[str] = None,
    title: Optional[str] = None,
    marking_lines: Optional[list[int]] = None,
    range_colors: Optional[list[str]] = None,
    option_modifications: Optional[dict] = None,
    variant: Optional[str] = None,
):
    """Create a chart with bars and lines."""
    data_mappings_to_tuples = await self._choose_data(order, data)

    if isinstance(y, str):
        y = [y]

    if not marking_lines:
        marking_lines = [0]

    common_options = get_common_echart_options()
    common_options["xAxis"][0]["boundaryGap"] = False
    # Calculate the width of the visual map with the label of the marking line
    # ( label  legend_ini - legend_end spaces for the legend colors)
    common_options["grid"]["right"] = (
        max(
            [
                len(
                    f" {marking_lines[i + 1]}   {marking_lines[i]} - {marking_lines[i + 1]}        "
                )
                for i in range(0, len(marking_lines) - 1)
            ]
        )
        * 6
    )
    common_options.update(
        {
            "visualMap": {
                "top": 50,
                "right": 10,
                "itemHeight": 14 if variant != "minimal" else 0,
                "showLabel": variant != "minimal",
                "outOfRange": {"color": "#999"},
                "pieces": [
                    {"min": marking_lines[i], "max": marking_lines[i + 1]}
                    for i in range(len(marking_lines) - 1)
                ]
                + [{"min": marking_lines[-1]}],
            },
        }
    )
    if range_colors:
        common_options["visualMap"]["color"] = range_colors[::-1]
    series_options = get_common_series_options()
    series_options.update(
        {
            "type": "line",
            "color": ShimokuPalette.CHART_C5.value,
            "markLine": {
                "data": [{"yAxis": val} for val in marking_lines],
                "label": {"show": variant != "minimal"},
                "itemStyle": {"color": ShimokuPalette.CHART_C5.value, "opacity": 0.5},
            },
        }
    )

    await self._create_trend_chart(
        data_mapping_to_tuples=data_mappings_to_tuples,
        axes=x,
        values=y,
        echart_options=common_options,
        series_options=series_options,
        order=order,
        x_axis_names=x_axis_name,
        y_axis_names=y_axis_name,
        title=title,
        rows_size=rows_size,
        cols_size=cols_size,
        padding=padding,
        option_modifications=option_modifications,
        variant=variant,
    )
