# AUTOGENERATED FILE! PLEASE DON'T EDIT HERE. EDIT THE SOURCE NOTEBOOKS INSTEAD
from .callbacks import Callback, Callbacks, Cbs
import k1lib; plt = k1lib.dep("matplotlib.pyplot")
from functools import partial
from typing import List, Tuple, Callable, Union
try: import torch; import torch.nn as nn; hasTorch = True
except:
    torch = k1lib.Object().withAutoDeclare(lambda: type("RandomClass", (object, ), {}))
    nn = k1lib.Object().withAutoDeclare(lambda: type("RandomClass", (object, ), {})); hasTorch = False
__all__ = ["HookParam"]
class ParamData(k1lib.Object):                                                   # ParamData
    def __init__(self):                                                          # ParamData
        super().__init__()                                                       # ParamData
        self.means = []; self.stds = []                                          # ParamData
        self.mins = []; self.maxs = []                                           # ParamData
    def update(self, torchParam:nn.Parameter):                                   # ParamData
        self.means.append(torchParam.mean().item())                              # ParamData
        self.stds.append(torchParam.std().item())                                # ParamData
        self.mins.append(torchParam.min().item())                                # ParamData
        self.maxs.append(torchParam.max().item())                                # ParamData
    def __len__(self): return len(self.means)                                    # ParamData
    def __repr__(self):                                                          # ParamData
        return f"""Param's saved data. Use...
- d.means: to get list of means
- d.stds: to get list of means
- d.mins: to get list of mins
- d.maxs: to get list of maxs"""                                                 # ParamData
class Param:                                                                     # Param
    def __init__(self, name:str, torchParam:nn.Parameter):                       # Param
        self.name = name                                                         # Param
        self.torchParam = torchParam                                             # Param
        self.data = ParamData()                                                  # Param
        self.every = k1lib.Every(3)                                              # Param
    def update(self):                                                            # Param
        if self.every(): self.data.update(self.torchParam.detach())              # Param
    def __repr__(self):                                                          # Param
        return f"""Param `{self.name}`. Use...
- p.torchParam: to get actual underlying parameter
- p.data: to get data stored
- cb.plot(): to quickly look at everything"""                                    # Param
@k1lib.patch(Cbs)                                                                # Param
class HookParam(Callback):                                                       # HookParam
    """Records means and stds of all parameters"""                               # HookParam
    def __init__(self):                                                          # HookParam
        ""                                                                       # HookParam
        super().__init__(); self.params:List[Param] = []                         # HookParam
    def __getitem__(self, idx:Union[int, slice]):                                # HookParam
        if type(idx) == int: return self.params[idx]                             # HookParam
        answer = HookParam(); answer.params = self.params[idx]; return answer    # HookParam
    def __len__(self): return len(self.params)                                   # HookParam
    def _selected(self, paramName:str):                                          # HookParam
        splits = paramName.split(".")                                            # HookParam
        try:                                                                     # HookParam
            mS = self.l.selector                                                 # HookParam
            for split in splits[:-1]: mS = mS[split]                             # HookParam
            return "HookParam" in mS and hasattr(mS, splits[-1])                 # HookParam
        except KeyError: return False                                            # HookParam
    def startRun(self):                                                          # HookParam
        if len(self) == 0: # set things up first time only                       # HookParam
            self.params = [Param(k, v) for k, v in self.l.model.named_parameters() if self._selected(k)] # HookParam
    def startBatch(self): [param.update() for param in self.params]              # HookParam
    def css(self, css:str):                                                      # HookParam
        """Creates a new HookParam object with selected modules. May be useful
for displaying a subset of the recorded data"""                                  # HookParam
        oldSelector = self.l.selector; answer = HookParam()                      # HookParam
        self.l.selector = k1lib.selector.select(self.l.model, css)               # HookParam
        answer.params = [param for param in self.params if self._selected(param.name)] # HookParam
        self.l.selector = oldSelector; return answer                             # HookParam
    def __repr__(self):                                                          # HookParam
        s = f", {len(self[0].data)} means and stds each" if len(self) > 0 else "" # HookParam
        names = "\n".join([f"  {i}. {p.name}" for i, p in enumerate(self)])      # HookParam
        return f"""{super()._reprHead}: {len(self)} params{s}:\n{names}\n
Use...
- p.plot(): to quickly look at everything
- p[i]: to view a single param
- p[a:b]: to get a new HookParam with selected params
- p.css("..."): to select a specific subset of modules only
{super()._reprCan}"""                                                            # HookParam
def plotF(params:Union[HookParam, Param, List[Param]], rangeSlice:slice):        # plotF
    if type(params) == Param: params = [params]                                  # plotF
    fields = params[0].data.state.keys(); step = rangeSlice.step or 1            # plotF
    fig, axes = plt.subplots(2, 2, figsize=(10, 6), dpi=100)                     # plotF
    axes = axes.flatten()                                                        # plotF
    for field, ax in zip(fields, axes):                                          # plotF
        for param in params:                                                     # plotF
            fieldData = param.data[field]                                        # plotF
            r = k1lib.Range(len(fieldData))[rangeSlice]                          # plotF
            ax.plot(r.range_[::step], fieldData[r.slice_][::step])               # plotF
        ax.set_title(field.capitalize())                                         # plotF
    plt.figlegend([p.name for p in params], loc='right')                         # plotF
@k1lib.patch(HookParam)                                                          # plotF
@k1lib.patch(Param)                                                              # plotF
def plot(self): return k1lib.viz.SliceablePlot(partial(plotF, self))             # plot
