# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/11_l2r.callbacks.ipynb.

# %% auto 0
__all__ = ['TrainEval', 'AvgSmoothMetric', 'TrackResults', 'Monitor', 'XParamScheduler', 'XLRFinder', 'ProgressBarCallback',
           'SaveCallBack', 'InterruptCallback', 'CSVLog', 'AMPMode', 'TestCallback', 'MemoryErrorTest',
           'MixedPrecision', 'GradientClip']

# %% ../../nbs/11_l2r.callbacks.ipynb 2
from fastai.torch_imports import *
from fastai.torch_core import *
from fastai.callback.core import *
from fastcore.all import *
from ..imports import *
from ..metrics import *

# %% ../../nbs/11_l2r.callbacks.ipynb 6
class TrainEval(Callback):
    "`Callback` that tracks the number of iterations done and properly sets training/eval mode"
    order, run_valid = -10, False
    
    def before_fit(self):
        "Set the iter and epoch counters to 0"
        self.learn.epoch = 0
        self.learn.train_iter,self.learn.pct_train = 0,0.
        
    def after_batch(self):
        "Update the iter counter (in training mode)"
        if not self.model.training: return
        self.learn.pct_train += 1./(self.n_iter*self.n_epochs)
        self.learn.train_iter += 1
        
    def before_train(self):
        "Set the model to training mode"
        self.learn.pct_train=self.epoch/self.n_epochs
        self.model.train()
        self.learn.training=True
        
    def before_validate(self):
        "Set the model to validation mode"
        self.model.eval()
        self.learn.training=False

# %% ../../nbs/11_l2r.callbacks.ipynb 8
from fastai.learner import Metric

# %% ../../nbs/11_l2r.callbacks.ipynb 9
class AvgSmoothMetric(Metric):
    "Smooth average of the losses (exponentially weighted with `beta`)"
    def __init__(self, beta=0.98): self.beta = beta
    def reset(self):               self.count,self.val = 0,tensor(0.)
    def accumulate(self, learn):
        self.count += 1
        self.val = torch.lerp(to_detach(learn.moi.mean()), self.val, self.beta)
    @property
    def value(self): return self.val/(1-self.beta**self.count)

# %% ../../nbs/11_l2r.callbacks.ipynb 11
from fastprogress.core import format_time  
import time

