# AUTOGENERATED FILE! PLEASE DON'T EDIT
from k1lib.callbacks import Callback, Cbs
import k1lib, numpy as np; from torch import nn
from k1lib import cli
_spacing = lambda s: f"{s}   "; # inserted at end of everything, if that element existed
_lcomp = 14; _lp1 = 8; _lp2 = 15; _lp3 = 14
class ComputationData:
    def __init__(self, cProfiler, mS:k1lib.selector.ModuleSelector):
        self.cProfiler = cProfiler; self.mS = mS; self.flop = 0
        self.handle = None; self.hook()
        self.flops = 0; self.tS = None # corresponding time selector
    def hook(self):
        def hk(m, i, o):
            i = k1lib.squeeze(i)
            if isinstance(m, nn.Linear): self.flop += i.numel() * m.out_features
            elif isinstance(m, nn.Conv2d):
                self.flop += m.out_channels * i.shape.numel() * np.prod(m.kernel_size)
            elif isinstance(m, (nn.LeakyReLU, nn.ReLU, nn.Sigmoid)):
                self.flop += i.numel()
        self.handle = self.mS.nn.register_forward_hook(hk)
    def unhook(self):
        self.cProfiler.totalFlop += self.flop; self.handle.remove()
    def __getstate__(self):
        answer = dict(self.__dict__)
        del answer["mS"]; del answer["cProfiler"]; return answer
    def __setstate__(self, state): self.__dict__.update(dict(state))
    def __str__(self):
        if self.flop <= 0: return ""
        a = _spacing(f"{k1lib.fmt.comp(self.flop)}".ljust(_lcomp))
        b = _spacing(f"{round(100 * self.flop / self.cProfiler.totalFlop)}%".rjust(_lp1))
        c = ""
        if self.cProfiler._tpAvailable:
            self.flops = self.flop / self.tS.data.time
            c = _spacing(f"{k1lib.fmt.compRate(self.flops)}".ljust(_lp2))
        d = ""
        if self.cProfiler.selected:
            if "_compProf_" in self.mS:
                d = f"{round(100 * self.flop / self.cProfiler.selectedTotalFlop)}%"
            d = _spacing(d.rjust(_lp3))
        return f"{a}{b}{c}{d}"
class ComputationProfiler(Callback):
    """Profiles computation. Only provide reports on well known
layers only, and thus can't really be universal. Example::

    l = k1lib.Learner.sample()
    l.cbs.add(Cbs.Profiler())
    # views table
    l.Profiler.computation
    # views table highlighted
    l.Profiler.computation.css("#lin1 > #lin")
"""
    def __init__(self, profiler:"Profiler"):
        super().__init__(); self.profiler = profiler
    def startRun(self):
        if not hasattr(self, "selector"): # if no selectors found
            self.selector = self.l.model.select("")
        for m in self.selector.modules(): m.data = ComputationData(self, m)
        self.selector.displayF = lambda m: (k1lib.fmt.txt.red if "_compProf_" in m else k1lib.fmt.txt.identity)(m.data)
        self.totalFlop = 0; self.selectedTotalFlop = None
    @property
    def selected(self): return self.selectedTotalFlop != None
    @property
    def _tpAvailable(self) -> bool:
        """Whether TimeProfiler's results are available"""
        try: self.profiler._time(); return True
        except Exception as e: return False
    def startStep(self): return True
    def _run(self):
        """Runs everything"""
        with self.cbs.context(), self.cbs.suspendEval():
            self.cbs.add(Cbs.Cpu()); self.l.run(1, 1)
        for m in self.selector.modules(): m.data.unhook()
    def detached(self): # time profiler integration, so that flops can be displayed
        if self._tpAvailable:
            for cS, tS in zip(self.selector.modules(), self.profiler.time.selector.modules()):
                cS.data.tS = tS # injecting dependency
    def css(self, css:str):
        """Selects a small part of the network to highlight. See also: :mod:`k1lib.selector`."""
        self.selector.parse(k1lib.selector.preprocess(css, "_compProf_"))
        self.selectedTotalFlop = 0
        for m in self.selector.modules():
            if "_compProf_" in m:
                self.selectedTotalFlop += m.data.flop
        print(self.__repr__())
        self.selector.clearProps(); self.selectedTotalFlop = None
    def __repr__(self):
        header = _spacing("computation".ljust(_lcomp))
        header += _spacing("% total".rjust(_lp1))
        header += _spacing("rate".ljust(_lp2)) if self._tpAvailable else ""
        header += _spacing("% selected".rjust(_lp3)) if self.selected else ""
        footer = _spacing(f"{k1lib.fmt.comp(self.totalFlop)}".ljust(_lcomp))
        footer += _spacing("".rjust(_lp1))
        footer += _spacing("".ljust(_lp2)) if self._tpAvailable else ""
        footer += _spacing(f"{k1lib.fmt.comp(self.selectedTotalFlop)}".rjust(_lp3)) if self.selected else ''
        footer = ("Total", footer)
        c = self.selector.__repr__(intro=False, header=header, footer=footer).split("\n") | cli.tab() | cli.join("\n")
        return f"""ComputationProfiler:\n{c}

The "rate" column will appear if integration with Profiler.time is
possible, showing actual ops/s

Can...
- cp.css("..."): highlights a particular part of the network
- cp.selector: to get internal k1lib.ModuleSelector object"""