# AUTOGENERATED FILE! PLEASE DON'T EDIT
from k1lib.callbacks import Callback, Cbs
import k1lib; from k1lib import cli
_li = 30
class IOData:
    def __init__(self, ioProfiler, mS:k1lib.selector.ModuleSelector):
        self.ioProfiler = ioProfiler; self.mS = mS
        self.iS = None; self.oS = None
        self.handle = None; self.hook()
    def hook(self):
        def hk(m, i, o):
            self.iS = list(k1lib.squeeze(i, True).shape)
            self.oS = list(k1lib.squeeze(o, True).shape)
        self.handle = self.mS.nn.register_forward_hook(hk)
    def unhook(self): self.handle.remove()
    def __getstate__(self):
        answer = dict(self.__dict__)
        del answer["mS"]; del answer["ioProfiler"]; return answer
    def __setstate__(self, state): self.__dict__.update(dict(state))
    def __str__(self):
        a = f"{self.iS}".ljust(_li); b = f"{self.oS}".ljust(_li)
        return f"{a}{b}"
class IOProfiler(Callback):
    """Gets input and output shapes of each layer.
Example::
    
    l = k1lib.Learner.sample()
    l.cbs.add(Cbs.Profiler())
    # views table
    l.Profiler.io
    # views table highlighted
    l.Profiler.io.css("#lin1")
"""
    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 = IOData(self, m)
        self.selector.displayF = lambda m: (k1lib.fmt.txt.red if "_ioProf_" in m else k1lib.fmt.txt.identity)(m.data)
    def startStep(self): return True
    def _run(self):
        """Runs everything"""
        with self.cbs.suspendEval(): self.l.run(1, 1)
        for m in self.selector.modules(): m.data.unhook()
    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, "_ioProf_"))
        print(self.__repr__()); self.selector.clearProps()
    def __repr__(self):
        header = "input shape".ljust(_li) + "output shape".ljust(_li)
        c = self.selector.__repr__(intro=False, header=header).split("\n") | cli.tab() | cli.join("\n")
        return f"""IOProfiler:\n{c}

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