# %% ../../nbs/11_l2r.callbacks.ipynb 12
class TrackResults(Callback):
    def __init__(self, train_metrics=False, beta=0.98): 
        store_attr()
        self.names = ['loss', 'ndcg', 'ndcg_at_6', 'acc']
        self.metric_names = L(self.names).map(lambda o: 'val_'+o)
        if train_metrics: self.metric_names = L(self.names).map(lambda o: 'train_'+o) + self.metric_names
        self.metric_names.append('time')
        self.metric_names = L('epoch') + self.metric_names
        self.smooth_moi = AvgSmoothMetric(beta=beta) # to maintain weighted avg of metric of interest 
    
    def before_fit(self): 
        self.lrs, self.mois, self.smooth_mois, self.losses_full, self.grads_full, self.metrics_full = [], [], [], [], defaultdict(list), defaultdict(list) 
        self.ioi = self.names.index(self.learn.save_call_back.monitor
                                  if hasattr(self.learn, 'save_call_back') else 'loss')
        self.smooth_moi.reset()
    
    def before_train(self): self._initialise_metrics()
    
    def before_validate(self): self._initialise_metrics()
        
    def after_train(self):
        self.losses_full.extend(self.losses)
        log = self._compute_epoch_mean()
        if self.train_metrics:
            self.metrics_full['trn'].append(log)
        # print(self.epoch, self.model.training, *log)
        self.log+=log
                
    def after_validate(self):
        log = self._compute_epoch_mean()
        if hasattr(self, 'metrics_full'):
            self.metrics_full['val'].append(log)
        # print(self.epoch if hasattr(self, 'epoch') else 0, self.model.training, *log)
        if getattr(self, 'log', False): self.log+=log

    def before_epoch(self):
        self.start_epoch = time.time()
        self.log = L(getattr(self, 'epoch', 0))

    def after_epoch(self): 
        self.log.append(format_time(time.time() - self.start_epoch))
        # import pdb; pdb.set_trace()
        self.logger(self.log)
            
    def _compute_epoch_mean(self):
        _li = [self.losses, self.ndcgs, self.ndcgs_at_6, self.accs]
        _li = [torch.stack(o) if o else torch.Tensor() for o in _li] 
        [self.losses, self.ndcgs, self.ndcgs_at_6, self.accs] = _li
        log = [round(o.mean().item(), 4) if o.sum() else "NA" for o in _li]
        return log
    
    def _initialise_metrics(self): self.losses, self.ndcgs, self.ndcgs_at_6, self.accs = [], [], [], []

    
    def after_batch(self):
        with torch.no_grad():
            loss = self.loss_func(self.preds, self.xb)
            if loss.mean().isnan(): 
                import pdb; pdb.set_trace()
            self.losses.append(loss.mean())
            if self.model.training:
                if self.train_metrics: 
                    self._compute_metrics()
                    # grab moi
                    self.learn.moi = getattr(self, "losses ndcgs ndcgs_at_6 accs".split()[self.ioi])[-1] # -1 is to get the value for the current batch
                    # do weighted avg
                    self.smooth_moi.accumulate(self.learn)
                    # record `learn` attribute
                    self.learn.smooth_moi = self.smooth_moi.value
                    # record moi attribute in `self`
                    self.mois.append(self.learn.moi)
                    # record smooth_moi attribute in `self`
                    self.smooth_mois.append(self.learn.smooth_moi)
                    # record lr in `self`
                    self.lrs.append(self.opt.hypers[-1]['lr'])
            else: self._compute_metrics()
                        
    def _compute_metrics(self):
        *_, _ndcg, _ndcg_at_k = ndcg(self.preds, self.xb, k=6)
        self.ndcgs.append(_ndcg.mean())
        self.ndcgs_at_6.append(_ndcg_at_k.mean())
        acc = accuracy(self.xb, self.model).mean()
        self.accs.append(acc.mean())
        
    def after_backward(self):
        for name,param in self.model.named_parameters():
            grad = param.grad.data.detach().clone()
            self.grads_full[name].append(grad)

# %% ../../nbs/11_l2r.callbacks.ipynb 13
class Monitor(Callback):
    order = 60
    def __init__(self, monitor='ndcg_at_6', comp=None, min_delta=0., reset_on_fit=False):
        if comp is None: comp = np.greater
        if comp == np.less: min_delta *= -1
        # store_attr()
        self.monitor,self.comp,self.min_delta,self.reset_on_fit,self.best= monitor,comp,min_delta,reset_on_fit,None
       
    def before_fit(self):
        if self.reset_on_fit or self.best is None: self.best = float('inf') if self.comp == np.less else -float('inf')
        assert self.monitor in self.track_results.names
        self.idx = list(self.track_results.names).index(self.monitor)
        
    def after_epoch(self):
        if self.track_results.metrics_full.get('val') is None: return 
        val = self.track_results.metrics_full.get('val')[-1][self.idx]
        if self.comp(val - self.min_delta if not isinstance(val, str) else self.best, self.best): self.best, self.new_best, = val, True
        else: self.new_best = False

# %% ../../nbs/11_l2r.callbacks.ipynb 15
@docs
class XParamScheduler(Callback):
    "Schedule hyper-parameters according to `scheds`"
    order,run_valid = 60,False

    def __init__(self, scheds): self.scheds = scheds
    def before_fit(self): self.hps = {p:[] for p in self.scheds.keys()}
    def before_batch(self): 
        if not self.model.training: return
        self._update_val(self.pct_train)

    def _update_val(self, pct):
        for n,f in self.scheds.items(): self.opt.set_hyper(n, f(pct))

    def after_batch(self):
        if not self.model.training: return
        for p in self.scheds.keys(): self.hps[p].append(self.opt.hypers[-1][p])

    def after_fit(self):
        if hasattr(self.learn, 'track_results') and hasattr(self, 'hps'): self.track_results.hps = self.hps

    _docs = {"before_fit": "Initialize container for hyper-parameters",
             "before_batch": "Set the proper hyper-parameters in the optimizer",
             "after_batch": "Record hyper-parameters of this batch",
             "after_fit": "Save the hyper-parameters in the track_results if there is one"}

