import time

import pandas   as pd
from ..HEP_PLOT import HepAxes  # noqa: F401
import matplotlib.pyplot as plt

class Tree:
    def __init__(self, rt):
        self.rt = rt
        self._rt_cache = {}
        self._keys = rt.keys()

        self.data_len = rt.num_entries if hasattr(rt, "num_entries") else 0
        if self.data_len < 1000:
            self._rt_df = rt.arrays(library="pd")
        else:
            self._rt_df = pd.concat([
                rt.arrays(entry_start=0, entry_stop=500, library="pd"),
                rt.arrays(entry_start=self.data_len-500, entry_stop=self.data_len, library="pd")
            ])

    def __str__(self):
        return str(self._rt_df)


    def keys(self):
        return self._keys


    def hist1d(self, keys, histtype="step", bins=40, label=None, stat=True, errors=True, **kwargs):
        '''
        :param keys:
        :param kwargs:
        :return:
        '''
        if keys not in self._keys:
            raise KeyError(f"Key ['{keys}'] not in the tree.")


        if keys in self._rt_cache:
            data = self._rt_cache[keys]
        else:
            data = self.rt[keys].array(library="np")
            self._rt_cache[keys] = data


        kwargs.setdefault("histtype", histtype)
        kwargs.setdefault("bins", bins)
        kwargs.setdefault("label", label or keys)

        fig, ax = plt.subplots()
        ax.hist(data, stat=stat, errors=errors, **kwargs)
        ax.set_xlabel(keys)
        ax.set_ylabel("Entries")
        ax.grid(True)
        ax.legend()
        ax.set_title(f"Histogram of {keys}")
        plt.tight_layout()
        plt.show(block=False)


        return fig
    
    def __getitem__(self, key):
        if key not in self._keys:
            raise KeyError(f"Key ['{key}'] not in the tree.")
        
        if key in self._rt_cache:
            return self._rt_cache[key]
        else:
            data = self.rt[key].array(library="np")
            self._rt_cache[key] = data
            return data




if __name__ == "__main__":
    pass
