from matplotlib.axes import Axes
from matplotlib.projections import register_projection
from matplotlib.ticker import AutoMinorLocator
import matplotlib.pyplot as plt
import numpy as np
from tqdm import tqdm
import matplotlib
import warnings
# matplotlib.use("tkAgg")  


LHCB_FONT_FAMILY = "Times New Roman"
LHCB_LINE_WIDTH = 1.0
LHCB_TEXT_SIZE = 14
LHCB_MAJOR_TICK_LENGTH = 8
LHCB_MINOR_TICK_LENGTH = 4
# Follow named ROOT color order: black, red, green, blue, yellow, magenta, cyan, purple.
LHCB_COLOR_CYCLE = ["black", "red", "green", "blue", "yellow", "magenta", "cyan", "purple"]


class HepAxes(Axes):
    name = "hep"

    def __init__(self, *args, **kwargs):
        self._hep_font_family = LHCB_FONT_FAMILY
        super().__init__(*args, **kwargs)
        self._apply_lhcb_style()
        self.hist_datas = []
        self._hist_artists = []
        self.grid(True)



    def _apply_lhcb_style(self):
        # ROOT lhcbStyle-like appearance, applied only to this projection.
        self.set_facecolor("white")
        self.figure.set_facecolor("white")

        # Add both major and minor ticks on x/y axes.
        self.xaxis.set_minor_locator(AutoMinorLocator())
        self.yaxis.set_minor_locator(AutoMinorLocator())
        self.minorticks_on()

        self.tick_params(
            axis="both",
            which="major",
            direction="in",
            top=True,
            right=True,
            width=LHCB_LINE_WIDTH,
            length=LHCB_MAJOR_TICK_LENGTH,
            labelsize=LHCB_TEXT_SIZE,
        )

        self.tick_params(
            axis="both",
            which="minor",
            direction="in",
            top=True,
            right=True,
            width=0.8 * LHCB_LINE_WIDTH,
            length=LHCB_MINOR_TICK_LENGTH,
        )

        for spine in self.spines.values():
            spine.set_linewidth(LHCB_LINE_WIDTH)

        self.set_prop_cycle(color=LHCB_COLOR_CYCLE)
        self.title.set_fontsize(1.2 * LHCB_TEXT_SIZE)

    def plot(self, *args, **kwargs):
        # Match lhcbStyle defaults for line and marker appearance.
        kwargs.setdefault("linestyle", "-")
        kwargs.setdefault("linewidth", LHCB_LINE_WIDTH)
        kwargs.setdefault("marker", "o")
        kwargs.setdefault("markersize", 4)

        return super().plot(*args, **kwargs)

    def set_title(self, label, *args, **kwargs):
        # Match lhcbStyle title size and weight.
        kwargs.setdefault("fontsize", 1.5 * LHCB_TEXT_SIZE)
        # kwargs.setdefault("fontweight", "bold")

        return super().set_title(label, *args, **kwargs)

    def set_xlabel(self, xlabel, *args, **kwargs):
        # Put x label on the right end of the axis.
        kwargs.setdefault("loc", "right")
        kwargs.setdefault("fontsize", 1.1 * LHCB_TEXT_SIZE)

        return super().set_xlabel(xlabel, *args, **kwargs)

    def set_ylabel(self, ylabel, *args, **kwargs):
        # Keep y label vertical and move it to the top of the axis.
        kwargs.setdefault("loc", "top")
        kwargs.setdefault("rotation", 90)
        kwargs.setdefault("fontsize", 1.1 * LHCB_TEXT_SIZE)

        return super().set_ylabel(ylabel, *args, **kwargs)
    
    def grid(self, visible=None, which="both", axis="both", **kwargs):
        if visible is False:
            if which in ("both", "major"):
                super().grid(visible=False, which="major", axis=axis)
            if which in ("both", "minor"):
                super().grid(visible=False, which="minor", axis=axis)
            return

        if which in ("both", "major"):
            major_kwargs = {
                "linestyle": "-",
                "linewidth": 0.8 * LHCB_LINE_WIDTH,
                "alpha": 0.5,
            }
            major_kwargs.update(kwargs)
            super().grid(visible=visible, which="major", axis=axis, **major_kwargs)

        if which in ("both", "minor"):
            minor_kwargs = {
                "linestyle": "--",
                "linewidth": 0.6 * LHCB_LINE_WIDTH,
                "alpha": 0.25,
            }
            minor_kwargs.update(kwargs)
            super().grid(visible=visible, which="minor", axis=axis, **minor_kwargs)


    def hist(self, *args, stat=True, errors=True, **kwargs):
        # Match lhcbStyle defaults for histogram appearance.
        kwargs.setdefault("histtype", "step")
        kwargs.setdefault("linewidth", LHCB_LINE_WIDTH)
        label = kwargs.get("label", None)
        if stat:
            stat_text = self._format_hist_stat(args)
            if label is None:
                kwargs["label"] = stat_text
            elif "$\\mu$=" not in label and "$\\sigma$=" not in label:
                kwargs["label"] = f"{label} {stat_text}"

        # remove stat and errors from kwargs before caching (will handle separately)
        self.hist_datas.append({"args": args, "kwargs": kwargs.copy(), "errors": errors})

        common_bins, common_range = self._get_common_hist_config()

        # Remove old histogram artists before full redraw.
        for artist in self._hist_artists:
            if hasattr(artist, "remove"):
                artist.remove()
        self._hist_artists = []

        # Reset property cycle so full redraw keeps stable color/marker order.
        self.set_prop_cycle(color=LHCB_COLOR_CYCLE)

        last_result = None
        for hist_data in self.hist_datas:
            draw_kwargs = hist_data["kwargs"].copy()
            draw_kwargs["bins"] = common_bins
            if common_range is not None:
                draw_kwargs["range"] = common_range
            else:
                draw_kwargs.pop("range", None)

            last_result = super().hist(*hist_data["args"], **draw_kwargs)
            self._hist_artists.extend(self._flatten_artists(last_result[2]))
            
            # Add error bars if requested
            errors_type = hist_data.get("errors", None)
            if errors_type is not False and errors_type is not None and last_result is not []:
                error_artists = self._draw_hist_errors(
                    hist_data["args"], last_result, errors_type, **draw_kwargs
                )
                self._hist_artists.extend(error_artists)


        self.set_title("Histogram")
        self.set_ylabel("Entries")
        # self.set_xlabel("Value")

        return last_result


    def twinx(self):
        warnings.warn(

            "\n\033[1;93m#### [NOT RECOMMENDED] HepAxes.twinx() have issues in current version ####\033[0m\n" \
            "\033[1;93mtwin_ax do not share the same bin with the main ax, so histograms drawn on twin_ax will not be perfectly aligned with those on the main ax!!!\033[0m\n" \
            "Will be fixed in future versions.",
            UserWarning,
            stacklevel=2,
        )
        # Create a twin axes with the same custom projection so hep-specific
        # features (e.g. hist stat/errors kwargs) work on the right axis too.
        twin = self._make_twin_axes(sharex=self, projection=self.name)
        twin.yaxis.tick_right()
        twin.yaxis.set_label_position('right')
        twin.yaxis.set_offset_position('right')
        twin.set_autoscalex_on(self.get_autoscalex_on())
        self.yaxis.tick_left()
        twin.xaxis.set_visible(False)
        twin.patch.set_visible(False)
        twin.xaxis.units = self.xaxis.units

        twin._hep_font_family = self._hep_font_family
        return twin

    def _format_hist_stat(self, hist_args):
        values = []
        for dataset in self._extract_hist_datasets(hist_args):
            finite = dataset[np.isfinite(dataset)]
            if finite.size:
                values.append(finite)

        if not values:
            return "(N=0), $\\mu$=nan, $\\sigma$=nan"

        stacked = np.concatenate(values)
        count = int(stacked.size)
        mean = float(np.mean(stacked))
        sigma = float(np.std(stacked))
        return f"(N={count})\n$\\mu$={mean:.2g}, $\\sigma$={sigma:.2g}"

    def _get_common_hist_config(self):
        explicit_bins = None
        explicit_range = None
        for hist_data in self.hist_datas:
            draw_kwargs = hist_data["kwargs"]
            if "bins" in draw_kwargs:
                explicit_bins = draw_kwargs["bins"]
            if "range" in draw_kwargs:
                explicit_range = draw_kwargs["range"]

        bins = explicit_bins if explicit_bins is not None else 10
        if explicit_range is not None:
            return bins, explicit_range

        values = []
        for hist_data in self.hist_datas:
            for dataset in self._extract_hist_datasets(hist_data["args"]):
                finite = dataset[np.isfinite(dataset)]
                if finite.size:
                    values.append(finite)

        if not values:
            return bins, None

        stacked = np.concatenate(values)
        x_min = float(np.min(stacked))
        x_max = float(np.max(stacked))
        if x_min == x_max:
            delta = 0.5 if x_min == 0.0 else abs(0.5 * x_min)
            x_min -= delta
            x_max += delta

        return bins, (x_min, x_max)

    def _extract_hist_datasets(self, hist_args):
        if not hist_args:
            return []

        x = hist_args[0]
        if isinstance(x, (list, tuple)):
            if not x:
                return []
            if np.isscalar(x[0]):
                return [np.asarray(x).ravel()]
            return [np.asarray(dataset).ravel() for dataset in x]

        x_array = np.asarray(x)
        if x_array.ndim <= 1:
            return [x_array.ravel()]

        return [x_array[:, i].ravel() for i in range(x_array.shape[1])]

    def _flatten_artists(self, maybe_nested):
        if isinstance(maybe_nested, (list, tuple)):
            artists = []
            for item in maybe_nested:
                artists.extend(self._flatten_artists(item))
            return artists

        return [maybe_nested]

    def _draw_hist_errors(self, hist_args, hist_result, errors, **kwargs):
        """Draw error bars on top of histogram. Returns list of artist objects."""
        # bins = kwargs["bins"]
        # data_range = kwargs["range"]
        if errors is True:
            errors_type = "stat"
        elif type(errors) is np.ndarray:
            errors_type = "custom"
        else:
            raise ValueError(f"Invalid errors argument: {errors}. Must be True, False, or array-like.")

        # Extract histogram values and bin edges
        hist_counts = hist_result[0]  # bin heights
        bin_edges = hist_result[1]    # bin edges
        
        # Compute bin centers
        bin_centers = (bin_edges[:-1] + bin_edges[1:]) / 2.0
        bin_width = bin_edges[1] - bin_edges[0]
        
        # Compute errors based on type
        if errors_type == "stat":
            errors = np.sqrt(np.maximum(hist_counts, 0))
            if "density" in kwargs and kwargs["density"]:
                # For density histograms, propagate Poisson errors from raw bin counts:
                # density_i = n_i / (N * bin_width_i), sigma_i = sqrt(n_i) / (N * bin_width_i)
                data = np.asarray(hist_args[0]).ravel()
                finite = data[np.isfinite(data)]
                raw_counts, _ = np.histogram(finite, bins=bin_edges)
                total = np.sum(raw_counts)
                if total > 0:
                    bin_widths = np.diff(bin_edges)
                    errors = np.sqrt(raw_counts) / (total * bin_widths)
                else:
                    errors = np.zeros_like(hist_counts, dtype=float)
        elif errors_type == "custom":
            if len(errors) != len(hist_counts):
                raise ValueError(f"Custom errors array length {len(errors)} does not match number of bins {len(hist_counts)}.")
            errors = np.asarray(errors)
        else:
            raise ValueError(f"Unknown errors type: {errors_type}")



        # Get current color from the last-drawn histogram line
        color = None
        for artist in self._flatten_artists(hist_result[2]):
            if hasattr(artist, "get_edgecolor"):
                color = artist.get_edgecolor()
                break
        
        # Draw error bars
        bin_width = bin_edges[1] - bin_edges[0]
        errorbar_artists = self.errorbar(
            bin_centers, hist_counts, yerr=errors, 
            fmt=".", elinewidth=LHCB_LINE_WIDTH, 
            color=color, alpha=0.7
        )
        artists = []

        # 手动画 caps（用数据坐标）
        cap_width = bin_width * 0.5  # 或者你想要的比例

        for x, y, err in zip(bin_centers, hist_counts, errors):
            cap_artistup = self.hlines(y + err, x - 0.5*cap_width, x + 0.5*cap_width,
                    color=color, linewidth=LHCB_LINE_WIDTH)
            cap_artistdown = self.hlines(y - err, x - 0.5*cap_width, x + 0.5*cap_width,
                    color=color, linewidth=LHCB_LINE_WIDTH)

            artists.append(cap_artistup)
            artists.append(cap_artistdown)


        # Return all artist objects from errorbar
        if hasattr(errorbar_artists, '__iter__'):
            artists.extend(self._flatten_artists(errorbar_artists))
        else:
            artists.append(errorbar_artists)
        return artists


    def draw(self, renderer, *args, **kwargs):
        # Apply font only to text artists in this axes.
        for text in self.findobj(match=lambda artist: hasattr(artist, "set_fontfamily")):
            text.set_fontfamily(self._hep_font_family)

        return super().draw(renderer, *args, **kwargs)


# Register the projection at import time so pyplot can resolve "hep".
register_projection(HepAxes)


if __name__ == "__main__":


    fig, ax = plt.subplots(subplot_kw={"projection": "hep"})
    ax2 = ax.twinx()
    # ax.plot([1, 2, 3], [1, 2, 3], label="Data 1")
    # ax.plot([1, 2, 3], [4, 5, 6], label="Data 2")
    ax.hist(np.random.randn(1000)**2, bins=100, histtype="step", label="Data 1", stat=True, errors=True)
    ax2.hist(np.random.randn(1000)+3, bins=100, histtype="step", label="Data 2", stat=True, errors=np.linspace(5, 20, 100))
    ax.set_xlabel("X-axis")
    ax.set_ylabel("Y-axis")
    ax.grid(True)

    ax.legend()

    ax.set_title("HepAxes Example")

    plt.tight_layout()
    plt.savefig("hep_axes_example.png", dpi=6000)
    plt.show()