# %% ../../nbs/11_l2r.callbacks.ipynb 21
from .learner import L2RLearner
from fastai.callback.schedule import *

# %% ../../nbs/11_l2r.callbacks.ipynb 22
@patch
def fit_one_cycle(self:L2RLearner, n_epoch, lr_max=None, div=25., div_final=1e5, pct_start=0.25, moms=None, cbs=None, start_epoch=0):
    "Fit `self.model` for `n_epoch` using the 1cycle policy."
    self.lr_max = lr_max
    self.opt = getattr(self, 'opt', None)
    if self.opt is None: self.create_opt()
    self.opt.set_hyper('lr', self.lr if lr_max is None else lr_max)
    lr_max = np.array([h['lr'] for h in self.opt.hypers])
    scheds = {'lr': combined_cos(pct_start, lr_max/div, lr_max, lr_max/div_final),
              'mom': combined_cos(pct_start, *(self.moms if moms is None else moms))}
    # self.add_cb(XParamScheduler(scheds)) 
    # pdb.set_trace()
    self.fit(n_epoch, cbs=XParamScheduler(scheds)+L(cbs), start_epoch=start_epoch)

# %% ../../nbs/11_l2r.callbacks.ipynb 24
@patch
def plot_sched(self:TrackResults, keys=None, figsize=None):
    keys = self.hps.keys() if keys is None else L(keys)
    rows,cols = (len(keys)+1)//2, min(2, len(keys))
    figsize = figsize or (6*cols,4*rows)
    _, axs = plt.subplots(rows, cols, figsize=figsize)
    axs = axs.flatten() if len(keys) > 1 else L(axs)
    for p,ax in zip(keys, axs):
        ax.plot(self.hps[p])
        ax.set_ylabel(p)

# %% ../../nbs/11_l2r.callbacks.ipynb 25
@patch
def fit_sgdr(self:L2RLearner, n_cycles, cycle_len, lr_max=None, cycle_mult=2, cbs=None, reset_opt=False, wd=None,
             start_epoch=0):
    "Fit `self.model` for `n_cycles` of `cycle_len` using SGDR."
    self.lr_max = lr_max
    self.opt = getattr(self, 'opt', None)
    if self.opt is None: self.create_opt()
    self.opt.set_hyper('lr', self.lr if lr_max is None else lr_max)
    lr_max = np.array([h['lr'] for h in self.opt.hypers])
    n_epoch = cycle_len * (cycle_mult**n_cycles-1)//(cycle_mult-1)
    pcts = [cycle_len * cycle_mult**i / n_epoch for i in range(n_cycles)]
    scheds = [SchedCos(lr_max, 0) for _ in range(n_cycles)]
    scheds = {'lr': combine_scheds(pcts, scheds)}
    self.fit(n_epoch, cbs=ParamScheduler(scheds)+L(cbs), reset_opt=reset_opt, wd=wd, start_epoch=start_epoch)

# %% ../../nbs/11_l2r.callbacks.ipynb 28
@docs
class XLRFinder(XParamScheduler):
    "Training with exponentially growing learning rate"
    run_after=TrackResults
    def __init__(self, start_lr=1e-7, end_lr=10, num_it=100, stop_div=True):
        if num_it < 6: num_it = 6
        self.scheds = {'lr': [SchedExp(s, e) for (s,e) in zip(start_lr,end_lr)
                             ] if is_listy(start_lr) else SchedExp(start_lr, end_lr)}
        self.num_it,self.stop_div = num_it,stop_div

    def before_fit(self):
        if hasattr(self, 'track_results'):
            moi = self.track_results.names[self.track_results.ioi]
            print(f'Smoothing {moi}')
        super().before_fit()
        path = self.path
        path.mkdir(parents=True, exist_ok=True)
        self.tmp_d = tempfile.TemporaryDirectory(dir=path)
        self.tmp_p = Path(self.tmp_d.name).stem
        self.learn.save(f'{self.tmp_p}/_tmp')
        self.best_moi = -float('inf')

    def before_batch(self): self._update_val(self.train_iter/self.num_it)

    def after_batch(self):
        super().after_batch()
        if self.smooth_moi > self.best_moi: self.best_moi = self.smooth_moi
        if 4*self.smooth_moi < self.best_moi and self.stop_div: raise CancelFitException()
        if self.train_iter >= self.num_it: raise CancelFitException()

    def before_validate(self): raise CancelValidException()

    def after_fit(self):
        self.learn.opt.zero_grad() # Needed before detaching the optimizer for future fits
        tmp_f = self.path/self.tmp_p/'_tmp.pth'
        if tmp_f.exists():
            self.learn.load(f'{self.tmp_p}/_tmp', with_opt=True)
            self.tmp_d.cleanup()
            
    def after_cancel_validate(self): pass

    _docs = {"before_fit": "Initialize container for hyper-parameters and save the model",
             "before_batch": "Set the proper hyper-parameters in the optimizer",
             "after_batch": "Record hyper-parameters of this batch and potentially stop training",
             "after_fit": "Save the hyper-parameters in the recorder if there is one and load the original model",
             "before_validate": "Skip the validation part of training",
             "after_cancel_validate": "pass `CancelValidException`",
             "run_after":""}

