# AUTOGENERATED FILE! PLEASE DON'T EDIT
from .callbacks import Callback, Callbacks, Cbs
import k1lib, time
try: import torch; hasTorch = True
except: hasTorch = False
__all__ = ["BatchLimit", "EpochLimit", "TimeLimit", "CancelOnExplosion",
           "CancelOnLowLoss", "CancelOnHighAccuracy", "CancelOnOverfit", "DontTrain",
           "GradientClipping", "GradientClippingNorm", "TrainOnly", "ValidOnly"]
@k1lib.patch(Cbs)
class BatchLimit(Callback):
    """Cancels the epoch after executed certain number of batches"""
    def __init__(self, limit:int):
        super().__init__(); self.order = 25
        self.limit = limit if limit != None else float("inf")
    def startEpoch(self): self.currentBatch = 0
    def startBatch(self):
        if self.currentBatch >= self.limit:
            raise k1lib.CancelEpochException(f"Batch {self.limit} reached")
    def endBatch(self): self.currentBatch += 1
@k1lib.patch(Cbs)
class EpochLimit(Callback):
    """Cancels the run after executed certain number of epochs"""
    def __init__(self, limit:int):
        super().__init__(); self.order = 25
        self.limit = limit if limit != None else float("inf")
    def startRun(self): self.currentEpoch = 0
    def startEpoch(self):
        if self.currentEpoch >= self.limit:
            raise k1lib.CancelRunException(f"Epoch {self.limit} reached!")
    def endEpoch(self): self.currentEpoch += 1
@k1lib.patch(Cbs)
class TimeLimit(Callback):
    """Cancels the run after a certain number of seconds have passed"""
    def __init__(self, seconds=30):
        super().__init__(); self.seconds = seconds if seconds != None else float("inf"); self.order = 25
    def startRun(self): self.startTime = time.time()
    def startBatch(self):
        if time.time() - self.startTime > self.seconds:
            raise k1lib.CancelRunException(f"Takes more than {self.seconds} seconds!")
@k1lib.patch(Cbs)
class CancelOnExplosion(Callback):
    """Cancels the run if any of the parameters are larger than a certain limit"""
    def __init__(self, limit:float=1e6):
        super().__init__(); self.order = 25
        self.limit = limit; self.triggered = False
    def startRun(self): self.triggered = False
    def startBatch(self):
        for p in self.l.model.parameters():
            o = p.detach()
            if o.max().float() > self.limit or o.min().float() < -self.limit:
                self.triggered = True
                raise k1lib.CancelRunException("Explosion detected!")
    def __repr__(self):
        return f"""{self._reprHead}, use...
- cb.triggered: to see if there was an explosion on the last run
- cb.progress: to see current progress at explosion time
{self._reprCan}"""
@k1lib.patch(Cbs)
class CancelOnLowLoss(Callback):
    " "
    def __init__(self, loss:float, epochMode:bool=False):
        """Cancels the run if loss is lower than amount specified.
Original class: :class:`~k1lib.callbacks.limits.CancelOnLowLoss`

:param epochMode: False if use batch loss, True if use valid epoch loss"""
        super().__init__(); self.order = 25; self.dependsOn = ["Loss"]
        self.loss = loss; self.epochMode = epochMode
    def startRun(self):
        if not hasattr(self.l.cbs, "Loss"):
            raise AttributeError("Learner does not have required `Loss` callback")
        self.v = self.cbs.Loss.valid; self.ve = self.cbs.Loss.epoch.valid # List[int]
    def endBatch(self):
        if self.epochMode:
            if len(self.ve) > 0 and self.ve[-1] < self.loss:
                raise k1lib.CancelRunException(f"Low loss {self.loss} ({self.ve[-3:]} actual) achieved!")
        elif len(self.v) and self.v[-1] < self.loss:
            raise k1lib.CancelRunException(f"Low loss {self.loss} ({self.v[-3:]} actual) achieved!")
