import matplotlib.pyplot as plt
import json
import tempfile
import webbrowser
import warnings
from matplotlib.image import AxesImage
from matplotlib.collections import QuadMesh, PolyCollection
try:
    from IPython.display import display, HTML
    _HAS_IPYTHON = True
except ImportError:
    _HAS_IPYTHON = False

class Ptoe:
    def __init__(self, plt_module):
        self._plt = plt_module

    def __getattr__(self, name):
        return getattr(self._plt, name)

    def rgba_to_echarts(self, color):
        if isinstance(color, str):
            return color

        if isinstance(color, (list, tuple)) and len(color) >= 3:
            r = int(color[0] * 255)
            g = int(color[1] * 255)
            b = int(color[2] * 255)
            a = color[3] if len(color) > 3 else 1.0
            return f"rgba({r},{g},{b},{a})"

        return None

    def to_json_safe(self, obj):
        if isinstance(obj, dict):
            return {k: self.to_json_safe(v) for k, v in obj.items()}

        if isinstance(obj, list):
            return [self.to_json_safe(v) for v in obj]

        if isinstance(obj, tuple):
            return [self.to_json_safe(v) for v in obj]

        try:
            import numpy as np
            if isinstance(obj, (np.integer,)):
                return int(obj)
            if isinstance(obj, (np.floating,)):
                return float(obj)
        except ImportError:
            pass

        return obj


    def to_py(self, v):
        if hasattr(v, "item"):
            return v.item()
        return v


    def render_html(self, option):
        safe_option = self.to_json_safe(option)
        html = f"""
    <!DOCTYPE html>
    <html>
    <head>
    <meta charset="utf-8" />
    <script src="https://cdn.jsdelivr.net/npm/echarts@5/dist/echarts.min.js"></script>
    </head>
    <body>
    <div id="chart" style="width: 100%; height: 600px;"></div>
    <script>
        const chart = echarts.init(document.getElementById('chart'));
        const option = {json.dumps(safe_option, ensure_ascii=False)};
        chart.setOption(option);
    </script>
    </body>
    </html>
    """
        with tempfile.NamedTemporaryFile("w", suffix=".html", delete=False, encoding="utf-8") as f:
            f.write(html)
            webbrowser.open(f.name)

    def render_inline(self, option):
        if not _HAS_IPYTHON:
            raise RuntimeError(
                "IPython is not available. "
                "Use inline=True only in Jupyter/Colab."
            )

        safe_option = self.to_json_safe(option)

        html = f"""
        <div id="chart" style="width: 100%; height: 500px;"></div>
        <script src="https://cdn.jsdelivr.net/npm/echarts@5/dist/echarts.min.js"></script>
        <script>
        const chart = echarts.init(document.getElementById('chart'));
        const option = {json.dumps(safe_option, ensure_ascii=False)};
        chart.setOption(option);
        </script>
        """
        display(HTML(html))


    def detect_chart_type(self, ax):
        if any(isinstance(im, AxesImage) for im in ax.images):
            return "heatmap"

        if any(isinstance(c, QuadMesh) for c in ax.collections):
            return "heatmap"

        if self.detect_boxplot(ax):
            return "boxplot"

        if any(isinstance(c, PolyCollection) for c in ax.collections):
            return "area"

        if ax.collections:
            return "scatter"

        if ax.lines:
            return "line"

        if ax.containers:
            return "bar"

        raise NotImplementedError("unsupported chart type")


    def is_stacked_bar(self, ax):
        if len(ax.containers) < 2:
            return False

        for container in ax.containers[1:]:
            for rect in container:
                if rect.get_y() > 0:
                    return True
        return False


    def extract_bar(self, ax):
        labels = [tick.get_text() for tick in ax.get_xticklabels()]
        _, legend_labels = ax.get_legend_handles_labels()

        series = []

        is_stacked = self.is_stacked_bar(ax) 

        for idx, container in enumerate(ax.containers):
            values = []

            color = None
            if len(container) > 0:
                color = self.rgba_to_echarts(container[0].get_facecolor())

            name = None
            if idx < len(legend_labels):
                name = legend_labels[idx]

            for rect in container:
                values.append(self.to_py(rect.get_height()))

            series.append({
                "type": "bar",
                "name": name,
                "data": values,
                "color": color,
                "x_labels": labels,
                "stack": "total" if is_stacked else None,
            })

        return series



    def extract_hist_enhanced(
        self,
        data,
        bins=None,
        density=False,
        cumulative=False,
        hist_range=None,
        color=None,
        alpha=None,
        edgecolor=None,
    ):
        import numpy as np

        data = np.asarray(data)

        counts, bin_edges = np.histogram(
            data,
            bins=bins,
            density=density,
             range=hist_range,
        )

        if cumulative:
            counts = np.cumsum(counts)

            if density:
                bin_widths = np.diff(bin_edges)
                counts = counts * bin_widths

        # category x-axis용 bin label 생성 (막대 개수 = bins 보장)
        x_labels = [
            f"{bin_edges[i]:.2f} ~ {bin_edges[i + 1]:.2f}"
            for i in range(len(bin_edges) - 1)
        ]

        series = {
            "type": "bar",
            "name": "hist",
            "data": [self.to_py(v) for v in counts],
            "x_labels": x_labels,
            "barWidth": "100%",
        }
  # --------------------------
  # 스타일 매핑
  # --------------------------

        item_style = {}

        if color:
            if alpha is not None:
                import matplotlib.colors as mcolors
                r, g, b = mcolors.to_rgb(color)
                r = int(r * 255)
                g = int(g * 255)
                b = int(b * 255)
                item_style["color"] = f"rgba({r},{g},{b},{alpha})"
            else:
                item_style["color"] = color

        if edgecolor:
            item_style["borderColor"] = edgecolor
            item_style["borderWidth"] = 1

        if item_style:
            series["itemStyle"] = item_style

        return [series]

    def extract_line(self, ax):
        series = []
        for line in ax.lines:
            x = list(line.get_xdata())
            y = list(line.get_ydata())
            color = self.rgba_to_echarts(line.get_color())
            label = line.get_label()

            drawstyle = line.get_drawstyle()
            step = None
            if drawstyle.startswith("steps"):
                if drawstyle == "steps-pre":
                    step = "start"
                elif drawstyle == "steps-mid":
                    step = "middle"
                elif drawstyle == "steps-post":
                    step = "end"

            item = {
                "type": "line",
                "name": label if not label.startswith("_") else None,
                "data": [
                    [self.to_py(xi), self.to_py(yi)]
                    for xi, yi in zip(x, y)
                ],
                "color": color,
            }

            if step:
                item["step"] = step

            series.append(item)

        return series



    def extract_scatter(self, ax):
        series = []
        for col in ax.collections:
            offsets = col.get_offsets()
            xs = offsets[:, 0].tolist()
            ys = offsets[:, 1].tolist()

            facecolors = col.get_facecolors()
            color = (
                self.rgba_to_echarts(tuple(facecolors[0]))
                if len(facecolors) > 0 else None
            )

            label = col.get_label()

            series.append({
                "type": "scatter",
                "name": label if not label.startswith("_") else None,
                "data": [
                    [self.to_py(xi), self.to_py(yi)]
                    for xi, yi in zip(xs, ys)
                ],
                "color": color,
            })
        return series

    def extract_area(self, ax):
        series = []

        for col in ax.collections:
            if not hasattr(col, "get_paths"):
                continue

            paths = col.get_paths()
            if not paths:
                continue

            verts = paths[0].vertices
            x = verts[:, 0].tolist()
            y = verts[:, 1].tolist()

            color = None
            facecolors = col.get_facecolors()
            if len(facecolors) > 0:
                color = self.rgba_to_echarts(tuple(facecolors[0]))

            series.append({
                "type": "line",
                "name": None,
                "data": [
                    [self.to_py(xi), self.to_py(yi)]
                    for xi, yi in zip(x, y)
                ],
                "color": color,
                "areaStyle": {},
            })

        return series



    def extract_boxplot(self, ax):
        lines = ax.lines
        series = []

        # box 하나당 6개 line
        BOX_LINE_COUNT = 6
        box_count = len(lines) // BOX_LINE_COUNT

        data = []
        labels = []

        for i in range(box_count):
            group = lines[i * BOX_LINE_COUNT:(i + 1) * BOX_LINE_COUNT]

            values = []

            for line in group:
                values.extend(line.get_ydata())

            # values 안에 min, q1, median, q3, max 다 있음
            values = sorted(set(values))

            if len(values) >= 5:
                data.append([
                    self.to_py(values[0]),  # min
                    self.to_py(values[1]),  # Q1
                    self.to_py(values[2]),  # median
                    self.to_py(values[3]),  # Q3
                    self.to_py(values[-1]), # max
                ])
                labels.append(f"Box {i+1}")

        series.append({
            "type": "boxplot",
            "name": "boxplot",
            "data": data,
            "x_labels": labels,
        })

        return series


    def detect_boxplot(self, ax):
        if len(ax.lines) < 5:
            return False


        vertical_lines = 0
        for line in ax.lines:
            x = line.get_xdata()
            y = line.get_ydata()
            if len(x) >= 2 and abs(x[0] - x[1]) < 1e-6:
                vertical_lines += 1

        return vertical_lines >= 2

    def extract_boxplot_enhanced(self, data):
        import numpy as np

        data = np.asarray(data)

        q1 = np.percentile(data, 25)
        median = np.percentile(data, 50)
        q3 = np.percentile(data, 75)
        iqr = q3 - q1

        lower_whisker = np.min(data[data >= q1 - 1.5 * iqr])
        upper_whisker = np.max(data[data <= q3 + 1.5 * iqr])

        series_data = [[
            self.to_py(lower_whisker),
            self.to_py(q1),
            self.to_py(median),
            self.to_py(q3),
            self.to_py(upper_whisker),
        ]]

        return [{
            "type": "boxplot",
            "name": "boxplot",
            "data": series_data,
            "x_labels": ["data"],
        }]



    def extract_heatmap(self, ax):
        # imshow 기준 (AxesImage)
        im = ax.images[0]
        arr = im.get_array()

        # numpy array → python list
        arr = arr.tolist()

        data = []
        for y in range(len(arr)):
            for x in range(len(arr[0])):
                data.append([x, y, arr[y][x]])

        return [{
            "type": "heatmap",
            "name": None,
            "data": data,
        }]



    def build_echarts_option(
        self,
        chart_type,
        title,
        xlabel,
        ylabel,
        x_labels,
        series_data,
        legend_labels,
    ):

        xaxis_type = (
            "category" if chart_type in ("hist", "bar", "boxplot")
            else self.infer_xaxis_type(series_data)
        )

        option = {
            "title": {"text": title} if title else {},
            "tooltip": {"trigger": "axis"},
            "legend": {"data": legend_labels} if legend_labels else {},
            "xAxis": {
                "type": xaxis_type,
                "name": xlabel,
                "data": x_labels if xaxis_type == "category" else None,
            },
            "yAxis": {
                "type": "value",
                "name": ylabel,
            },
            "series": [],
        }

        if chart_type == "hist":
            option["xAxis"]["axisLabel"] = {
                "interval": 0,
                "rotate": 45,
            }

        if chart_type == "heatmap":
            values = [v[2] for v in series_data[0]["data"]]
            option["visualMap"] = {
                "min": min(values),
                "max": max(values),
                "calculable": True,
                "orient": "horizontal",
                "left": "center",
                "bottom": 0,
            }
            option["xAxis"]["type"] = "category"
            option["yAxis"]["type"] = "category"

        for s in series_data:
            series_item = {
                "type": s["type"],
                "name": s.get("name"),
                "data": s["data"],
            }

            if "barWidth" in s:
                series_item["barWidth"] = s["barWidth"]

            if "itemStyle" in s:
                series_item["itemStyle"] = s["itemStyle"]

            elif s.get("color"):
                series_item["itemStyle"] = {"color": s["color"]}

            if s.get("stack"):
                series_item["stack"] = s["stack"]

            option["series"].append(series_item)

        return option




    def infer_xaxis_type(self, series_data):
        if not series_data:
            return "value"

        data = series_data[0].get("data", [])
        if not data:
            return "value"

        first = data[0]
        if isinstance(first, (list, tuple)):
            x = first[0]
        else:
            x = first

        return "category" if isinstance(x, str) else "value"
        

    def show(
        self,
        inline=False,
        data=None,
        chart_type=None,
        bins=None,
        density=False,
        cumulative=False,
        hist_range=None,
        color=None,
        alpha=None,
        edgecolor=None,
    ):
        fig = self._plt.gcf()

        # -----------------------------
        # 1) chart_type + data 모드
        # -----------------------------
        if chart_type and data is not None:

            title = ""
            xlabel = ""
            ylabel = ""
            legend_labels = []

            if chart_type == "hist":
                series_data = self.extract_hist_enhanced(
                    data=data,
                    bins=bins,
                    density=density,
                    cumulative=cumulative,
                    hist_range=hist_range,
                    color=color,
                    alpha=alpha,
                    edgecolor=edgecolor,
                )

            elif chart_type == "boxplot":
                series_data = self.extract_boxplot_enhanced(data)

            else:
                raise ValueError("Unsupported enhanced chart_type")

        # -----------------------------
        # 2) matplotlib 감싸기 모드
        # -----------------------------
        else:
            if len(fig.axes) != 1:
                raise NotImplementedError("only single axes supported")

            ax = fig.axes[0]

            title = ax.get_title()
            xlabel = ax.get_xlabel()
            ylabel = ax.get_ylabel()
            _, legend_labels = ax.get_legend_handles_labels()

            chart_type = self.detect_chart_type(ax)

            if chart_type == "boxplot":
                series_data = self.extract_boxplot(ax)
            elif chart_type == "bar":
                series_data = self.extract_bar(ax)
            elif chart_type == "line":
                series_data = self.extract_line(ax)
            elif chart_type == "scatter":
                series_data = self.extract_scatter(ax)
            elif chart_type == "area":
                series_data = self.extract_area(ax)
            elif chart_type == "heatmap":
                series_data = self.extract_heatmap(ax)
            else:
                raise NotImplementedError(chart_type)

        option = self.build_echarts_option(
            chart_type=chart_type,
            title=title,
            xlabel=xlabel,
            ylabel=ylabel,
            x_labels=series_data[0].get("x_labels"),
            series_data=series_data,
            legend_labels=legend_labels,
        )

        if inline:
            self.render_inline(option)
        else:
            self.render_html(option)

        self._plt.close(fig)