# %% ../../nbs/11_l2r.callbacks.ipynb 29
@patch
def plot_xlr_find(self:TrackResults, skip_end=5, return_fig=True, suggestions=None, nms=None, **kwargs):
    "Plot the result of an LR Finder test (won't work if you didn't do `learn.xlr_find()` before)"
    lrs  = self.lrs  if skip_end==0 else self.lrs[:-skip_end]
    mois = L(self.mois if skip_end==0 else self.mois[:-skip_end]).map(Tensor.cpu)
    smooth_mois = L(self.smooth_mois if skip_end==0 else self.smooth_mois[:-skip_end]).map(Tensor.cpu)
    fig, (ax1, ax2) = plt.subplots(1,2)
    ax1.plot(lrs, smooth_mois)
    ax1.set_ylabel("Smoothened MOI")
    ax1.set_xlabel("Learning Rate")
    ax1.set_xscale('log')
    if suggestions:
        colors = plt.rcParams['axes.prop_cycle'].by_key()['color'][1:]
        for (val, idx), nm, color in zip(suggestions, nms, colors):
            ax1.plot(val, idx, 'o', label=nm, c=color)
        ax1.legend(loc='best')
    ax2.plot(range(len(mois)), mois, label='MOI')
    ax2.plot(range(len(smooth_mois)), smooth_mois, label='SMOI')
    ax2.legend(loc='best')