@k1lib.patch(Cbs)
class CancelOnHighAccuracy(Callback):
    """Cancels the run if accuracy is higher than the amount specified"""
    def __init__(self, accuracy:float):
        super().__init__(); self.order = 25
        self.accuracy = accuracy; self.dependsOn = ["Accuracy"]
    def endBatch(self):
        if not hasattr(self.l, "Accuracy"): raise AttributeError("Learner does not have `Accuracy` callback")
        a = self.l.Accuracy.valid[-1]
        if a > self.accuracy:
            raise k1lib.CancelRunException(f"High accuracy {self.accuracy} ({a} actual) achieved!")
@k1lib.patch(Cbs)
class CancelOnOverfit(Callback):
    def __init__(self, ratio:float=1.2, alpha:float=0.99, after:int=10):
        """Cancels the run if overfit is detected.

:param ratio: Max ratio between the lowest loss and the current loss before cancelling the run
:param alpha: Moving average's alpha, used for both minLoss and loss estimates
:param after: After how many epochs should the overfit detection be activated?"""
        super().__init__(); self.ratio = ratio
        self.minLoss = k1lib.MovingAvg(alpha=alpha, debias=True)
        self.loss = k1lib.MovingAvg(alpha=alpha, debias=True)
        self.count = 0; self.after = after
    def startRun(self): self.count = 0
    def endEpoch(self): self.count += 1
    def endBatch(self):
        if not self.l.model.training:
            loss = self.l.loss; self.loss(loss)
            if self.loss.value < self.minLoss.value or self.minLoss.value == 0: self.minLoss(self.loss.value)
            if self.count > self.after and self.loss.value > self.minLoss.value * self.ratio:
                raise k1lib.CancelRunException(f"Overfit detected! Smoothed min loss: {self.minLoss.value}, loss: {loss}")
@k1lib.patch(Cbs)
class DontTrain(Callback):
    """Don't allow the network to train at all"""
    def startBackward(self): return True
    def startStep(self): return True
if hasTorch:
    from torch.nn.utils import clip_grad_value_
    @k1lib.patch(Cbs)
    class GradientClipping(Callback):
        """Clips gradient to a specific max value"""
        def __init__(self, value:float): super().__init__(); self.value = value
        def startStep(self):
            clip_grad_value_(self.l.model.parameters(), self.value)
else:
    class GradientClipping(Callback): pass
if hasTorch:
    from torch.nn.utils import clip_grad_norm_
    @k1lib.patch(Cbs)
    class GradientClippingNorm(Callback):
        """Clips gradient to a specific max_norm value. Can choose to lump
all params together or do each separately.

See also: :class:`~k1lib.callbacks.limits.GradientClipping` callback."""
        def __init__(self, max_norm:float, each:bool=True):
            super().__init__(); self.max_norm = max_norm; self.each = each
        def startStep(self):
            if self.each:
                for m in self.l.model.parameters():
                    clip_grad_norm_(m, self.max_norm)
            else: clip_grad_norm_(self.l.model.parameters(), self.max_norm)
else:
    class GradientClippingNorm(Callback): pass
@k1lib.patch(Cbs)
class TrainOnly(Callback):
    " "
    def __init__(self, cb):
        """Only executes specified callback when training. This modifies the callback's
``suspended`` variable, so it may interfere with :meth:`k1lib.callbacks.callbacks.Callbacks.suspend`
by setting it to different values while in the context."""
        super().__init__(); self.cb = cb
    def startBatch(self):
        self.cb.suspended = not self.l.model.training
@k1lib.patch(Cbs)
class ValidOnly(Callback):
    " "
    def __init__(self, cb):
        """Same as :class:`TrainOnly`, but only executes specified callback when doing
validation."""
        super().__init__(); self.cb = cb
    def startBatch(self):
        self.cb.suspended = self.l.model.training