text: "Callbacks provides extensibility to the `basic_train` loop. See `train` for examples of custom callbacks."
  from .basic_data import *
  from .torch_core import *
  import torch.distributed as dist
  
  __all__ = ['AverageMetric', 'Callback', 'CallbackHandler', 'OptimWrapper', 'SmoothenValue', 'Scheduler', 'annealing_cos', 'CallbackList',
             'annealing_exp', 'annealing_linear', 'annealing_no', 'annealing_poly']
  
  class OptimWrapper():
      "Basic wrapper around `opt` to simplify hyper-parameters changes."
      def __init__(self, opt:optim.Optimizer, wd:Floats=0., true_wd:bool=False, bn_wd:bool=True):
          assert not isinstance(opt, OptimWrapper)
          self.opt,self.true_wd,self.bn_wd = opt,true_wd,bn_wd
          self.opt_keys = list(self.opt.param_groups[0].keys())
          self.opt_keys.remove('params')
          self.read_defaults()
          self.wd = wd
  
      @classmethod
      def create(cls, opt_func:Union[type,Callable], lr:Union[float,Tuple,List], layer_groups:ModuleList, wd:Floats=0., 
                 true_wd:bool=False, bn_wd:bool=True)->optim.Optimizer:
          "Create an `optim.Optimizer` from `opt_func` with `lr`. Set lr on `layer_groups`."
          split_params = split_no_wd_params(layer_groups)
          opt = opt_func([{'params': p, 'lr':0} for p in split_params])
          opt = cls(opt, wd=wd, true_wd=true_wd, bn_wd=bn_wd)
          opt.lr,opt.opt_func = listify(lr, layer_groups),opt_func
          return opt
  
      def new(self, layer_groups:Collection[nn.Module], split_no_wd:bool=True):
          "Create a new `OptimWrapper` from `self` with another `layer_groups` but the same hyper-parameters."
          opt_func = getattr(self, 'opt_func', self.opt.__class__)
          res = self.create(opt_func, self.lr, layer_groups, wd=self.wd, true_wd=self.true_wd, bn_wd=self.bn_wd)
          res.mom,res.beta = self.mom,self.beta
          return res
  
      def new_with_params(self, param_groups:Collection[Collection[nn.Parameter]]):
          "Create a new `OptimWrapper` from `self` with another `layer_groups` but the same hyper-parameters."
          opt_func = getattr(self, 'opt_func', self.opt.__class__)
          opt = opt_func([{'params': p, 'lr':0} for p in param_groups])
          opt = self.__class__(opt, wd=self.wd, true_wd=self.true_wd, bn_wd=self.bn_wd)
          opt.lr,opt.opt_func,opt.mom,opt.beta = self.lr,opt_func,self.mom,self.beta
          return opt
  
      def __repr__(self)->str:
          return f'OptimWrapper over {repr(self.opt)}.\nTrue weight decay: {self.true_wd}'
  
      #Pytorch optimizer methods
      def step(self)->None:
          "Set weight decay and step optimizer."
          # weight decay outside of optimizer step (AdamW)
          if self.true_wd:
              for lr,wd,pg1,pg2 in zip(self._lr,self._wd,self.opt.param_groups[::2],self.opt.param_groups[1::2]):
                  for p in pg1['params']: p.data.mul_(1 - wd*lr)
                  if self.bn_wd:
                      for p in pg2['params']: p.data.mul_(1 - wd*lr)
              self.set_val('weight_decay', listify(0, self._wd))
          self.opt.step()
  
      def zero_grad(self)->None:
          "Clear optimizer gradients."
          self.opt.zero_grad()
  
      #Passthrough to the inner opt.
      def __getattr__(self, k:str)->Any: return getattr(self.opt, k, None)
      def __setstate__(self,data:Any): self.__dict__.update(data)
  
      def clear(self):
          "Reset the state of the inner optimizer."
          sd = self.state_dict()
          sd['state'] = {}
          self.load_state_dict(sd)
  
      @property
      def n_params(self): return sum([len(pg['params']) for pg in self.opt.param_groups])
  
      #Hyperparameters as properties
      @property
      def lr(self)->float: return self._lr[-1]
      @lr.setter
      def lr(self, val:float)->None:
          self._lr = self.set_val('lr', listify(val, self._lr))
  
      @property
      def mom(self)->float:return self._mom[-1]
      @mom.setter
      def mom(self, val:float)->None:
          if 'momentum' in self.opt_keys: self.set_val('momentum', listify(val, self._mom))
          elif 'betas' in self.opt_keys:  self.set_val('betas', (listify(val, self._mom), self._beta))
          self._mom = listify(val, self._mom)
  
      @property
      def beta(self)->float: return None if self._beta is None else self._beta[-1]
      @beta.setter
      def beta(self, val:float)->None:
          "Set beta (or alpha as makes sense for given optimizer)."
          if val is None: return
          if 'betas' in self.opt_keys:    self.set_val('betas', (self._mom, listify(val, self._beta)))
          elif 'alpha' in self.opt_keys:  self.set_val('alpha', listify(val, self._beta))
          self._beta = listify(val, self._beta)
  
      @property
      def wd(self)->float: return self._wd[-1]
      @wd.setter
      def wd(self, val:float)->None:
          "Set weight decay."
          if not self.true_wd: self.set_val('weight_decay', listify(val, self._wd), bn_groups=self.bn_wd)
          self._wd = listify(val, self._wd)
  
      #Helper functions
      def read_defaults(self)->None:
          "Read the values inside the optimizer for the hyper-parameters."
          self._beta = None
          if 'lr' in self.opt_keys: self._lr = self.read_val('lr')
          if 'momentum' in self.opt_keys: self._mom = self.read_val('momentum')
          if 'alpha' in self.opt_keys: self._beta = self.read_val('alpha')
          if 'betas' in self.opt_keys: self._mom,self._beta = self.read_val('betas')
          if 'weight_decay' in self.opt_keys: self._wd = self.read_val('weight_decay')
          reserved_names = ['params', 'lr', 'momentum', 'alpha', 'betas', 'weight_decay']
          stat_names = [n for n in self.opt_keys if n not in reserved_names]
          self._stats = {n:self.read_val(n) for n in stat_names}
  
      def get_stat(self, name:str)->float: 
          if name in ['lr', 'mom', 'beta', 'wd']: return getattr(self, name)
          else: return self._stats[name][-1]
      def set_stat(self, name:str, value:Union[float, Collection[float]])->None:
          if name in ['lr', 'mom', 'beta', 'wd']: setattr(self, name, value)
          else:
              val = listify(value, self._stats[name])
              self.set_val(name, val)
              self._stats[name] = val
  
      def set_val(self, key:str, val:Any, bn_groups:bool=True)->Any:
          "Set `val` inside the optimizer dictionary at `key`."
          if is_tuple(val): val = [(v1,v2) for v1,v2 in zip(*val)]
          for v,pg1,pg2 in zip(val,self.opt.param_groups[::2],self.opt.param_groups[1::2]):
              pg1[key] = v
              if bn_groups: pg2[key] = v
          return val
  
      def read_val(self, key:str) -> Union[List[float],Tuple[List[float],List[float]]]:
          "Read a hyperparameter `key` in the optimizer dictionary."
          val = [pg[key] for pg in self.opt.param_groups[::2]]
          if is_tuple(val[0]): val = [o[0] for o in val], [o[1] for o in val]
          return val
      
      def get_state(self):
          "Return the inner state minus the layer groups."
          return {'opt_state':self.opt.state_dict(), 'lr':self._lr, 'wd':self._wd, 'beta':self._beta, 'mom':self._mom,
                  'opt_func':self.opt_func, 'true_wd':self.true_wd, 'bn_wd':self.bn_wd}
  
      @classmethod
      def load_with_state_and_layer_group(cls, state:dict, layer_groups:Collection[nn.Module]):
          res = cls.create(state['opt_func'], state['lr'], layer_groups, wd=state['wd'], true_wd=state['true_wd'], 
                       bn_wd=state['bn_wd'])
          res._mom,res._beta = state['mom'],state['beta']
          res.load_state_dict(state['opt_state'])
          return res
  
  class Callback():
      "Base class for callbacks that want to record values, dynamically change learner params, etc."
      _order=0
      def on_train_begin(self, **kwargs:Any)->None:
          "To initialize constants in the callback."
          pass
      def on_epoch_begin(self, **kwargs:Any)->None:
          "At the beginning of each epoch."
          pass
      def on_batch_begin(self, **kwargs:Any)->None:
          "Set HP before the output and loss are computed."
          pass
      def on_loss_begin(self, **kwargs:Any)->None:
          "Called after forward pass but before loss has been computed."
          pass
      def on_backward_begin(self, **kwargs:Any)->None:
          "Called after the forward pass and the loss has been computed, but before backprop."
          pass
      def on_backward_end(self, **kwargs:Any)->None:
          "Called after backprop but before optimizer step. Useful for true weight decay in AdamW."
          pass
      def on_step_end(self, **kwargs:Any)->None:
          "Called after the step of the optimizer but before the gradients are zeroed."
          pass
      def on_batch_end(self, **kwargs:Any)->None:
          "Called at the end of the batch."
          pass
      def on_epoch_end(self, **kwargs:Any)->None:
          "Called at the end of an epoch."
          pass
      def on_train_end(self, **kwargs:Any)->None:
          "Useful for cleaning up things and saving files/models."
          pass
      def jump_to_epoch(self, epoch)->None:
          "To resume training at `epoch` directly."
          pass
  
      def get_state(self, minimal:bool=True):
          "Return the inner state of the `Callback`, `minimal` or not."
          to_remove = ['exclude', 'not_min'] + getattr(self, 'exclude', []).copy()
          if minimal: to_remove += getattr(self, 'not_min', []).copy()
          return {k:v for k,v in self.__dict__.items() if k not in to_remove}
  
      def  __repr__(self):
          attrs = func_args(self.__init__)
          to_remove = getattr(self, 'exclude', [])
          list_repr = [self.__class__.__name__] + [f'{k}: {getattr(self, k)}' for k in attrs if k != 'self' and k not in to_remove]
          return '\n'.join(list_repr)
  
  class SmoothenValue():
      "Create a smooth moving average for a value (loss, etc) using `beta`."
      def __init__(self, beta:float):
          self.beta,self.n,self.mov_avg = beta,0,0
  
      def add_value(self, val:float)->None:
          "Add `val` to calculate updated smoothed value."
          self.n += 1
          self.mov_avg = self.beta * self.mov_avg + (1 - self.beta) * val
          self.smooth = self.mov_avg / (1 - self.beta ** self.n)
  
  CallbackList = Collection[Callback]
  
  def _get_init_state(): return {'epoch':0, 'iteration':0, 'num_batch':0, 'skip_validate': False}
  
  @dataclass
  class CallbackHandler():
      "Manage all of the registered `callbacks` and `metrics`, smoothing loss by momentum `beta`."
      callbacks:CallbackList=None
      metrics:CallbackList=None
      beta:float=0.98
  
      def __post_init__(self)->None:
          "Initialize smoother and learning stats."
          self.callbacks = ifnone(self.callbacks, [])
          self.metrics = ifnone(self.metrics, [])
          self.metrics = [(met if isinstance(met, Callback) else AverageMetric(met)) for met in self.metrics]
          self.callbacks = sorted(self.callbacks, key=lambda o: getattr(o, '_order', 0))
          self.smoothener = SmoothenValue(self.beta)
          self.state_dict:Dict[str,Union[int,float,Tensor]]=_get_init_state()
  
      def _call_and_update(self, cb, cb_name, **kwargs)->None:
          "Call `cb_name` on `cb` and update the inner state."
          new = ifnone(getattr(cb, f'on_{cb_name}')(**self.state_dict, **kwargs), dict())
          for k,v in new.items():
              if k not in self.state_dict:
                  raise Exception(f"{k} isn't a valid key in the state of the callbacks.")
              else: self.state_dict[k] = v
      
      def __call__(self, cb_name, call_mets=True, **kwargs)->None:
          "Call through to all of the `CallbakHandler` functions."
          if call_mets: 
              for met in self.metrics: self._call_and_update(met, cb_name, **kwargs)
          for cb in self.callbacks: self._call_and_update(cb, cb_name, **kwargs)
  
      def set_dl(self, dl:DataLoader):
          "Set the current `dl` used."
          if hasattr(self, 'cb_dl'): self.callbacks.remove(self.cb_dl)
          if isinstance(dl.dataset, Callback):
              self.callbacks.append(dl.dataset)
              self.cb_dl = dl.dataset
  
      def on_train_begin(self, epochs:int, pbar:PBar, metrics:MetricFuncList)->None:
          "About to start learning."
          self.state_dict = _get_init_state()
          self.state_dict.update(dict(n_epochs=epochs, pbar=pbar, metrics=metrics))
          names = [(met.name if hasattr(met, 'name') else camel2snake(met.__class__.__name__)) for met in self.metrics]
          self('train_begin', metrics_names=names)
          if self.state_dict['epoch'] != 0:
              self.state_dict['pbar'].main_bar.total -= self.state_dict['epoch']
              for cb in self.callbacks: cb.jump_to_epoch(self.state_dict['epoch'])
  
      def on_epoch_begin(self)->None:
          "Handle new epoch."
          self.state_dict['num_batch'],self.state_dict['stop_training'] = 0,False
          self('epoch_begin')
  
      def on_batch_begin(self, xb:Tensor, yb:Tensor, train:bool=True)->Tuple[Any,Any]:
          "Handle new batch `xb`,`yb` in `train` or validation."
          self.state_dict.update(dict(last_input=xb, last_target=yb, train=train, 
              stop_epoch=False, skip_step=False, skip_zero=False, skip_bwd=False))
          self('batch_begin', call_mets = not self.state_dict['train'])
          return self.state_dict['last_input'], self.state_dict['last_target']
  
      def on_loss_begin(self, out:Tensor)->Any:
          "Handle start of loss calculation with model output `out`."
          self.state_dict['last_output'] = out
          self('loss_begin', call_mets=False)
          return self.state_dict['last_output']
  
      def on_backward_begin(self, loss:Tensor)->Tuple[Any,Any]:
          "Handle gradient calculation on `loss`."
          self.smoothener.add_value(loss.float().detach().cpu())
          self.state_dict['last_loss'], self.state_dict['smooth_loss'] = loss, self.smoothener.smooth
          self('backward_begin', call_mets=False)
          return self.state_dict['last_loss'], self.state_dict['skip_bwd']
  
      def on_backward_end(self)->Any:
          "Handle end of gradient calculation."
          self('backward_end', call_mets=False)
          return self.state_dict['skip_step']
  
      def on_step_end(self)->Any:
          "Handle end of optimization step."
          self('step_end', call_mets=False)
          return self.state_dict['skip_zero']
  
      def on_batch_end(self, loss:Tensor)->Any:
          "Handle end of processing one batch with `loss`."
          self.state_dict['last_loss'] = loss
          self('batch_end', call_mets = not self.state_dict['train'])
          if self.state_dict['train']:
              self.state_dict['iteration'] += 1
              self.state_dict['num_batch'] += 1
          return self.state_dict['stop_epoch']
  
      def on_epoch_end(self, val_loss:Tensor)->bool:
          "Epoch is done, process `val_loss`."
          self.state_dict['last_metrics'] = [val_loss] if val_loss is not None else [None]
          self('epoch_end', call_mets = val_loss is not None)
          self.state_dict['epoch'] += 1
          return self.state_dict['stop_training']
  
      def on_train_end(self, exception:Union[bool,Exception])->None:
          "Handle end of training, `exception` is an `Exception` or False if no exceptions during training."
          self('train_end', exception=exception)
          
      @property
      def skip_validate(self): return self.state_dict['skip_validate']
  
  class AverageMetric(Callback):
      "Wrap a `func` in a callback for metrics computation."
      def __init__(self, func):
          # If func has a __name__ use this one else it should be a partial
          name = func.__name__ if hasattr(func, '__name__') else func.func.__name__
          self.func, self.name = func, name
          self.world = num_distrib()
  
      def on_epoch_begin(self, **kwargs):
          "Set the inner value to 0."
          self.val, self.count = 0.,0
  
      def on_batch_end(self, last_output, last_target, **kwargs):
          "Update metric computation with `last_output` and `last_target`."
          if not is_listy(last_target): last_target=[last_target]
          self.count += first_el(last_target).size(0)
          val = self.func(last_output, *last_target)
          if self.world:
              val = val.clone()
              dist.all_reduce(val, op=dist.ReduceOp.SUM)
              val /= self.world
          self.val += first_el(last_target).size(0) * val.detach().cpu()
  
      def on_epoch_end(self, last_metrics, **kwargs):
          "Set the final result in `last_metrics`."
          return add_metrics(last_metrics, self.val/self.count)
  
  def annealing_no(start:Number, end:Number, pct:float)->Number:
      "No annealing, always return `start`."
      return start
  def annealing_linear(start:Number, end:Number, pct:float)->Number:
      "Linearly anneal from `start` to `end` as pct goes from 0.0 to 1.0."
      return start + pct * (end-start)
  def annealing_exp(start:Number, end:Number, pct:float)->Number:
      "Exponentially anneal from `start` to `end` as pct goes from 0.0 to 1.0."
      return start * (end/start) ** pct
  def annealing_cos(start:Number, end:Number, pct:float)->Number:
      "Cosine anneal from `start` to `end` as pct goes from 0.0 to 1.0."
      cos_out = np.cos(np.pi * pct) + 1
      return end + (start-end)/2 * cos_out
  
  def do_annealing_poly(start:Number, end:Number, pct:float, degree:Number)->Number:
      "Helper function for `anneal_poly`."
      return end + (start-end) * (1-pct)**degree
  def annealing_poly(degree:Number)->Number:
      "Anneal polynomically from `start` to `end` as pct goes from 0.0 to 1.0."
      return functools.partial(do_annealing_poly, degree=degree)
  
  class Scheduler():
      "Used to \"step\" from start,end (`vals`) over `n_iter` iterations on a schedule defined by `func`"
      def __init__(self, vals:StartOptEnd, n_iter:int, func:Optional[AnnealFunc]=None):
          self.start,self.end = (vals[0],vals[1]) if is_tuple(vals) else (vals,0)
          self.n_iter = max(1,n_iter)
          if func is None: self.func = annealing_linear if is_tuple(vals) else annealing_no
          else:          self.func = func
          self.n = 0
          
      def restart(self): self.n = 0
  
      def step(self)->Number:
          "Return next value along annealed schedule."
          self.n += 1
          return self.func(self.start, self.end, self.n/self.n_iter)
  
      @property
      def is_done(self)->bool:
          "Return `True` if schedule completed."
          return self.n >= self.n_iter
  
  