# AUTOGENERATED FILE! PLEASE DON'T EDIT HERE. EDIT THE SOURCE NOTEBOOKS INSTEAD
from .callbacks import Callback, Callbacks, Cbs
import k1lib, os, warnings
from typing import Callable
try: import torch; hasTorch = True
except: hasTorch = False
__all__ = ["Autosave", "DontTrainValid", "InspectLoss", "ModifyLoss", "Cpu", "Cuda",
           "DType", "InspectBatch", "ModifyBatch", "InspectOutput", "ModifyOutput",
           "Beep", "OnProgress"]
@k1lib.patch(Cbs)
class Autosave(Callback):                                                        # Autosave
    """Autosaves 3 versions of the network to disk"""                            # Autosave
    def __init__(self): super().__init__(); self.order = 23                      # Autosave
    def endRun(self):                                                            # Autosave
        os.system("mv autosave-1.pth autosave-0.pth")                            # Autosave
        os.system("mv autosave-2.pth autosave-1.pth")                            # Autosave
        self.l.save("autosave-2.pth")                                            # Autosave
@k1lib.patch(Cbs)                                                                # Autosave
class DontTrainValid(Callback):                                                  # DontTrainValid
    """If is not training, then don't run m.backward() and opt.step().
The core training loop in k1lib.Learner don't specifically do this,
cause there may be some weird cases where you want to also train valid."""       # DontTrainValid
    def _common(self):                                                           # DontTrainValid
        if not self.l.model.training: return True                                # DontTrainValid
    def startBackward(self): return self._common()                               # DontTrainValid
    def startStep(self): return self._common()                                   # DontTrainValid
@k1lib.patch(Cbs)                                                                # DontTrainValid
class InspectLoss(Callback):                                                     # InspectLoss
    """Expected `f` to take in 1 float."""                                       # InspectLoss
    def __init__(self, f): super().__init__(); self.f = f; self.order = 15       # InspectLoss
    def endLoss(self): self.f(self.loss.detach())                                # InspectLoss
@k1lib.patch(Cbs)                                                                # InspectLoss
class ModifyLoss(Callback):                                                      # ModifyLoss
    """Expected `f` to take in 1 float and return 1 float."""                    # ModifyLoss
    def __init__(self, f): super().__init__(); self.f = f; self.order = 13       # ModifyLoss
    def endLoss(self): self.l.loss = self.f(self.loss)                           # ModifyLoss
@k1lib.patch(Cbs)                                                                # ModifyLoss
class Cuda(Callback):                                                            # Cuda
    """Moves batch and model to the default GPU"""                               # Cuda
    def startRun(self): self.l.model.cuda()                                      # Cuda
    def startBatch(self):                                                        # Cuda
        try: self.l.xb = self.l.xb.cuda()                                        # Cuda
        except Exception as e: warnings.warn(f"xb can't be moved to the GPU: {e}") # Cuda
        try: self.l.yb = self.l.yb.cuda()                                        # Cuda
        except Exception as e: warnings.warn(f"yb can't be moved to the GPU: {e}") # Cuda
@k1lib.patch(Cbs)                                                                # Cuda
class Cpu(Callback):                                                             # Cpu
    """Moves batch and model to CPU"""                                           # Cpu
    def startRun(self): self.l.model.cpu()                                       # Cpu
    def startBatch(self):                                                        # Cpu
        self.l.xb = self.l.xb.cpu()                                              # Cpu
        self.l.yb = self.l.yb.cpu()                                              # Cpu
@k1lib.patch(Cbs)                                                                # Cpu
class DType(Callback):                                                           # DType
    """Moves batch and model to a specified data type"""                         # DType
    def __init__(self, dtype): super().__init__(); self.dtype = dtype            # DType
    def startRun(self): self.l.model = self.l.model.to(self.dtype)               # DType
    def startBatch(self):                                                        # DType
        self.l.xb = self.l.xb.to(self.dtype)                                     # DType
        self.l.yb = self.l.yb.to(self.dtype)                                     # DType
@k1lib.patch(Cbs)                                                                # DType
class InspectBatch(Callback):                                                    # InspectBatch
    """Expected `f` to take in 2 tensors."""                                     # InspectBatch
    def __init__(self, f:callable): super().__init__(); self.f = f; self.order = 15 # InspectBatch
    def startBatch(self): self.f(self.l.xb, self.l.yb)                           # InspectBatch
@k1lib.patch(Cbs)                                                                # InspectBatch
class ModifyBatch(Callback):                                                     # ModifyBatch
    """Modifies xb and yb on the fly. Expected `f`
    to take in 2 tensors and return 2 tensors."""                                # ModifyBatch
    def __init__(self, f): super().__init__(); self.f = f; self.order = 13       # ModifyBatch
    def startBatch(self): self.l.xb, self.l.yb = self.f(self.l.xb, self.l.yb)    # ModifyBatch
@k1lib.patch(Cbs)                                                                # ModifyBatch
class InspectOutput(Callback):                                                   # InspectOutput
    """Expected `f` to take in 1 tensor."""                                      # InspectOutput
    def __init__(self, f): super().__init__(); self.f = f; self.order = 15       # InspectOutput
    def endPass(self): self.f(self.y)                                            # InspectOutput
@k1lib.patch(Cbs)                                                                # InspectOutput
class ModifyOutput(Callback):                                                    # ModifyOutput
    """Modifies output on the fly. Expected `f` to take
in 1 tensor and return 1 tensor"""                                               # ModifyOutput
    def __init__(self, f): super().__init__(); self.f = f; self.order = 13       # ModifyOutput
    def endPass(self): self.l.y = self.f(self.y)                                 # ModifyOutput
@k1lib.patch(Cbs)                                                                # ModifyOutput
class Beep(Callback):                                                            # Beep
    """Plays a beep sound when the run is over"""                                # Beep
    def endRun(self): k1lib.beep()                                               # Beep
@k1lib.patch(Cbs)                                                                # Beep
class OnProgress(Callback):                                                      # OnProgress
    """Triggers a specific function once when reached a certain
progress.

:param f: function that takes in a :class:`~k1lib.Learner` object
:param progress: progress to trigger the function, from 0 to 1"""                # OnProgress
    def __init__(self, f:Callable[["k1lib.Learner"], None], progress:float):     # OnProgress
        super().__init__(); self.f = f; self.progress = progress; self.ran = False # OnProgress
    def startRun(self): self.ran = False                                         # OnProgress
    def startBatch(self):                                                        # OnProgress
        if (not self.ran) and self.l.progress > self.progress:                   # OnProgress
            self.ran = True; self.f(self.l)                                      # OnProgress