# %% ../../nbs/11_l2r.callbacks.ipynb 30
@patch
def xrl_find(self:L2RLearner, start_lr=1e-5, end_lr=1e-1, num_it=400, stop_div=True, show_plot=True, suggest_funcs=(valley,)):
    "Launch a mock training to find a good learning rate and return suggestions based on `suggest_funcs` as a named tuple"
    n_epochs = num_it//len(self.dls.train) + 1
    cb=XLRFinder(start_lr=start_lr, end_lr=end_lr, num_it=num_it, stop_div=stop_div)
    self.fit(n_epochs, cbs=cb)
    if suggest_funcs is not None:
        lrs, smooth_mois = tensor(self.track_results.lrs[num_it//10:-5]),tensor(self.track_results.smooth_mois[num_it//10:-5])
        nan_idxs = torch.nonzero(torch.isnan(smooth_mois.view(-1)))
        if len(nan_idxs):
            drop_idx = min(nan_idxs)
            lrs = lrs[:drop_idx]
            smooth_mois = smooth_mois[:drop_idx]
        _suggestions, nms = [], []
        for func in tuplify(suggest_funcs):
            nms.append(func.__name__ if not isinstance(func, partial) else func.func.__name__) # deal with partials
            _suggestions.append(func(lrs, smooth_mois, num_it))
            
        # pdb.set_trace()
        SuggestedLRs = collections.namedtuple('SuggestedLRs', nms)
        lrs, pnts = [], []
        for lr, pnt in _suggestions:
            lrs.append(lr)
            pnts.append(pnt)   
        if show_plot: self.track_results.plot_xlr_find(suggestions=pnts, nms=nms)
        return SuggestedLRs(*lrs)
    
    elif show_plot: self.track_results.plot_xlr_find()
        

# %% ../../nbs/11_l2r.callbacks.ipynb 32
class ProgressBarCallback(Callback):
    order = 70
    
    def before_fit(self):
        self.mbar = master_bar(range(self.n_epochs))
        if self.learn.logger != noop:
            self.old_logger,self.learn.logger = self.logger,self._write_stats
            # metric_names = L(self.track_results.names).map(lambda o: 'train_'+o) + L(self.track_results.names).map(lambda o: 'val_'+o)
            # metric_names.append('time')
            # metric_names = L('epoch') + metric_names
            self._write_stats(self.track_results.metric_names)
        
    def before_epoch(self):
        if getattr(self, 'mbar', False): self.mbar.update(self.epoch)
        
    def before_train(self): self._launch_pbar()
    
    def before_validate(self): self._launch_pbar()
        
    def _launch_pbar(self):
        self.pbar = progress_bar(self.dl, parent=getattr(self, 'mbar', None), leave=False)
        self.pbar.update(0)
        
    def after_batch(self):
        self.pbar.update(self.iter_num+1)
    
    def after_train(self):
        self.pbar.on_iter_end()
        
    def after_validate(self):
        self.pbar.on_iter_end()
        
    def after_fit(self):
        if getattr(self, 'mbar', False):
            self.mbar.on_iter_end()
            delattr(self, 'mbar')
        if hasattr(self, 'old_logger'): self.learn.logger = self.old_logger

    def _write_stats(self, log):
        if getattr(self, 'mbar', False): self.mbar.write([f'{l:.6f}' if isinstance(l, float) else str(l) for l in log], table=True)
    

    

# %% ../../nbs/11_l2r.callbacks.ipynb 34
class SaveCallBack(Monitor):
    order = Monitor.order+1
    def __init__(self, 
        fname, 
        monitor='ndcg_at_6', 
        comp=None, 
        min_delta=0., 
        reset_on_fit=False,
    ):
        super().__init__(monitor=monitor, comp=comp, min_delta=min_delta, reset_on_fit=reset_on_fit)
        self.last_saved_path = None
        store_attr('fname')
        
    @property
    def best(self): return self._best
    @best.setter    
    def best(self, b): self._best = b
    
    def after_epoch(self):
        super().after_epoch()
        if self.new_best:
            print(f'Better model found at epoch {self.epoch} with {self.monitor} value: {self.best}.')
            self.learn.save(self.fname)
    
    # def after_fit(self):
        # if self.best_at_end: self.learn.load(self.fname)

# %% ../../nbs/11_l2r.callbacks.ipynb 36
class InterruptCallback(Callback):
    def __init__(self, epoch):
        self._interupt_before = epoch
    def before_epoch(self):
        if self.epoch == self._interupt_before:
            raise CancelFitException

# %% ../../nbs/11_l2r.callbacks.ipynb 38
class CSVLog(Callback):
    "Log the results displayed in `learn.path/fname`"
    order=60
    def __init__(self, fname='history.csv', append=False):
        self.fname,self.append = Path(fname),append

    def read_log(self):
        "Convenience method to quickly access the log."
        return pd.read_csv(self.path/self.fname)

    def before_fit(self):
        "Prepare file with metric names."
        self.path.parent.mkdir(parents=True, exist_ok=True)
        self.file = (self.path/self.fname).open('a' if self.append else 'w')
        self.file.write(','.join(self.track_results.metric_names) + '\n')
        self.old_logger,self.learn.logger = self.logger,self._write_line

    def _write_line(self, log):
        "Write a line with `log` and call the old logger."
        self.file.write(','.join([str(t) for t in log]) + '\n')
        self.file.flush()
        os.fsync(self.file.fileno())
        self.old_logger(log)

    def after_fit(self):
        "Close the file and clean up."
        self.file.close()
        self.learn.logger = self.old_logger

# %% ../../nbs/11_l2r.callbacks.ipynb 40
from torch.cuda.amp import GradScaler, autocast

# %% ../../nbs/11_l2r.callbacks.ipynb 41
class AMPMode(Enum):
    "Automatic mixed precision modes for ease of completion"
    FP16 = 'fp16'
    BF16 = 'bf16'

# %% ../../nbs/11_l2r.callbacks.ipynb 42
class TestCallback(Callback):
    order = 1000
    def before_batch(self):
        # import pdb; pdb.set_trace()
        ic(self.opt.hypers)
    # def after_epoch(self):
        # ic(self.opt.hypers)

# %% ../../nbs/11_l2r.callbacks.ipynb 43
class MemoryErrorTest(Callback):
    order=SaveCallBack.order-1
    def after_batch(self):
        if self.iter_num == 4:
            # raise torch.cuda.OutOfMemoryError("Memory !!!")
            raise CancelTrainException()
    def after_cancel_train(self): pass

# %% ../../nbs/11_l2r.callbacks.ipynb 44
@delegates(GradScaler)
class MixedPrecision(Callback):
    "Mixed precision training using Pytorch's Automatic Mixed Precision (AMP)"
    order = 10
    def __init__(self,
        amp_mode:str|AMPMode=AMPMode.FP16, # Mixed Precision training mode. Supports fp16 and bf16.
        **kwargs
    ):
        amp_mode = AMPMode(amp_mode)
        store_attr(names='amp_mode')
        self.kwargs = kwargs

    def before_fit(self):
        if self.amp_mode == AMPMode.BF16:
            if torch.cuda.is_available() and not torch.cuda.is_bf16_supported():
                raise ValueError("Unsupported GPU for bfloat16 mixed precision training")
            dtype = torch.bfloat16
        elif self.amp_mode == AMPMode.FP16:
            dtype = torch.float16
        else:
            raise ValueError(f"Unrecognized precision: {self.amp_mode}")
        # `GradScaler` is not needed for bfloat16 as fp32 and bf16 have the same range
        self.kwargs['enabled'] = dtype == torch.float16
        self.autocast,self.learn.scaler,self.scales = autocast(dtype=dtype), GradScaler(**self.kwargs),L()

    def before_batch(self): 
        if self.model.training: 
            self.autocast.__enter__()
    def after_pred(self):
        pass
        # assert self.learn.preds.dtype == torch.float16
        # self.learn.preds = to_float(self.preds)
    def after_loss(self): self.autocast.__exit__(None, None, None)
    def before_backward(self): 
        # pass
        # import pdb; pdb.set_trace()
        # self.learn.preds = self.scaler.scale(self.preds)
        assert self.learn.preds.dtype == torch.float16
        # assert self.learn.lambda_i.dtype == torch.float16
        assert self.learn.srtd_preds.dtype == torch.float16
        self.learn.srtd_preds = self.scaler.scale(self.srtd_preds)
        # self.learn.lambda_i = self.scaler.scale(self.lambda_i)
    # def after_backward(self):
    #     print(self.track_results.grads_full['token_factors.weight']
    #     import pdb; pdb.set_trace()
    def before_step(self):
        "Use `self` as a fake optimizer. `self.skipped` will be set to True `after_step` if gradients overflow."
        self.skipped=True
        self.scaler.step(self)
        if self.skipped: raise CancelStepException()
        self.scales.append(self.scaler.get_scale())
        # pass
    def after_step(self): self.learn.scaler.update()
    def after_cancel_step(self): pass
    def after_fit(self): self.autocast,self.learn.scaler,self.scales = None,None,None

    @property
    def param_groups(self):
        "Pretend to be an optimizer for `GradScaler`"
        return self.opt.param_groups
    def step(self, *args, **kwargs):
        "Fake optimizer step to detect whether this batch was skipped from `GradScaler`"
        self.skipped=False
    def after_epoch(self): 
        torch.cuda.empty_cache()
    def after_batch(self): 
        torch.cuda.empty_cache()

# %% ../../nbs/11_l2r.callbacks.ipynb 47
class GradientClip(Callback):
    "Clip norm of gradients"
    order=MixedPrecision.order+1
    def __init__(self,max_norm:float=1., norm_type:float=2.0): store_attr()
    def before_step(self): 
        # import pdb; pdb.set_trace()
        nn.utils.clip_grad_norm_(self.model.parameters(), self.max_norm, self.norm_type)
   
