text: # 
  # controller.py
  # 
  # Author(s):
  # Matteo Spallanzani <spmatteo@iis.ee.ethz.ch>
  # 
  # Copyright (c) 2020-2021 ETH Zurich.
  # 
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  # 
  # http://www.apache.org/licenses/LICENSE-2.0
  # 
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # 
  
  from functools import partial
  from typing import NamedTuple
  from collections import OrderedDict
  import torch
  import torch.nn as nn
  import copy
  
  from .utils import almost_symm_quant
  from .ops import ANAModule
  from .ops import ANAConv1d, ANAConv2d, ANAConv3d, ANALinear, ANAActivation
  from ..controller import Controller
  
  from typing import Union, List
  
  
  
  # Coefficients generated with this code:
  # https://colab.research.google.com/drive/1IH8TwCfkvcMkHx4tHldMNgsvBIgCpTeM?usp=sharing
  # first entry is c1, second is c2 with alpha_w^* = c1 * sqrt(Ew2) - c2 * Ew1
  _sawb_symm_lut = {
        3: [ 2.58072,  1.68371],
        4: [ 3.22824,  2.19780],
        7: [ 6.99146,  6.34798],
        8: [ 7.57111,  6.96333],
       15: [12.04287, 12.07040],
       16: [12.16637, 12.19222],
       31: [14.71474, 15.13775],
       32: [15.18684, 15.72128],
       63: [20.78714, 22.60234],
       64: [20.54678, 22.29264],
      127: [27.15792, 30.50641],
      128: [27.75226, 31.25157],
      255: [35.48044, 40.89530],
      256: [34.97231, 40.26019],
  }
  
  _sawb_asymm_lut = {
        3: [ 2.58072,  1.68371],
        4: [ 4.74932,  3.32860],
        7: [ 6.99146,  6.34798],
        8: [ 8.26900,  7.36509],
       15: [12.04287, 12.07040],
       16: [12.06485, 11.86833],
       31: [14.71474, 15.13775],
       32: [15.47646, 15.98616],
       63: [20.78714, 22.60234],
       64: [21.96366, 24.03408],
      127: [27.15792, 30.50641],
      128: [28.26410, 31.87972],
      255: [35.48044, 40.89530],
      256: [36.15292, 41.73554],
  }
  
  def bws(t: int, tstart: int, tend: int, alpha: int) -> float:
      """Bounded-window scheduling function.
  
      This function returns 1.0 for `t <= tstart`, 0.0 for `tend <= t`, and
      `\lambda^{\alpha}` for `tstart < t < tend`, where `\lambda =
      \frac{tend - t}{tend - tstart}`.
      """
  
      # assert 0 <= tstart < tend
      # assert 0 <= alpha
  
      t          = min(max(t, tstart), tend)
      multiplier = (float(tend - t) / float(tend - tstart)) ** alpha
  
      return multiplier
  
  
  def uws(t: int, tstart: int, eps: float, alpha: int) -> float:
      """Unbounded-window scheduling function."""
  
      # assert 0 <= tstart
      # assert 0.0 <= eps
      # assert 0 <= alpha
  
      t          = max(t, tstart)
      multiplier = (1 / float(t - tstart + 1)) ** alpha
      multiplier = 0.0 if multiplier < eps else multiplier
  
      return multiplier
  
  
  _AC_MAPPER = {'bws': bws, 'uws': uws}
  
  
  class ANATimer(object):
  
      def __init__(self,
                   mi_spec:    dict,
                   sigma_spec: dict):
  
          self._mi = None
          
          assert ANATimer.check_spec(mi_spec, is_sigma=False)
          self._mi_beta = mi_spec['beta']
          self._mi_fun  = partial(_AC_MAPPER[mi_spec['fun']], **mi_spec['kwargs'])
  
          self._sigma = None
          assert ANATimer.check_spec(sigma_spec, is_sigma=True)
          self._sigma_beta = sigma_spec['beta']
          self._sigma_fun  = partial(_AC_MAPPER[sigma_spec['fun']], **sigma_spec['kwargs'])
  
      @property
      def mi(self) -> Union[None, float]:
          return self._mi
  
      @property
      def sigma(self) -> Union[None, float]:
          return self._sigma
  
      @staticmethod
      def check_spec(spec: dict, is_sigma: bool) -> bool:
  
          check_beta = isinstance(spec['beta'], float) and (True if not is_sigma else (0.0 <= spec['beta']))
          spec_type  = spec['fun']
          check_fun  = spec_type in set(_AC_MAPPER.keys())
  
          if check_beta and check_fun:
              spec_kwargs = spec['kwargs']
              if spec_type == 'bws':
                  check_tstart = isinstance(spec_kwargs['tstart'], int) and (0 <= spec_kwargs['tstart'])
                  check_tend   = isinstance(spec_kwargs['tend'], int)   and (spec_kwargs['tstart'] < spec_kwargs['tend'])
                  check_alpha  = isinstance(spec_kwargs['alpha'], int)  and (0 <= spec_kwargs['alpha'])
                  is_correct = check_tstart and check_tend and check_alpha
              elif spec_type == 'uws':
                  check_tstart = isinstance(spec_kwargs['tstart'], int) and (0 <= spec_kwargs['tstart'])
                  check_eps    = isinstance(spec_kwargs['eps'], float)  and (0.0 <= spec_kwargs['eps'])
                  check_alpha  = isinstance(spec_kwargs['alpha'], int)  and (0 <= spec_kwargs['alpha'])
                  is_correct   = check_tstart and check_eps and check_alpha
              else:
                  raise NotImplemented
          else:
              is_correct = False
  
          return is_correct
  
      def step(self, t: int) -> None:
          self._mi    = self._mi_beta    * self._mi_fun(t)
          self._sigma = self._sigma_beta * self._sigma_fun(t)
  
  
  Timer2Modules = NamedTuple('Timer2Modules', [('timer',   ANATimer),
                                               ('modules', List[torch.nn.Module])])
  
  
  class ANAController(Controller):
      """The training controller for the *additive noise annealing* algorithm.
  
      This object sets and updates the hyper-parameters of the additive noise
      annealing (ANA) algorithm. These hyper-parameters are the means and the
      standard deviations of uni-variate probability distributions describing
      random variables that are added to the arguments of the quantizers. The
      ANA algorithm achieves quantization by gradually *annealing* these
      distributions to Dirac deltas centered at zero; the annealing has to be
      intended as convergence in distribution.
  
      For simplicity, ANA ``torch.nn.Module``s add the same noise to all the
      components of their input ``torch.Tensor``s. This reduces the number of
      hyper-parameters required by each ANA module to two (the mean and the
      standard deviation of the distribution).
      """
  
      def __init__(self, module: torch.nn.Module, ctrl_spec: List):
  
          self._global_step   = -1
          self._timer2modules = []
  
          self._n2m = self.get_modules(module)
  
          if isinstance(module, nn.DataParallel):
              # the network is wrapped inside an nn.DataParallel module:
              # this requires to resolve an additional naming layer
              ctrl_spec = copy.deepcopy(ctrl_spec)  # TODO: find a better workaround to avoid overwriting the original controller specifications
              for timer_spec in ctrl_spec:
                  timer_spec['modules'] = ['.'.join(['module', m]) for m in timer_spec['modules']]
  
          # verify that each ANA module is linked to exactly one timer (i.e., check that ANA targets form a partition of the collection of ANA modules)
          all_ana_module_names = set(self._n2m.keys())
          ana_timer_targets    = []
          for timer_spec in ctrl_spec:
              ana_timer_target = set(timer_spec['modules'])
              assert ana_timer_target.issubset(all_ana_module_names)  # non-ANA module specified as ANA target
              ana_timer_targets.append(ana_timer_target)
          assert len(all_ana_module_names) == sum([len(ana_timer_target) for ana_timer_target in ana_timer_targets])
          assert len(all_ana_module_names) == len(set.union(*ana_timer_targets))
  
          # build timers and link them to the specified ANA modules
          for timer_spec in ctrl_spec:
              timer   = ANATimer(timer_spec['mi'], timer_spec['sigma'])
              modules = list(map(lambda n: self._n2m[n], set(timer_spec['modules'])))
              self._timer2modules.append(Timer2Modules(timer=timer, modules=modules))
  
      @staticmethod
      def get_modules(module: torch.nn.Module, parent_name: str = '', n2m: OrderedDict = OrderedDict()):
  
          for n, m in module.named_children():
              if len(list(m.children())) == 0:  # ``Module`` is not ``nn.Sequential`` or other container type; i.e., this is a "leaf" ``Module``
                  if isinstance(m, ANAModule):
                      n2m.update({parent_name + n: m})
              else:
                  ANAController.get_modules(m, parent_name=''.join([parent_name, n, '.']), n2m=n2m)
  
          return n2m
  
      @property
      def n2m(self) -> dict:
          return self._n2m
  
      def state_dict(self) -> dict:
          return {'_global_step': self._global_step}
  
      def step_pre_training_epoch(self, *args, **kwargs) -> None:
  
          self._global_step += 1 #count on the training epocs
  
          for timer, modules in self._timer2modules:
              timer.step(self._global_step) # reduces the mean and the variance during training 
              for m in modules:
                  # remove noise
  
                  if not m.started and type(m)!= ANAActivation:
                      self.reset_clip_bounds(m, init=True)
                      m.started |= True                                                         # values
                  m.set_noise(timer.mi, timer.sigma) # Set the nodes agains in the module 
  
      def step_pre_validation_epoch(self, *args, **kwargs) -> None:
  
          for timer, modules in self._timer2modules:
  
              for m in modules:
                  m.set_noise(0.0, 0.0)
  
      def reset_clip_bounds(self, m: Union[ANAConv1d, ANAConv2d, ANAConv3d, ANALinear], method: str = None, init : bool = False):
          w = m.weight.data                   # save the weights in a variable w
          if method is None:                  # Check if there is a method in configuration file
              method = m.init_clip            # if not assign the default one of the layer
          # calcoliamo le statistiche per layer      
          reduce_dims = tuple(range(len(w.shape)))           
          # Usiamo il metodo max per iniziare 
          if method=='max':
              if m.signed:
                  # Take the maximum of the first slice
                  max_val = torch.amax(w.abs(), dim=reduce_dims)                      # Returns the maximum value of each slice of the input tensor in the given dimension(s) dim.
                  # if symm_wts is true, we do "almost symmetric" quantization in the case of an even n_levels (to account for e.g. int8 range going from -128 to 127)
                  min_val, max_val = almost_symm_quant(max_val, m.n_levels)           # Compute the bounds as you can see in the function almost_symm_quant
              else:                                                                   # If you don't want symmetric clipping values, just take the maximum and the minimum absolute 
                  max_val = torch.amax(w, dim=reduce_dims)                        
                  min_val = torch.amin(w, dim=reduce_dims)
          
          elif method == 'std':
                  max_val = w.mean(dim=reduce_dims) + w.std(dim=reduce_dims) * m.nb_std
                  if m.signed:
                      # 'std' initialization is inherently symmetrical, so use the "almost_symmetrical" quantization anyway
                      min_val, max_val = almost_symm_quant(max_val, m.n_levels)
                  else:
                      min_val = w.mean(dim=reduce_dims) - w.std(dim=reduce_dims) * m.nb_std
  
          elif method[0:4] == 'sawb':
              symm = method[5:] == 'symm'
              # mean absolute weights: E[|W|] - either channel-wise or over the whole tensor depending on reduce_dims
              e_w_abs = m.weight.data.abs().mean(dim=reduce_dims)
              e_w2 = torch.mean(m.weight.data**2, dim=reduce_dims)
              if symm:
                  coeffs = _sawb_symm_lut[m.n_levels]                 # coefficients of the regression model explained in the article  Bridging the Accuracy Gap for 2-bit Quantized Neural Networks (QNN) 
                                                                      # used to minimize argmax_{alpha} ||w_q - w||
              else:
                  coeffs = _sawb_asymm_lut[m.n_levels]
              alpha = coeffs[0] * e_w2.sqrt() - coeffs[1] * e_w_abs
  
  
              # calculate the min/max bounds as well for the cases where SAWB
              # produces a negative alpha
              max_val_mm = torch.amax(w.abs(), dim=reduce_dims)
              # if symm_wts is true, we do "almost symmetric" quantization in the case of an even n_levels (to account for e.g. int8 range going from -128 to 127)
              min_val_mm, max_val_mm = almost_symm_quant(max_val_mm, m.n_levels)
  
              if symm:
                  min_val = -alpha
                  max_val = alpha
              else:
                  min_val, max_val = almost_symm_quant(alpha, m.n_levels)
  
              # where alpha is negative, use min/max bounds
              min_val, max_val = torch.where(alpha < 0, min_val_mm, min_val), torch.where(alpha < 0, max_val_mm, max_val)
          else: # method == 'const' In this method we just take the input values as 
              min_val = torch.ones_like(m.clip_lo.data) * self.init_clip_lo
              max_val = torch.ones_like(m.clip_hi.data) * self.init_clip_hi
  
  
          # Reshape the found bounds into a good format
          m.clip_hi.data = m.expand_bounds(max_val)
          m.clip_lo.data = m.expand_bounds(min_val)
      