import time

import numpy as np
import uproot
from matplotlib import pyplot as plt

from .Tree import Tree
from ..HEP_PLOT import HepAxes  # noqa: F401

class ROOT:
    def __init__(self, filename=None):
        self.rt_file = None
        self.load(filename=filename)

    def __str__(self):
        rt_info = ""
        if self.rt_file is None:
            return "No ROOT file is currently loaded"
        else:
            rt_info += f"File: {self.filename} is opened\n"

        rt_info += self._directory_tree()

        return rt_info

    def __getitem__(self, key):

        return Tree(self.rt_file[key])

    def _directory_tree(self, max_depth=None):
        """Return a pretty string of the ROOT file structure; limit depth with max_depth (None = no limit)."""
        if self.rt_file is None:
            return "No ROOT file loaded"

        lines = []

        def _recurse(directory, prefix="", level=0):
            if max_depth is not None and level >= max_depth:
                return
            items = list(directory.items(recursive=False))
            for idx, (key, obj) in enumerate(items):
                name = key.split(";")[0]
                connector = "├── " if idx < len(items) - 1 else "└── "
                
                lines.append(f"{prefix}{connector}{name}")
                
                if isinstance(obj, uproot.reading.ReadOnlyDirectory):
                    child_prefix = prefix + ("│   " if idx < len(items) - 1 else "    ")
                    _recurse(obj, child_prefix, level + 1)

        _recurse(self.rt_file)
        return "\n".join(lines) if lines else "(empty file)"

    def load(self, filename):
        self.filename = filename
        try:
            print(f"Loading ROOT file: {filename}")
            self.rt_file = uproot.open(filename)
            print(self.rt_file)
        except Exception as e:
            print(f"Error loading file: {e}")
            self.rt_file = None









if __name__ == "__main__":
    # rt = ROOT("/home/easonfu/pyproj/FUCKROOT/large_10G.root")
    rt = ROOT("/home/easonfu/pyproj/FUCKROOT/tuple_reco_baseline.root")
    print(rt)

    rt_layer1 = rt["FakeClusteringUP_46bab5f1/layer0;1"]
    # rt_layer1 = rt["tree"]


    print(rt_layer1)
    print(rt_layer1.keys())

    fig, ax = plt.subplots(subplot_kw={"projection": "hep"})

    # ax2 = ax.twinx()
    # normalize
    # rt_layer1.hist1d("var_1", errors=True, ax=ax)
    # rt_layer1.hist1d("x", errors=True, ax=ax)
    # rt_layer1.hist1d("y", errors=True, ax=ax)
    rt_layer1.hist1d("z", errors=True, ax=ax, bins=100)

    ax.set_xlabel("X Label")
    # ax.set_ylabel("Y Label")

    plt.show()

    # for I in range(10):
    #     t0 = time.time()
    #     print(f"Time taken: {time.time() - t0:.2f} seconds")






