# AUTOGENERATED FILE! PLEASE DON'T EDIT HERE. EDIT THE SOURCE NOTEBOOKS INSTEAD
from .callbacks import Callback, Callbacks, Cbs
import k1lib, time, math, logging, numpy as np
from functools import partial; plt = k1lib.dep("matplotlib.pyplot")
try: import torch; from torch import nn; hasTorch = True
except: hasTorch = False
__all__ = ["Profiler"]
if hasTorch:
    import k1lib.callbacks.profilers as ps
    ComputationProfiler = ps.computation.ComputationProfiler
    IOProfiler = ps.io.IOProfiler
    MemoryProfiler = ps.memory.MemoryProfiler
    TimeProfiler = ps.time.TimeProfiler
else:
    class ComputationProfiler: pass
    class IOProfiler: pass
    class MemoryProfiler: pass
    class TimeProfiler: pass
@k1lib.patch(Cbs)
class Profiler(Callback):                                                        # Profiler
    """Profiles memory, time, and computational complexity of the network. See over
:mod:`k1lib.callbacks.profilers` for more details on each of these profilers"""  # Profiler
    def __init__(self):                                                          # Profiler
        super().__init__(); self.clear(); self.dependsOn=["Recorder"]            # Profiler
    def clear(self):                                                             # Profiler
        """Clears every child profilers"""                                       # Profiler
        self._mpCache=None; self._tpCache=None                                   # Profiler
        self._cpCache=None; self._ioCache=None                                   # Profiler
    def _memory(self): # do this to quickly debug, cause if not, Callback will just raise AttributeError on .memory # Profiler
        if self._mpCache != None: return self._mpCache                           # Profiler
        with self.cbs.context():                                                 # Profiler
            mp = MemoryProfiler(); self.cbs.add(mp)                              # Profiler
            mp._run(); self._mpCache = mp; return mp                             # Profiler
    @property                                                                    # Profiler
    def memory(self) -> MemoryProfiler:                                          # Profiler
        """Gets :class:`~k1lib.callbacks.profilers.memory.MemoryProfiler`"""     # Profiler
        return self._memory()                                                    # Profiler
    def _computation(self):                                                      # Profiler
        if self._cpCache != None: return self._cpCache                           # Profiler
        with self.cbs.context():                                                 # Profiler
            cp = ComputationProfiler(self); self.cbs.add(cp)                     # Profiler
            cp._run(); self._cpCache = cp; return cp                             # Profiler
    @property                                                                    # Profiler
    def computation(self) -> ComputationProfiler:                                # Profiler
        """Gets :class:`~k1lib.callbacks.profilers.computation.ComputationProfiler`""" # Profiler
        return self._computation()                                               # Profiler
    def _time(self):                                                             # Profiler
        if self._tpCache != None: return self._tpCache                           # Profiler
        with self.cbs.context():                                                 # Profiler
            tp = TimeProfiler(); self.cbs.add(tp)                                # Profiler
            tp._run(); self._tpCache = tp; return tp                             # Profiler
    @property                                                                    # Profiler
    def time(self) -> TimeProfiler:                                              # Profiler
        """Gets :class:`~k1lib.callbacks.profilers.time.TimeProfiler`"""         # Profiler
        return self._time()                                                      # Profiler
    def _io(self):                                                               # Profiler
        if self._ioCache != None: return self._ioCache                           # Profiler
        with self.cbs.context():                                                 # Profiler
            io = IOProfiler(); self.cbs.add(io)                                  # Profiler
            io._run(); self._ioCache = io; return io                             # Profiler
    @property                                                                    # Profiler
    def io(self) -> IOProfiler:                                                  # Profiler
        """Gets :class:`~k1lib.callbacks.profilers.io.IOProfiler`"""             # Profiler
        return self._io()                                                        # Profiler
    def __repr__(self):                                                          # Profiler
        return f"""{self._reprHead}, can...
- p.memory: to profile module memory requirements
- p.time: to profile module execution times
- p.computation: to estimate module computation
- p.io: to get input and output shapes of
{self._reprCan}"""                                                               # Profiler
