"""Contains the classes that deal with constant temperature dynamics.

Copyright (C) 2013, Joshua More and Michele Ceriotti

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program. If not, see <http.//www.gnu.org/licenses/>.

Contains the algorithms which propagate the thermostatting steps in the constant
temperature ensembles. Includes the new GLE thermostat, which can be used to
run PI+GLE dynamics, reducing the number of path integral beads required.

Classes:
   Thermostat: Base thermostat class with the generic methods and attributes.
   ThermoLangevin: Holds the algorithms for a langevin thermostat.
   ThermoPILE_L: Holds the algorithms for a path-integral langevin equation
      thermostat, with a thermostat coupled directly to the
      centroid coordinate of each bead.
   ThermoPILE_G: Holds the algorithms for a path-integral langevin equation
      thermostat, with a thermostat coupled to the kinetic energy for
      the entire system.
   ThermoSVR: Holds the algorithms for a stochastic velocity rescaling
      thermostat.
   ThermoGLE: Holds the algorithms for a generalised langevin equation
      thermostat.
   ThermoNMGLE: Holds the algorithms for a generalised langevin equation
      thermostat in the normal mode representation.
   ThermoNMGLEG: Holds the algorithms for a generalised langevin equation
      thermostat in the normal mode representation, with kinetic energy as
      well as potential energy sampling optimization.
"""

__all__ = ['Thermostat', 'ThermoLangevin', 'ThermoPILE_L', 'ThermoPILE_G',
           'ThermoSVR', 'ThermoGLE', 'ThermoNMGLE', 'ThermoNMGLEG']

import numpy as np
from ipi.utils.depend   import *
from ipi.utils.units    import *
from ipi.utils.mathtools import matrix_exp, stab_cholesky, root_herm
from ipi.utils.prng import Random
from ipi.utils.messages import verbosity, warning, info
from ipi.engine.beads import Beads
from ipi.engine.normalmodes import NormalModes

class Thermostat(dobject):
   """Base thermostat class.

   Gives the standard methods and attributes needed in all the thermostat
   classes.

   Attributes:
      prng: A pseudo random number generator object.
      ndof: The number of degrees of freedom that the thermostat will be
         attached to.

   Depend objects:
      dt: The time step used in the algorithms. Depends on the simulation dt.
      temp: The simulation temperature. Higher than the system temperature by
         a factor of the number of beads. Depends on the simulation temp.
      ethermo: The total energy exchanged with the bath due to the thermostat.
      p: The momentum vector that the thermostat is coupled to. Depends on the
         beads p object.
      m: The mass vector associated with p. Depends on the beads m object.
      sm: The square root of the mass vector.
   """

   def __init__(self, temp = 1.0, dt = 1.0, ethermo=0.0):
      """Initialises Thermostat.

      Args:
         temp: The simulation temperature. Defaults to 1.0.
         dt: The simulation time step. Defaults to 1.0.
         ethermo: The initial heat energy transferred to the bath.
            Defaults to 0.0. Will be non-zero if the thermostat is
            initialised from a checkpoint file.
      """

      dset(self,"temp",   depend_value(name='temp', value=temp))
      dset(self,"dt",     depend_value(name='dt', value=dt))
      dset(self,"ethermo",depend_value(name='ethermo',value=ethermo))

   def bind(self, beads=None, atoms=None, pm=None, prng=None, fixdof=None):
      """Binds the appropriate degrees of freedom to the thermostat.

      This takes an object with degrees of freedom, and makes their momentum
      and mass vectors members of the thermostat. It also then creates the
      objects that will hold the data needed in the thermostat algorithms
      and the dependency network.

      Args:
         beads: An optional beads object to take the mass and momentum vectors
            from.
         atoms: An optional atoms object to take the mass and momentum vectors
            from.
         pm: An optional tuple containing a single momentum value and its
            conjugate mass.
         prng: An optional pseudo random number generator object. Defaults to
            Random().
         fixdof: An optional integer which can specify the number of constraints
            applied to the system. Defaults to zero.

      Raises:
         TypeError: Raised if no appropriate degree of freedom or object
            containing a momentum vector is specified for
            the thermostat to couple to.
      """

      if prng is None:
         warning("Initializing thermostat from standard random PRNG", verbosity.medium)
         self.prng = Random()
      else:
         self.prng = prng

      if not beads is None:
         dset(self,"p",beads.p.flatten())
         dset(self,"m",beads.m3.flatten())
      elif not atoms is None:
         dset(self,"p",dget(atoms, "p"))
         dset(self,"m",dget(atoms, "m3"))
      elif not pm is None:
         dset(self,"p",pm[0])
         dset(self,"m",pm[1])
      else:
         raise TypeError("Thermostat.bind expects either Beads, Atoms, NormalModes, or a (p,m) tuple to bind to")

      if fixdof is None:
         self.ndof = len(self.p)
      else:
         self.ndof = float(len(self.p) - fixdof)

      dset(self, "sm",
         depend_array(name="sm", value=np.zeros(len(dget(self,"m"))),
            func=self.get_sm, dependencies=[dget(self,"m")]))

   def get_sm(self):
      """Retrieves the square root of the mass matrix.

      Returns:
         A vector of the square root of the mass matrix with one value for
         each degree of freedom.
      """

      return np.sqrt(self.m)

   def step(self):
      """Dummy thermostat step."""

      pass

class ThermoLangevin(Thermostat):
   """Represents a langevin thermostat.

   Depend objects:
      tau: Thermostat damping time scale. Larger values give a less strongly
         coupled thermostat.
      T: Coefficient of the diffusive contribution of the thermostat, i.e. the
         drift back towards equilibrium. Depends on tau and the time step.
      S: Coefficient of the stochastic contribution of the thermostat, i.e.
         the uncorrelated Gaussian noise. Depends on T and the temperature.
   """

   def get_T(self):
      """Calculates the coefficient of the overall drift of the velocities."""

      return np.exp(-0.5*self.dt/self.tau)

   def get_S(self):
      """Calculates the coefficient of the white noise."""

      return np.sqrt(Constants.kb*self.temp*(1 - self.T**2))

   def __init__(self, temp = 1.0, dt = 1.0, tau = 1.0, ethermo=0.0):
      """Initialises ThermoLangevin.

      Args:
         temp: The simulation temperature. Defaults to 1.0.
         dt: The simulation time step. Defaults to 1.0.
         tau: The thermostat damping timescale. Defaults to 1.0.
         ethermo: The initial heat energy transferred to the bath.
            Defaults to 0.0. Will be non-zero if the thermostat is
            initialised from a checkpoint file.
      """

      super(ThermoLangevin,self).__init__(temp, dt, ethermo)

      dset(self,"tau",depend_value(value=tau,name='tau'))
      dset(self,"T",
         depend_value(name="T",func=self.get_T,
            dependencies=[dget(self,"tau"), dget(self,"dt")]))
      dset(self,"S",
         depend_value(name="S",func=self.get_S,
            dependencies=[dget(self,"temp"), dget(self,"T")]))

   def step(self):
      """Updates the bound momentum vector with a langevin thermostat."""

      p = depstrip(self.p).copy()
      sm = depstrip(self.sm)

      p /= sm

      self.ethermo += np.dot(p,p)*0.5
      p *= self.T
      p += self.S*self.prng.gvec(len(p))
      self.ethermo -= np.dot(p,p)*0.5

      p *= sm

      self.p = p

class ThermoPILE_L(Thermostat):
   """Represents a PILE thermostat with a local centroid thermostat.

   Attributes:
      _thermos: The list of the different thermostats for all the ring polymer
         normal modes.
      nm: A normal modes object to attach the thermostat to.
      prng: Random number generator used in the stochastic integration
         algorithms.

   Depend objects:
      tau: Centroid thermostat damping time scale. Larger values give a
         less strongly coupled centroid thermostat.
      tauk: Thermostat damping time scale for the non-centroid normal modes.
         Depends on the ring polymer spring constant, and thus the simulation
         temperature.
      pilescale: A float used to reduce the intensity of the PILE thermostat if
         required.
   """

   def __init__(self, temp = 1.0, dt = 1.0, tau = 1.0, ethermo=0.0, scale=1.0):
      """Initialises ThermoPILE_L.

      Args:
         temp: The simulation temperature. Defaults to 1.0.
         dt: The simulation time step. Defaults to 1.0.
         tau: The centroid thermostat damping timescale. Defaults to 1.0.
         ethermo: The initial conserved energy quantity. Defaults to 0.0. Will
            be non-zero if the thermostat is initialised from a checkpoint file.
         scale: A float used to reduce the intensity of the PILE thermostat if
            required.

      Raises:
         TypeError: Raised if the thermostat is used with any object other than
            a beads object, so that we make sure that the objects needed for the
            normal mode transformation exist.
      """

      super(ThermoPILE_L,self).__init__(temp,dt,ethermo)
      dset(self,"tau",depend_value(value=tau,name='tau'))
      dset(self,"pilescale",depend_value(value=scale,name='pilescale'))

   def bind(self, nm=None, prng=None, bindcentroid=True, fixdof=None):
      """Binds the appropriate degrees of freedom to the thermostat.

      This takes a beads object with degrees of freedom, and makes its momentum
      and mass vectors members of the thermostat. It also then creates the
      objects that will hold the data needed in the thermostat algorithms
      and the dependency network.

      Gives the interface for both the PILE_L and PILE_G thermostats, which
      only differ in their treatment of the centroid coordinate momenta.

      Args:
         nm: An optional normal mode object to take the mass and momentum
            vectors from.
         prng: An optional pseudo random number generator object. Defaults to
            Random().
         bindcentroid: An optional boolean which decides whether a Langevin
            thermostat is attached to the centroid mode of each atom
            separately, or the total kinetic energy. Defaults to True, which
            gives a thermostat bound to each centroid momentum.
         fixdof: An optional integer which can specify the number of constraints
            applied to the system. Defaults to zero.

      Raises:
         TypeError: Raised if no appropriate degree of freedom or object
            containing a momentum vector is specified for
            the thermostat to couple to.
      """

      if nm is None or not type(nm) is NormalModes:
         raise TypeError("ThermoPILE_L.bind expects a NormalModes argument to bind to")
      if prng is None:
         self.prng = Random()
      else:
         self.prng = prng

      prev_ethermo = self.ethermo

      # creates a set of thermostats to be applied to individual normal modes
      self._thermos = [ ThermoLangevin(temp=1, dt=1, tau=1) for b in range(nm.nbeads) ]
      # optionally does not bind the centroid, so we can re-use all of this
      # in the PILE_G case
      if not bindcentroid:
         self._thermos[0] = None

      self.nm = nm

      dset(self,"tauk",
         depend_array(name="tauk", value=np.zeros(nm.nbeads-1,float),
            func=self.get_tauk, dependencies=[dget(self,"pilescale"), dget(nm,"dynomegak")] ) )

      # must pipe all the dependencies in such a way that values for the nm thermostats
      # are automatically updated based on the "master" thermostat
      def make_taugetter(k):
         return lambda: self.tauk[k-1]
      it = 0
      for t in self._thermos:
         if t is None:
            it += 1
            continue
         if it > 0:
            fixdof = None # only the centroid thermostat may have constraints

         # bind thermostat t to the it-th bead
         t.bind(pm=(nm.pnm[it,:],nm.dynm3[it,:]),prng=self.prng, fixdof=fixdof)
         # pipes temp and dt
         deppipe(self,"temp", t, "temp")
         deppipe(self,"dt", t, "dt")

         # for tau it is slightly more complex
         if it == 0:
            deppipe(self,"tau", t, "tau")
         else:
            # Here we manually connect _thermos[i].tau to tauk[i].
            # Simple and clear.
            dget(t,"tau").add_dependency(dget(self,"tauk"))
            dget(t,"tau")._func = make_taugetter(it)
         dget(self,"ethermo").add_dependency(dget(t,"ethermo"))
         it += 1

      # since the ethermo will be "delegated" to the normal modes thermostats,
      # one has to split
      # any previously-stored value between the sub-thermostats
      if bindcentroid:
         for t in self._thermos:
            t.ethermo = prev_ethermo/nm.nbeads
         dget(self,"ethermo")._func = self.get_ethermo;
         # if we are not binding the centroid just yet, this bit of the piping
         # is delegated to the function which is actually calling this

   def get_tauk(self):
      """Computes the thermostat damping time scale for the non-centroid
      normal modes.

      Returns:
         An array with the damping time scales for the non-centroid modes.
      """

      # Also include an optional scaling factor to reduce the intensity of NM thermostats
      return  np.array([ self.pilescale/(2*self.nm.dynomegak[k])  for k in range(1,len(self._thermos)) ])

   def get_ethermo(self):
      """Computes the total energy transferred to the heat bath for all the
      thermostats.
      """

      et = 0.0;
      for t in self._thermos:
         et += t.ethermo
      return et

   def step(self):
      """Updates the bound momentum vector with a PILE thermostat."""

      # super-cool! just loop over the thermostats! it's as easy as that!
      for t in self._thermos:
         t.step()

class ThermoSVR(Thermostat):
   """Represents a stochastic velocity rescaling thermostat.

   Depend objects:
      tau: Centroid thermostat damping time scale. Larger values give a
         less strongly coupled centroid thermostat.
      K: Scaling factor for the total kinetic energy. Depends on the
         temperature.
      et: Parameter determining the strength of the thermostat coupling.
         Depends on tau and the time step.
   """

   def get_et(self):
      """Calculates the damping term in the propagator."""

      return np.exp(-0.5*self.dt/self.tau)

   def get_K(self):
      """Calculates the average kinetic energy per degree of freedom."""

      return Constants.kb*self.temp*0.5

   def __init__(self, temp = 1.0, dt = 1.0, tau = 1.0, ethermo=0.0):
      """Initialises ThermoSVR.

      Args:
         temp: The simulation temperature. Defaults to 1.0.
         dt: The simulation time step. Defaults to 1.0.
         tau: The thermostat damping timescale. Defaults to 1.0.
         ethermo: The initial conserved energy quantity. Defaults to 0.0. Will
            be non-zero if the thermostat is initialised from a checkpoint file.
      """

      super(ThermoSVR,self).__init__(temp,dt,ethermo)

      dset(self,"tau",depend_value(value=tau,name='tau'))
      dset(self,"et",
         depend_value(name="et",func=self.get_et,
            dependencies=[dget(self,"tau"), dget(self,"dt")]))
      dset(self,"K",
         depend_value(name="K",func=self.get_K, dependencies=[dget(self,"temp")]))

   def step(self):
      """Updates the bound momentum vector with a stochastic velocity rescaling
      thermostat. See G Bussi, D Donadio, M Parrinello,
      Journal of Chemical Physics 126, 014101 (2007)
      """

      K = np.dot(depstrip(self.p),depstrip(self.p)/depstrip(self.m))*0.5

      # rescaling is un-defined if the KE is zero
      if K == 0.0:
         return

      # gets the stochastic term (basically a Gamma distribution for the kinetic energy)
      r1 = self.prng.g
      if (self.ndof-1)%2 == 0:
         rg = 2.0*self.prng.gamma((self.ndof-1)/2)
      else:
         rg = 2.0*self.prng.gamma((self.ndof-2)/2) + self.prng.g**2

      alpha2 = self.et + self.K/K*(1 - self.et)*(r1**2 + rg) + 2.0*r1*np.sqrt(self.K/K*self.et*(1 - self.et))
      alpha = np.sqrt(alpha2)
      if (r1 + np.sqrt(2*K/self.K*self.et/(1 - self.et))) < 0:
         alpha *= -1

      self.ethermo += K*(1 - alpha2)
      self.p *= alpha

class ThermoPILE_G(ThermoPILE_L):
   """Represents a PILE thermostat with a global centroid thermostat.

   Simply replaces the Langevin thermostat for the centroid normal mode with
   a global velocity rescaling thermostat.
   """

   def __init__(self, temp = 1.0, dt = 1.0, tau = 1.0, ethermo=0.0, scale = 1.0):
      """Initialises ThermoPILE_G.

      Args:
         temp: The simulation temperature. Defaults to 1.0.
         dt: The simulation time step. Defaults to 1.0.
         tau: The centroid thermostat damping timescale. Defaults to 1.0.
         ethermo: The initial conserved energy quantity. Defaults to 0.0. Will
            be non-zero if the thermostat is initialised from a checkpoint file.
         scale: A float used to reduce the intensity of the PILE thermostat if
            required.
      """

      super(ThermoPILE_G,self).__init__(temp,dt,tau,ethermo)
      dset(self,"pilescale",depend_value(value=scale,name='pilescale'))

   def bind(self, nm=None, prng=None, fixdof=None):
      """Binds the appropriate degrees of freedom to the thermostat.

      This takes a beads object with degrees of freedom, and makes its momentum
      and mass vectors members of the thermostat. It also then creates the
      objects that will hold the data needed in the thermostat algorithms
      and the dependency network.

      Uses the PILE_L bind interface, with bindcentroid set to false so we can
      specify that thermostat separately, by binding a global
      thermostat to the centroid mode.

      Args:
         beads: An optional beads object to take the mass and momentum vectors
            from.
         prng: An optional pseudo random number generator object. Defaults to
            Random().
         fixdof: An optional integer which can specify the number of constraints
            applied to the system. Defaults to zero.

      """

      # first binds as a local PILE, then substitutes the thermostat on the centroid
      prev_ethermo = self.ethermo
      super(ThermoPILE_G,self).bind(nm=nm,prng=prng,bindcentroid=False, fixdof=fixdof)

      #centroid thermostat
      self._thermos[0] = ThermoSVR(temp=1, dt=1, tau=1)

      t = self._thermos[0]
      t.bind(pm=(nm.pnm[0,:],nm.dynm3[0,:]),prng=self.prng, fixdof=fixdof)
      deppipe(self,"temp", t, "temp")
      deppipe(self,"dt", t, "dt")
      deppipe(self,"tau", t, "tau")
      dget(self,"ethermo").add_dependency(dget(t,"ethermo"))

      # splits any previous ethermo between the thermostats, and finishes to bind ethermo to the sum function
      for t in self._thermos:
         t.ethermo = prev_ethermo/nm.nbeads
      dget(self,"ethermo")._func = self.get_ethermo;

class ThermoGLE(Thermostat):
   """Represents a GLE thermostat.

   This is similar to a langevin thermostat, in that it uses Gaussian random
   numbers to simulate a heat bath acting on the system, but simulates a
   non-Markovian system by using a Markovian formulation in an extended phase
   space. This allows for a much greater degree of flexibility, and this
   thermostat, properly fitted, can give the an approximation to the correct
   quantum ensemble even for a classical, 1-bead simulation. More reasonably,
   using this thermostat allows for a far smaller number of replicas of the
   system to be used, as the convergence of the properties
   of the system is accelerated with respect to number of beads when PI+GLE
   are used in combination. (See M. Ceriotti, D. E. Manolopoulos, M. Parinello,
   J. Chem. Phys. 134, 084104 (2011)).

   Attributes:
      ns: The number of auxilliary degrees of freedom.
      s: An array holding all the momenta, including the ones for the
         auxilliary degrees of freedom.

   Depend objects:
      A: Drift matrix giving the damping time scales for all the different
         degrees of freedom.
      C: Static covariance matrix.
         Satisfies A.C + C.transpose(A) = B.transpose(B), where B is the
         diffusion matrix, giving the strength of the coupling of the system
         with the heat bath, and thus the size of the stochastic
         contribution of the thermostat.
      T: Matrix for the diffusive contribution of the thermostat, i.e. the
         drift back towards equilibrium. Depends on A and the time step.
      S: Matrix for the stochastic contribution of the thermostat, i.e.
         the uncorrelated Gaussian noise. Depends on C and T.
   """

   def get_T(self):
      """Calculates the matrix for the overall drift of the velocities."""

      return matrix_exp(-0.5*self.dt*self.A)

   def get_S(self):
      """Calculates the matrix for the coloured noise."""

      SST = Constants.kb*(self.C - np.dot(self.T,np.dot(self.C,self.T.T)))

      # Uses a symetric decomposition rather than Cholesky, since it is more stable
      return root_herm(SST)

   def get_C(self):
      """Calculates C from temp (if C is not set explicitly)"""

      rC = np.identity(self.ns + 1,float)*self.temp
      return rC[:]

   def __init__(self, temp = 1.0, dt = 1.0, A = None, C = None, ethermo=0.0):
      """Initialises ThermoGLE.

      Args:
         temp: The simulation temperature. Defaults to 1.0.
         dt: The simulation time step. Defaults to 1.0.
         A: An optional matrix giving the drift matrix. Defaults to a single
            value of 1.0.
         C: An optional matrix giving the covariance matrix. Defaults to an
            identity matrix times temperature with the same dimensions as the
            total number of degrees of freedom in the system.
         ethermo: The initial heat energy transferred to the bath.
            Defaults to 0.0. Will be non-zero if the thermostat is
            initialised from a checkpoint file.
      """

      super(ThermoGLE,self).__init__(temp,dt,ethermo)

      if A is None:
         A = np.identity(1,float)
      dset(self,"A",depend_value(value=A.copy(),name='A'))

      self.ns = len(self.A) - 1;

      # now, this is tricky. if C is taken from temp, then we want it to be updated
      # as a depend of temp. Otherwise, we want it to be an independent beast.
      if C is None:
         C = np.identity(self.ns+1,float)*self.temp
         dset(self,"C",
            depend_value(name='C', func=self.get_C,
               dependencies=[dget(self,"temp")]))
      else:
         dset(self,"C",depend_value(value=C.copy(),name='C'))

      dset(self,"T",
         depend_value(name="T",func=self.get_T,
            dependencies=[dget(self,"A"), dget(self,"dt")]))
      dset(self,"S",
         depend_value(name="S",func=self.get_S,
            dependencies=[dget(self,"C"), dget(self,"T")]))

      self.s = np.zeros(0)

   def bind(self, beads=None, atoms=None, pm=None, prng=None, fixdof=None):
      """Binds the appropriate degrees of freedom to the thermostat.

      This takes an object with degrees of freedom, and makes their momentum
      and mass vectors members of the thermostat. It also then creates the
      objects that will hold the data needed in the thermostat algorithms
      and the dependency network.

      Args:
         beads: An optional beads object to take the mass and momentum vectors
            from.
         atoms: An optional atoms object to take the mass and momentum vectors
            from.
         pm: An optional tuple containing a single momentum value and its
            conjugate mass.
         prng: An optional pseudo random number generator object. Defaults to
            Random().
         fixdof: An optional integer which can specify the number of constraints
            applied to the system. Defaults to zero.

      Raises:
         TypeError: Raised if no appropriate degree of freedom or object
            containing a momentum vector is specified for
            the thermostat to couple to.
      """

      super(ThermoGLE,self).bind(beads,atoms,pm,prng,fixdof)

      # allocates, initializes or restarts an array of s's
      if self.s.shape != (self.ns + 1, len(dget(self,"m"))):
         if len(self.s) > 0:
            warning("Mismatch in GLE s array size on restart, will reinitialise to free particle.", verbosity.low)
         self.s = np.zeros((self.ns + 1, len(dget(self,"m"))))

         # Initializes the s vector in the free-particle limit
         info(" GLE additional DOFs initialised to the free-particle limit.", verbosity.low)
         SC = stab_cholesky(self.C*Constants.kb)
         self.s[:] = np.dot(SC, self.prng.gvec(self.s.shape))
      else:
         info("GLE additional DOFs initialised from input.", verbosity.medium)

   def step(self):
      """Updates the bound momentum vector with a GLE thermostat"""

      p = depstrip(self.p).copy()

      self.s[0,:] = self.p/self.sm

      self.ethermo += np.dot(self.s[0],self.s[0])*0.5
      self.s[:] = np.dot(self.T,self.s) + np.dot(self.S,self.prng.gvec(self.s.shape))
      self.ethermo -= np.dot(self.s[0],self.s[0])*0.5

      self.p = self.s[0]*self.sm

class ThermoNMGLE(Thermostat):
   """Represents a 'normal-modes' GLE thermostat.

   An extension to the GLE thermostat which is applied in the
   normal modes representation, and which allows to use a different
   GLE for each normal mode

   Attributes:
      ns: The number of auxilliary degrees of freedom.
      nb: The number of beads.
      s: An array holding all the momenta, including the ones for the
         auxilliary degrees of freedom.

   Depend objects:
      A: Drift matrix giving the damping time scales for all the different
         degrees of freedom (must contain nb terms).
      C: Static covariance matrix.
         Satisfies A.C + C.transpose(A) = B.transpose(B), where B is the
         diffusion matrix, giving the strength of the coupling of the system
         with the heat bath, and thus the size of the stochastic
         contribution of the thermostat.
   """

   def get_C(self):
      """Calculates C from temp (if C is not set explicitly)."""

      rv = np.ndarray((self.nb, self.ns+1, self.ns+1), float)
      for b in range(0,self.nb):
         rv[b] = np.identity(self.ns + 1,float)*self.temp
      return rv[:]

   def __init__(self, temp = 1.0, dt = 1.0, A = None, C = None, ethermo=0.0):
      """Initialises ThermoGLE.

      Args:
         temp: The simulation temperature. Defaults to 1.0.
         dt: The simulation time step. Defaults to 1.0.
         A: An optional matrix giving the drift matrix. Defaults to a single
            value of 1.0.
         C: An optional matrix giving the covariance matrix. Defaults to an
            identity matrix times temperature with the same dimensions as the
            total number of degrees of freedom in the system.
         ethermo: The initial heat energy transferred to the bath.
            Defaults to 0.0. Will be non-zero if the thermostat is
            initialised from a checkpoint file.
      """

      super(ThermoNMGLE,self).__init__(temp,dt,ethermo)

      if A is None:
         A = np.identity(1,float)
      dset(self,"A",depend_value(value=A.copy(),name='A'))

      self.nb = len(self.A)
      self.ns = len(self.A[0]) - 1;

      # now, this is tricky. if C is taken from temp, then we want it to be
      # updated as a depend of temp.
      # Otherwise, we want it to be an independent beast.
      if C is None:
         dset(self,"C",depend_value(name='C', func=self.get_C, dependencies=[dget(self,"temp")]))
      else:
         dset(self,"C",depend_value(value=C.copy(),name='C'))

   def bind(self, nm=None, prng=None, fixdof=None):
      """Binds the appropriate degrees of freedom to the thermostat.

      This takes an object with degrees of freedom, and makes their momentum
      and mass vectors members of the thermostat. It also then creates the
      objects that will hold the data needed in the thermostat algorithms
      and the dependency network. Actually, this specific thermostat requires
      being called on a beads object.

      Args:
         nm: An optional normal modes object to take the mass and momentum
            vectors from.
         prng: An optional pseudo random number generator object. Defaults to
            Random().
         fixdof: An optional integer which can specify the number of constraints
            applied to the system. Defaults to zero.

      Raises:
         TypeError: Raised if no beads object is specified for
            the thermostat to couple to.
      """

      if nm is None or not type(nm) is NormalModes:
         raise TypeError("ThermoNMGLE.bind expects a NormalModes argument to bind to")

      if prng is None:
         self.prng = Random()
      else:
         self.prng = prng

      if (nm.nbeads != self.nb):
         raise IndexError("The parameters in nm_gle options correspond to a bead number "+str(self.nb)+ " which does not match the number of beads in the path" + str(nm.nbeads) )

      # allocates, initializes or restarts an array of s's
      if self.s.shape != (self.nb, self.ns + 1, nm.natoms *3) :
         if len(self.s) > 0:
            warning("Mismatch in GLE s array size on restart, will reinitialise to free particle.", verbosity.low)
         self.s = np.zeros((self.nb, self.ns + 1, nm.natoms*3))

         # Initializes the s vector in the free-particle limit
         info(" GLE additional DOFs initialised to the free-particle limit.", verbosity.low)
         for b in range(self.nb):
            SC = stab_cholesky(self.C[b]*Constants.kb)
            self.s[b] = np.dot(SC, self.prng.gvec(self.s[b].shape))
      else:
         info("GLE additional DOFs initialised from input.", verbosity.medium)

      prev_ethermo = self.ethermo

      # creates a set of thermostats to be applied to individual normal modes
      self._thermos = [ThermoGLE(temp=1, dt=1, A=self.A[b], C=self.C[b]) for b in range(self.nb)]

      # must pipe all the dependencies in such a way that values for the nm
      # thermostats are automatically updated based on the "master" thermostat
      def make_Agetter(k):
         return lambda: self.A[k]
      def make_Cgetter(k):
         return lambda: self.C[k]

      it = 0
      for t in self._thermos:
         t.s = self.s[it]  # gets the s's as a slice of self.s
         t.bind(pm=(nm.pnm[it,:],nm.dynm3[it,:]), prng=self.prng) # bind thermostat t to the it-th normal mode

         # pipes temp and dt
         deppipe(self,"temp", t, "temp")
         deppipe(self,"dt", t, "dt")

         # here we pipe the A and C of individual NM to the "master" arrays
         dget(t,"A").add_dependency(dget(self,"A"))
         dget(t,"A")._func = make_Agetter(it)
         dget(t,"C").add_dependency(dget(self,"C"))
         dget(t,"C")._func = make_Cgetter(it)
         dget(self,"ethermo").add_dependency(dget(t,"ethermo"))
         it += 1

      # since the ethermo will be "delegated" to the normal modes thermostats,
      # one has to split
      # any previously-stored value between the sub-thermostats
      for t in self._thermos:
         t.ethermo = prev_ethermo/self.nb

      dget(self,"ethermo")._func = self.get_ethermo;

   def step(self):
      """Updates the thermostat in NM representation by looping over the
      individual DOFs.
      """

      for t in self._thermos:
         t.step()

   def get_ethermo(self):
      """Computes the total energy transferred to the heat bath for all the nm
      thermostats.
      """

      et = 0.0;
      for t in self._thermos:
         et += t.ethermo
      return et

class ThermoNMGLEG(ThermoNMGLE):
   """Represents a 'normal-modes' GLE thermostat + SVR.

   An extension to the above NMGLE thermostat which also adds a stochastic velocity
   rescaling to the centroid.

   Depend objects:
      tau: Thermostat damping time scale. Larger values give a less strongly
         coupled thermostat.
   """

   def __init__(self, temp = 1.0, dt = 1.0, A = None, C = None, tau=1.0, ethermo=0.0):

      super(ThermoNMGLEG,self).__init__(temp, dt, A, C, ethermo)
      dset(self,"tau",depend_value(value=tau,name='tau'))

   def bind(self, nm=None, prng=None, fixdof=None):
      """Binds the appropriate degrees of freedom to the thermostat.

      This takes an object with degrees of freedom, and makes their momentum
      and mass vectors members of the thermostat. It also then creates the
      objects that will hold the data needed in the thermostat algorithms
      and the dependency network. Actually, this specific thermostat requires
      being called on a beads object.

      Args:
         nm: An optional normal modes object to take the mass and momentum
            vectors from.
         prng: An optional pseudo random number generator object. Defaults to
            Random().
         fixdof: An optional integer which can specify the number of constraints
            applied to the system. Defaults to zero.
      """

      super(ThermoNMGLEG,self).bind(nm, prng, fixdof)

      t = ThermoSVR(self.temp, self.dt, self.tau)

      t.bind(pm=(nm.pnm[0,:],nm.dynm3[0,:]), prng=self.prng) # bind global thermostat to centroid

      # pipes temp and dt
      deppipe(self,"temp", t, "temp")
      deppipe(self,"dt", t, "dt")
      deppipe(self,"tau", t, "tau")

      dget(self,"ethermo").add_dependency(dget(t,"ethermo"))
      self._thermos.append(t)


from rl_games.common import a2c_common
from rl_games.algos_torch import torch_ext

from rl_games.algos_torch import central_value
from rl_games.common import common_losses
from rl_games.common import datasets
from rl_games.algos_torch import ppg_aux
from rl_games.common.ewma_model import EwmaModel

from torch import optim
import torch
from torch import nn
import numpy as np
import gym

class A2CAgent(a2c_common.ContinuousA2CBase):
    def __init__(self, base_name, params):
        a2c_common.ContinuousA2CBase.__init__(self, base_name, params)
        obs_shape = self.obs_shape
        build_config = {
            'actions_num' : self.actions_num,
            'input_shape' : obs_shape,
            'num_seqs' : self.num_actors * self.num_agents,
            'value_size': self.env_info.get('value_size',1),
            'normalize_value' : self.normalize_value,
            'normalize_input': self.normalize_input,
        }

        self.model = self.network.build(build_config)
        self.model.to(self.ppo_device)
        self.states = None
        if self.ewma_ppo:
            self.ewma_model = EwmaModel(self.model, ewma_decay=0.889)
        self.init_rnn_from_model(self.model)
        self.last_lr = float(self.last_lr)
        self.bound_loss_type = self.config.get('bound_loss_type', 'bound') # 'regularisation' or 'bound'
        self.optimizer = optim.Adam(self.model.parameters(), float(self.last_lr), eps=1e-08, weight_decay=self.weight_decay)

        if self.has_central_value:
            cv_config = {
                'state_shape' : self.state_shape,
                'value_size' : self.value_size,
                'ppo_device' : self.ppo_device,
                'num_agents' : self.num_agents,
                'horizon_length' : self.horizon_length,
                'num_actors' : self.num_actors,
                'num_actions' : self.actions_num,
                'seq_len' : self.seq_len,
                'normalize_value' : self.normalize_value,
                'network' : self.central_value_config['network'],
                'config' : self.central_value_config,
                'writter' : self.writer,
                'max_epochs' : self.max_epochs,
                'multi_gpu' : self.multi_gpu
            }
            self.central_value_net = central_value.CentralValueTrain(**cv_config).to(self.ppo_device)

        self.use_experimental_cv = self.config.get('use_experimental_cv', True)
        self.dataset = datasets.PPODataset(self.batch_size, self.minibatch_size, self.is_discrete, self.is_rnn, self.ppo_device, self.seq_len)
        if self.normalize_value:
            self.value_mean_std = self.central_value_net.model.value_mean_std if self.has_central_value else self.model.value_mean_std
        if 'phasic_policy_gradients' in self.config:
            self.has_phasic_policy_gradients = True
            self.ppg_aux_loss = ppg_aux.PPGAux(self, self.config['phasic_policy_gradients'])
        self.has_value_loss = (self.has_central_value and self.use_experimental_cv) \
                            or (not self.has_phasic_policy_gradients and not self.has_central_value)
        self.algo_observer.after_init(self)
    def update_epoch(self):
        self.epoch_num += 1
        return self.epoch_num

    def save(self, fn):
        state = self.get_full_state_weights()
        torch_ext.save_checkpoint(fn, state)

    def restore(self, fn):
        checkpoint = torch_ext.load_checkpoint(fn)
        self.set_full_state_weights(checkpoint)

    def get_masked_action_values(self, obs, action_masks):
        assert False

    def calc_gradients(self, input_dict):
        value_preds_batch = input_dict['old_values']
        old_action_log_probs_batch = input_dict['old_logp_actions']
        advantage = input_dict['advantages']
        old_mu_batch = input_dict['mu']
        old_sigma_batch = input_dict['sigma']
        return_batch = input_dict['returns']
        actions_batch = input_dict['actions']

        # if self.modality_dict is not None:
        #     obs_batch = {}
        #     for modality_name in self.modality_dict:
        #         obs_batch[modality_name] = input_dict[modality_name]
        #     obs_batch = self._preproc_obs(obs_batch)
        #
        # else:
        obs_batch = input_dict['obs']
        obs_batch = self._preproc_obs(obs_batch)

        lr_mul = 1.0
        curr_e_clip = lr_mul * self.e_clip

        batch_dict = {
            'is_train': True,
            'prev_actions': actions_batch,
            'obs' : obs_batch,
        }

        rnn_masks = None
        if self.is_rnn:
            rnn_masks = input_dict['rnn_masks']
            batch_dict['rnn_states'] = input_dict['rnn_states']
            batch_dict['seq_length'] = self.seq_len

        with torch.cuda.amp.autocast(enabled=self.mixed_precision):
            res_dict = self.model(batch_dict)
            action_log_probs = res_dict['prev_neglogp']
            values = res_dict['values']
            entropy = res_dict['entropy']
            mu = res_dict['mus']
            sigma = res_dict['sigmas']

            if self.ewma_ppo:
                ewma_dict = self.ewma_model(batch_dict)
                proxy_neglogp = ewma_dict['prev_neglogp']
                a_loss = common_losses.decoupled_actor_loss(old_action_log_probs_batch, action_log_probs, proxy_neglogp, advantage, curr_e_clip)
                old_action_log_probs_batch = proxy_neglogp # to get right statistic later
                old_mu_batch = ewma_dict['mus']
                old_sigma_batch = ewma_dict['sigmas']
            else:
                a_loss = common_losses.actor_loss(old_action_log_probs_batch, action_log_probs, advantage, self.ppo, curr_e_clip)

            if self.has_value_loss:
                c_loss = common_losses.critic_loss(value_preds_batch, values, curr_e_clip, return_batch, self.clip_value)
            else:
                c_loss = torch.zeros(1, device=self.ppo_device)
            if self.bound_loss_type == 'regularisation':
                b_loss = self.reg_loss(mu)
            elif self.bound_loss_type == 'bound':
                b_loss = self.bound_loss(mu)
            else:
                b_loss = torch.zeros(1, device=self.ppo_device)
            losses, sum_mask = torch_ext.apply_masks([a_loss.unsqueeze(1), c_loss, entropy.unsqueeze(1), b_loss.unsqueeze(1)], rnn_masks)
            a_loss, c_loss, entropy, b_loss = losses[0], losses[1], losses[2], losses[3]

            loss = a_loss + 0.5 * c_loss * self.critic_coef - entropy * self.entropy_coef + b_loss * self.bounds_loss_coef

            if self.multi_gpu:
                self.optimizer.zero_grad()
            else:
                for param in self.model.parameters():
                    param.grad = None

        self.scaler.scale(loss).backward()
        #TODO: Refactor this ugliest code of they year
        if self.truncate_grads:
            if self.multi_gpu:
                self.optimizer.synchronize()
                self.scaler.unscale_(self.optimizer)
                nn.utils.clip_grad_norm_(self.model.parameters(), self.grad_norm)
                with self.optimizer.skip_synchronize():
                    self.scaler.step(self.optimizer)
                    self.scaler.update()
            else:
                self.scaler.unscale_(self.optimizer)
                nn.utils.clip_grad_norm_(self.model.parameters(), self.grad_norm)
                self.scaler.step(self.optimizer)
                self.scaler.update()
        else:
            self.scaler.step(self.optimizer)
            self.scaler.update()

        with torch.no_grad():
            reduce_kl = rnn_masks is None
            kl_dist = torch_ext.policy_kl(mu.detach(), sigma.detach(), old_mu_batch, old_sigma_batch, reduce_kl)
            if rnn_masks is not None:
                kl_dist = (kl_dist * rnn_masks).sum() / rnn_masks.numel()  #/ sum_mask

        if self.ewma_ppo:
            self.ewma_model.update()

        self.diagnostics.mini_batch(self,
        {
            'values' : value_preds_batch,
            'returns' : return_batch,
            'new_neglogp' : action_log_probs,
            'old_neglogp' : old_action_log_probs_batch,
            'masks' : rnn_masks
        }, curr_e_clip, 0)

        self.train_result = (a_loss, c_loss, entropy, \
            kl_dist, self.last_lr, lr_mul, \
            mu.detach(), sigma.detach(), b_loss)

    def train_actor_critic(self, input_dict):
        self.calc_gradients(input_dict)
        return self.train_result

    def reg_loss(self, mu):
        if self.bounds_loss_coef is not None:
            reg_loss = (mu*mu).sum(axis=-1)
        else:
            reg_loss = 0
        return reg_loss

    def bound_loss(self, mu):
        if self.bounds_loss_coef is not None:
            soft_bound = 1.1
            mu_loss_high = torch.clamp_min(mu - soft_bound, 0.0)**2
            mu_loss_low = torch.clamp_max(mu + soft_bound, 0.0)**2
            b_loss = (mu_loss_low + mu_loss_high).sum(axis=-1)
        else:
            b_loss = 0
        return b_loss


import pymlir
from ..utils.mlir_shell import mlir_quant
from ..utils.math_function import cal_sqnr
import cv2
import numpy as np
import sys
import os
import copy
import math
import time
from tqdm import tqdm
from cvi_toolkit.utils.log_setting import setup_logger
logger = setup_logger('root')

tpu_skip_op = ['tpu.input', 'tpu.quant', 'tpu.layer_norm', 'tpu.softmax'
               'tpu.split', 'tpu.reshape', 'tpu.gru']

class MixQuantModel:
    def __init__(self, fp32_mlir, all_bf16=False,
                 calib_table=None, mix_table=None):
        self.fp32_mlir = fp32_mlir
        self.all_bf16 = all_bf16
        self.calib_table = calib_table
        self.mix_table = mix_table
        self.quanted_mlir_file = '{}.quanted.tune.mlir'.format(fp32_mlir)
        self.model = self._build()

    def _build(self):
        ret = mlir_quant(self.fp32_mlir, self.quanted_mlir_file, "cv183x", "tmp.csv",
                         self.all_bf16, self.calib_table, self.mix_table)
        if ret != 0:
            raise RuntimeError("generate quanted mlir model failed")
        model = pymlir.module()
        model.load(self.quanted_mlir_file)
        return model

    def infer(self, data):
        if type(data) == np.lib.npyio.NpzFile:
            for k, v in data.items():
                self.model.set_tensor(k, v)
            self.model.invoke()
        else:
            self.model.run(data)
        outputs = {}
        for name in self.model.get_output_details():
            outputs[name] = self.model.get_tensor(name)
        return outputs

    def clean(self):
        try:
            mlir_weight_file = self.model.get_weight_file_path()
            del self.model
            os.remove(self.quanted_mlir_file)
            os.remove(mlir_weight_file)
        except:
            pass

class MixPrecSearcher(object):
    def __init__(self, mlir_file, calib_table, image_file_list=[],
                 skip_op=tpu_skip_op, precrocess_func=None, input_num=10):

        if len(image_file_list) == 0:
            raise RuntimeError("Please Set image file list!")

        self.images = image_file_list
        self.fp32_mlir = mlir_file
        self.calib_table = calib_table
        self.skip_ops = skip_op
        self.preprocess_func = precrocess_func
        self.input_data_buffer = dict()
        logger.info("[*] selected images->")
        for i, image in enumerate(self.images):
            logger.info("**** <{}> {}".format(i, image))

    def _dump_all_ops_to_file(self, target_file, ops, excludes=[]):
        with open(target_file, 'w') as f:
            for op in ops:
                if op['name'] not in excludes:
                    f.write(op['name'] + "\n")

    def _is_npz(self, image):
        return True if image.split('.')[-1] == 'npz' else False

    def _fetch_data(self):
        for image in self.images:
            if image not in self.input_data_buffer:
                if self._is_npz(image):
                    x = np.load(image)
                else:
                    x = self.preprocess_func(image)
                self.input_data_buffer[image] = x
            else:
                x = self.input_data_buffer[image]
            yield x

    def _loss(self, preds, gt_preds):
        ret = 0
        cnt = 0
        for op_name in gt_preds:
            a = gt_preds[op_name]
            b = preds[op_name]
            loss = cal_sqnr(a, b)
            if not math.isinf(loss):
                ret += -loss * a.size
                cnt += a.size

        if ret == 0 and cnt == 0:
            return -math.inf
        else:
            return ret / cnt

    def run(self):
        loss_list = list()
        predictions_gt = list()

        # set all layer for bf16
        bf16_model = MixQuantModel(self.fp32_mlir, all_bf16=True)
        for data in self._fetch_data():
            outputs = bf16_model.infer(data)
            predictions_gt.append(outputs)

        self.op_info = bf16_model.model.op_info

        pbar = tqdm(self.op_info)
        for op in pbar:
            pbar.set_description("Processing {}".format(op['name']))
            if op['type'] in self.skip_ops:
                continue

            mix_table = "tmp_mix_table.txt"
            self._dump_all_ops_to_file(mix_table, self.op_info, [op['name']])
            mix_model = MixQuantModel(self.fp32_mlir, False,
                                      self.calib_table, mix_table)
            loss = 0
            for idx, data in enumerate(self._fetch_data()):
                outputs = mix_model.infer(data)
                loss += self._loss(outputs, predictions_gt[idx])
            mix_model.clean()
            loss_list.append((op['name'], loss / len(self.images)))

        bf16_model.clean()
        return sorted(loss_list, key=lambda x: x[1], reverse=True)


"""High-throughput combinational logic timing simulators.

These simulators work similarly to :py:class:`~kyupy.logic_sim.LogicSim`.
They propagate values through the combinational circuit from (pseudo) primary inputs to (pseudo) primary outputs.
Instead of propagating logic values, these simulators propagate signal histories (waveforms).
They are designed to run many simulations in parallel and while their latencies are quite high, they can achieve
high throughput.

The simulators are not event-based and are not capable of simulating sequential circuits directly.

Two simulators are available: :py:class:`WaveSim` runs on the CPU, and the derived class
:py:class:`WaveSimCuda` runs on the GPU.
"""

import math
from bisect import bisect, insort_left

import numpy as np

from . import numba, cuda, hr_bytes

TMAX = np.float32(2 ** 127)
"""A large 32-bit floating point value used to mark the end of a waveform."""
TMAX_OVL = np.float32(1.1 * 2 ** 127)
"""A large 32-bit floating point value used to mark the end of a waveform that
may be incomplete due to an overflow."""
TMIN = np.float32(-2 ** 127)
"""A large negative 32-bit floating point value used at the beginning of waveforms that start with logic-1."""

class Heap:
    def __init__(self):
        self.chunks = dict()  # map start location to chunk size
        self.released = list()  # chunks that were released
        self.current_size = 0
        self.max_size = 0

    def alloc(self, size):
        for idx, loc in enumerate(self.released):
            if self.chunks[loc] == size:
                del self.released[idx]
                return loc
            if self.chunks[loc] > size:  # split chunk
                chunksize = self.chunks[loc]
                self.chunks[loc] = size
                self.chunks[loc + size] = chunksize - size
                self.released[idx] = loc + size  # move released pointer: loc -> loc+size
                return loc
        # no previously released chunk; make new one
        loc = self.current_size
        self.chunks[loc] = size
        self.current_size += size
        self.max_size = max(self.max_size, self.current_size)
        return loc

    def free(self, loc):
        size = self.chunks[loc]
        if loc + size == self.current_size:  # end of managed area, remove chunk
            del self.chunks[loc]
            self.current_size -= size
            # check and remove prev chunk if free
            if len(self.released) > 0:
                prev = self.released[-1]
                if prev + self.chunks[prev] == self.current_size:
                    chunksize = self.chunks[prev]
                    del self.chunks[prev]
                    del self.released[-1]
                    self.current_size -= chunksize
            return
        released_idx = bisect(self.released, loc)
        if released_idx < len(self.released) and loc + size == self.released[released_idx]:  # next chunk is free, merge
            chunksize = size + self.chunks[loc + size]
            del self.chunks[loc + size]
            self.chunks[loc] = chunksize
            size = self.chunks[loc]
            self.released[released_idx] = loc
        else:
            insort_left(self.released, loc)  # put in a new release
        if released_idx > 0:  # check if previous chunk is free
            prev = self.released[released_idx - 1]
            if prev + self.chunks[prev] == loc:  # previous chunk is adjacent to freed one, merge
                chunksize = size + self.chunks[prev]
                del self.chunks[loc]
                self.chunks[prev] = chunksize
                del self.released[released_idx]

    def __repr__(self):
        r = []
        for loc in sorted(self.chunks.keys()):
            size = self.chunks[loc]
            released_idx = bisect(self.released, loc)
            is_released = released_idx > 0 and len(self.released) > 0 and self.released[released_idx - 1] == loc
            r.append(f'{loc:5d}: {"free" if is_released else "used"} {size}')
        return "\n".join(r)

class WaveSim:
    """A waveform-based combinational logic timing simulator running on CPU.

    :param circuit: The circuit to simulate.
    :param timing: The timing annotation of the circuit (see :py:func:`kyupy.sdf.DelayFile.annotation` for details)
    :param sims: The number of parallel simulations.
    :param wavecaps: The number of floats available in each waveform. Waveforms are encoding the signal switching
        history by storing transition times. The waveform capacity roughly corresponds to the number of transitions
        that can be stored. A capacity of ``n`` can store at least ``n-2`` transitions. If more transitions are
        generated during simulation, the latest glitch is removed (freeing up two transition times) and an overflow
        flag is set. If an integer is given, all waveforms are set to that same capacity. With an array of length
        ``len(circuit.lines)`` the capacity can be controlled for each intermediate waveform individually.
    :param strip_forks: If enabled, the simulator will not evaluate fork nodes explicitly. This saves simulation time
        by reducing the number of nodes to simulate, but (interconnect) delay annotations of lines read by fork nodes
        are ignored.
    :param keep_waveforms: If disabled, memory of intermediate signal waveforms will be re-used. This greatly reduces
        memory footprint, but intermediate signal waveforms become unaccessible after a propagation.
    """
    def __init__(self, circuit, timing, sims=8, wavecaps=16, strip_forks=False, keep_waveforms=True):
        self.circuit = circuit
        self.sims = sims
        self.overflows = 0
        self.interface = list(circuit.interface) + [n for n in circuit.nodes if 'dff' in n.kind.lower()]

        self.lst_eat_valid = False

        self.cdata = np.zeros((len(self.interface), sims, 7), dtype='float32')

        if isinstance(wavecaps, int):
            wavecaps = [wavecaps] * len(circuit.lines)

        intf_wavecap = 4  # sufficient for storing only 1 transition.

        # indices for state allocation table (sat)
        self.zero_idx = len(circuit.lines)
        self.tmp_idx = self.zero_idx + 1
        self.ppi_offset = self.tmp_idx + 1
        self.ppo_offset = self.ppi_offset + len(self.interface)
        self.sat_length = self.ppo_offset + len(self.interface)

        # translate circuit structure into self.ops
        ops = []
        interface_dict = dict((n, i) for i, n in enumerate(self.interface))
        for n in circuit.topological_order():
            if n in interface_dict:
                inp_idx = self.ppi_offset + interface_dict[n]
                if len(n.outs) > 0 and n.outs[0] is not None:  # first output of a PI/PPI
                    ops.append((0b1010, n.outs[0].index, inp_idx, self.zero_idx))
                if 'dff' in n.kind.lower():  # second output of DFF is inverted
                    if len(n.outs) > 1 and n.outs[1] is not None:
                        ops.append((0b0101, n.outs[1].index, inp_idx, self.zero_idx))
                else:  # if not DFF, no output is inverted.
                    for o_line in n.outs[1:]:
                        if o_line is not None:
                            ops.append((0b1010, o_line.index, inp_idx, self.zero_idx))
            else:  # regular node, not PI/PPI or PO/PPO
                o0_idx = n.outs[0].index if len(n.outs) > 0 and n.outs[0] is not None else self.tmp_idx
                i0_idx = n.ins[0].index if len(n.ins) > 0 and n.ins[0] is not None else self.zero_idx
                i1_idx = n.ins[1].index if len(n.ins) > 1 and n.ins[1] is not None else self.zero_idx
                kind = n.kind.lower()
                if kind == '__fork__':
                    if not strip_forks:
                        for o_line in n.outs:
                            ops.append((0b1010, o_line.index, i0_idx, i1_idx))
                elif kind.startswith('nand'):
                    ops.append((0b0111, o0_idx, i0_idx, i1_idx))
                elif kind.startswith('nor'):
                    ops.append((0b0001, o0_idx, i0_idx, i1_idx))
                elif kind.startswith('and'):
                    ops.append((0b1000, o0_idx, i0_idx, i1_idx))
                elif kind.startswith('or'):
                    ops.append((0b1110, o0_idx, i0_idx, i1_idx))
                elif kind.startswith('xor'):
                    ops.append((0b0110, o0_idx, i0_idx, i1_idx))
                elif kind.startswith('xnor'):
                    ops.append((0b1001, o0_idx, i0_idx, i1_idx))
                elif kind.startswith('not') or kind.startswith('inv') or kind.startswith('ibuf'):
                    ops.append((0b0101, o0_idx, i0_idx, i1_idx))
                elif kind.startswith('buf') or kind.startswith('nbuf'):
                    ops.append((0b1010, o0_idx, i0_idx, i1_idx))
                elif kind.startswith('__const1__') or kind.startswith('tieh'):
                    ops.append((0b0101, o0_idx, i0_idx, i1_idx))
                elif kind.startswith('__const0__') or kind.startswith('tiel'):
                    ops.append((0b1010, o0_idx, i0_idx, i1_idx))
                else:
                    print('unknown gate type', kind)
        self.ops = np.asarray(ops, dtype='int32')

        # create a map from fanout lines to stem lines for fork stripping
        stems = np.zeros(self.sat_length, dtype='int32') - 1  # default to -1: 'no fanout line'
        if strip_forks:
            for f in circuit.forks.values():
                prev_line = f.ins[0]
                while prev_line.driver.kind == '__fork__':
                    prev_line = prev_line.driver.ins[0]
                stem_idx = prev_line.index
                for ol in f.outs:
                    stems[ol] = stem_idx

        # calculate level (distance from PI/PPI) and reference count for each line
        levels = np.zeros(self.sat_length, dtype='int32')
        ref_count = np.zeros(self.sat_length, dtype='int32')
        level_starts = [0]
        current_level = 1
        for i, op in enumerate(self.ops):
            # if we fork-strip, always take the stems for determining fan-in level
            i0_idx = stems[op[2]] if stems[op[2]] >= 0 else op[2]
            i1_idx = stems[op[3]] if stems[op[3]] >= 0 else op[3]
            if levels[i0_idx] >= current_level or levels[i1_idx] >= current_level:
                current_level += 1
                level_starts.append(i)
            levels[op[1]] = current_level  # set level of the output line
            ref_count[i0_idx] += 1
            ref_count[i1_idx] += 1
        self.level_starts = np.asarray(level_starts, dtype='int32')
        self.level_stops = np.asarray(level_starts[1:] + [len(self.ops)], dtype='int32')

        # state allocation table. maps line and interface indices to self.state memory locations
        self.sat = np.zeros((self.sat_length, 3), dtype='int')
        self.sat[:, 0] = -1

        h = Heap()

        # allocate and keep memory for special fields
        self.sat[self.zero_idx] = h.alloc(intf_wavecap), intf_wavecap, 0
        self.sat[self.tmp_idx] = h.alloc(intf_wavecap), intf_wavecap, 0
        ref_count[self.zero_idx] += 1
        ref_count[self.tmp_idx] += 1

        # allocate and keep memory for PI/PPI, keep memory for PO/PPO (allocated later)
        for i, n in enumerate(self.interface):
            if len(n.outs) > 0:
                self.sat[self.ppi_offset + i] = h.alloc(intf_wavecap), intf_wavecap, 0
                ref_count[self.ppi_offset + i] += 1
            if len(n.ins) > 0:
                i0_idx = stems[n.ins[0]] if stems[n.ins[0]] >= 0 else n.ins[0]
                ref_count[i0_idx] += 1

        # allocate memory for the rest of the circuit
        for op_start, op_stop in zip(self.level_starts, self.level_stops):
            free_list = []
            for op in self.ops[op_start:op_stop]:
                # if we fork-strip, always take the stems
                i0_idx = stems[op[2]] if stems[op[2]] >= 0 else op[2]
                i1_idx = stems[op[3]] if stems[op[3]] >= 0 else op[3]
                ref_count[i0_idx] -= 1
                ref_count[i1_idx] -= 1
                if ref_count[i0_idx] <= 0: free_list.append(self.sat[i0_idx, 0])
                if ref_count[i1_idx] <= 0: free_list.append(self.sat[i1_idx, 0])
                o_idx = op[1]
                cap = wavecaps[o_idx]
                self.sat[o_idx] = h.alloc(cap), cap, 0
            if not keep_waveforms:
                for loc in free_list:
                    h.free(loc)

        # copy memory location and capacity from stems to fanout lines
        for lidx, stem in enumerate(stems):
            if stem >= 0:  # if at a fanout line
                self.sat[lidx] = self.sat[stem]

        # copy memory location to PO/PPO area
        for i, n in enumerate(self.interface):
            if len(n.ins) > 0:
                self.sat[self.ppo_offset + i] = self.sat[n.ins[0]]

        # pad timing
        self.timing = np.zeros((self.sat_length, 2, 2))
        self.timing[:len(timing)] = timing

        # allocate self.state
        self.state = np.zeros((h.max_size, sims), dtype='float32') + TMAX

        m1 = np.array([2 ** x for x in range(7, -1, -1)], dtype='uint8')
        m0 = ~m1
        self.mask = np.rollaxis(np.vstack((m0, m1)), 1)

    def __repr__(self):
        total_mem = self.state.nbytes + self.sat.nbytes + self.ops.nbytes + self.cdata.nbytes
        return f'<WaveSim {self.circuit.name} sims={self.sims} ops={len(self.ops)} ' + \
               f'levels={len(self.level_starts)} mem={hr_bytes(total_mem)}>'

    def get_line_delay(self, line, polarity):
        """Returns the current delay of the given ``line`` and ``polarity`` in the simulation model."""
        return self.timing[line, 0, polarity]

    def set_line_delay(self, line, polarity, delay):
        """Sets a new ``delay`` for the given ``line`` and ``polarity`` in the simulation model."""
        self.timing[line, 0, polarity] = delay

    def assign(self, vectors, time=0.0, offset=0):
        """Assigns new values to the primary inputs and state-elements.

        :param vectors: The values to assign preferably in 8-valued logic. The values are converted to
            appropriate waveforms with or one transition (``RISE``, ``FALL``) no transitions
            (``ZERO``, ``ONE``, and others).
        :type vectors: :py:class:`~kyupy.logic.BPArray`
        :param time: The transition time of the generated waveforms.
        :param offset: The offset into the vector set. The vector assigned to the first simulator is
            ``vectors[offset]``.
        """
        nvectors = min(len(vectors) - offset, self.sims)
        for i in range(len(self.interface)):
            ppi_loc = self.sat[self.ppi_offset + i, 0]
            if ppi_loc < 0: continue
            for p in range(nvectors):
                vector = p + offset
                a = vectors.data[i, :, vector // 8]
                m = self.mask[vector % 8]
                toggle = 0
                if len(a) <= 2:
                    if a[0] & m[1]:
                        self.state[ppi_loc, p] = TMIN
                        toggle += 1
                else:
                    if a[1] & m[1]:
                        self.state[ppi_loc, p] = TMIN
                        toggle += 1
                    if (a[2] & m[1]) and ((a[0] & m[1]) != (a[1] & m[1])):
                        self.state[ppi_loc + toggle, p] = time
                        toggle += 1
                self.state[ppi_loc + toggle, p] = TMAX

    def propagate(self, sims=None, sd=0.0, seed=1):
        """Propagates all waveforms from the (pseudo) primary inputs to the (pseudo) primary outputs.

        :param sims: Number of parallel simulations to execute. If None, all available simulations are performed.
        :param sd: Standard deviation for injection of random delay variation. Active, if value is positive.
        :param seed: Random seed for delay variations.
        """
        sims = min(sims or self.sims, self.sims)
        for op_start, op_stop in zip(self.level_starts, self.level_stops):
            self.overflows += level_eval(self.ops, op_start, op_stop, self.state, self.sat, 0, sims,
                                         self.timing, sd, seed)
        self.lst_eat_valid = False

    def wave(self, line, vector):
        # """Returns the desired waveform from the simulation state. Only valid, if simulator was
        # instantiated with ``keep_waveforms=True``."""
        if line < 0:
            return [TMAX]
        mem, wcap, _ = self.sat[line]
        if mem < 0:
            return [TMAX]
        return self.state[mem:mem + wcap, vector]

    def wave_ppi(self, i, vector):
        return self.wave(self.ppi_offset + i, vector)

    def wave_ppo(self, o, vector):
        return self.wave(self.ppo_offset + o, vector)

    def capture(self, time=TMAX, sd=0.0, seed=1, cdata=None, offset=0):
        """Simulates a capture operation at all state-elements and primary outputs.

        The capture analyzes the propagated waveforms at and around the given capture time and returns
        various results for each capture operation.

        :param time: The desired capture time. By default, a capture of the settled value is performed.
        :param sd: A standard deviation for uncertainty in the actual capture time.
        :param seed: The random seed for a capture with uncertainty.
        :param cdata: An array to copy capture data into (optional). See the return value for details.
        :param offset: An offset into the supplied capture data array.
        :return: The capture data as numpy array.

            The 3-dimensional capture data array contains for each interface node (axis 0),
            and each test (axis 1), seven values:

            0. Probability of capturing a 1 at the given capture time (same as next value, if no
               standard deviation given).
            1. A capture value decided by random sampling according to above probability and given seed.
            2. The final value (assume a very late capture time).
            3. True, if there was a premature capture (capture error), i.e. final value is different
               from captured value.
            4. Earliest arrival time. The time at which the output transitioned from its initial value.
            5. Latest stabilization time. The time at which the output transitioned to its final value.
            6. Overflow indicator. If non-zero, some signals in the input cone of this output had more
               transitions than specified in ``wavecaps``. Some transitions have been discarded, the
               final values in the waveforms are still valid.
        """
        for i, node in enumerate(self.interface):
            if len(node.ins) == 0: continue
            for p in range(self.sims):
                self.cdata[i, p] = self.capture_wave(self.ppo_offset + i, p, time, sd, seed)
        if cdata is not None:
            assert offset < cdata.shape[1]
            cap_dim = min(cdata.shape[1] - offset, self.sims)
            cdata[:, offset:cap_dim + offset] = self.cdata[:, 0:cap_dim]
        self.lst_eat_valid = True
        return self.cdata

    def reassign(self, time=0.0):
        """Re-assigns the last capture to the appropriate pseudo-primary inputs. Generates a new set of
        waveforms at the PPIs that start with the previous final value of that PPI, and transitions at the
        given time to the value captured in a previous simulation. :py:func:`~WaveSim.capture` must be called
        prior to this function. The final value of each PPI is taken from the randomly sampled concrete logic
        values in the capture data.

        :param time: The transition time at the inputs (usually 0.0).
        """
        for i in range(len(self.interface)):
            ppi_loc = self.sat[self.ppi_offset + i, 0]
            ppo_loc = self.sat[self.ppo_offset + i, 0]
            if ppi_loc < 0 or ppo_loc < 0: continue
            for sidx in range(self.sims):
                ival = self.val(self.ppi_offset + i, sidx, TMAX) > 0.5
                oval = self.cdata[i, sidx, 1] > 0.5
                toggle = 0
                if ival:
                    self.state[ppi_loc, sidx] = TMIN
                    toggle += 1
                if ival != oval:
                    self.state[ppi_loc + toggle, sidx] = time
                    toggle += 1
                self.state[ppi_loc + toggle, sidx] = TMAX

    def eat(self, line, vector):
        eat = TMAX
        for t in self.wave(line, vector):
            if t >= TMAX: break
            if t <= TMIN: continue
            eat = min(eat, t)
        return eat

    def lst(self, line, vector):
        lst = TMIN
        for t in self.wave(line, vector):
            if t >= TMAX: break
            if t <= TMIN: continue
            lst = max(lst, t)
        return lst

    def lst_ppo(self, o, vector):
        if not self.lst_eat_valid:
            self.capture()
        return self.cdata[o, vector, 5]

    def toggles(self, line, vector):
        tog = 0
        for t in self.wave(line, vector):
            if t >= TMAX: break
            if t <= TMIN: continue
            tog += 1
        return tog

    def _vals(self, idx, vector, times, sd=0.0):
        s_sqrt2 = sd * math.sqrt(2)
        m = 0.5
        accs = [0.0] * len(times)
        values = [0] * len(times)
        for t in self.wave(idx, vector):
            if t >= TMAX: break
            for idx, time in enumerate(times):
                if t < time:
                    values[idx] = values[idx] ^ 1
            m = -m
            if t <= TMIN: continue
            if s_sqrt2 > 0:
                for idx, time in enumerate(times):
                    accs[idx] += m * (1 + math.erf((t - time) / s_sqrt2))
        if (m < 0) and (s_sqrt2 > 0):
            for idx, time in enumerate(times):
                accs[idx] += 1
        if s_sqrt2 == 0:
            return values
        return accs

    def vals(self, line, vector, times, sd=0):
        return self._vals(line, vector, times, sd)

    def val(self, line, vector, time=TMAX, sd=0):
        return self.capture_wave(line, vector, time, sd)[0]

    def vals_ppo(self, o, vector, times, sd=0):
        return self._vals(self.ppo_offset + o, vector, times, sd)

    def val_ppo(self, o, vector, time=TMAX, sd=0):
        if not self.lst_eat_valid:
            self.capture(time, sd)
        return self.cdata[o, vector, 0]

    def capture_wave(self, line, vector, time=TMAX, sd=0.0, seed=1):
        s_sqrt2 = sd * math.sqrt(2)
        m = 0.5
        acc = 0.0
        eat = TMAX
        lst = TMIN
        tog = 0
        ovl = 0
        val = int(0)
        final = int(0)
        for t in self.wave(line, vector):
            if t >= TMAX:
                if t == TMAX_OVL:
                    ovl = 1
                break
            m = -m
            final ^= 1
            if t < time:
                val ^= 1
            if t <= TMIN: continue
            if s_sqrt2 > 0:
                acc += m * (1 + math.erf((t - time) / s_sqrt2))
            eat = min(eat, t)
            lst = max(lst, t)
            tog += 1
        if s_sqrt2 > 0:
            if m < 0:
                acc += 1
            if acc >= 0.99:
                val = 1
            elif acc > 0.01:
                seed = (seed << 4) + (vector << 20) + (line-self.ppo_offset << 1)
                seed = int(0xDEECE66D) * seed + 0xB
                seed = int(0xDEECE66D) * seed + 0xB
                rnd = float((seed >> 8) & 0xffffff) / float(1 << 24)
                val = rnd < acc
            else:
                val = 0
        else:
            acc = val

        return acc, val, final, (val != final), eat, lst, ovl

@numba.njit
def level_eval(ops, op_start, op_stop, state, sat, st_start, st_stop, line_times, sd, seed):
    overflows = 0
    for op_idx in range(op_start, op_stop):
        op = ops[op_idx]
        for st_idx in range(st_start, st_stop):
            overflows += wave_eval(op, state, sat, st_idx, line_times, sd, seed)
    return overflows

@numba.njit
def rand_gauss(seed, sd):
    clamp = 0.5
    if sd <= 0.0:
        return 1.0
    while True:
        x = -6.0
        for _ in range(12):
            seed = int(0xDEECE66D) * seed + 0xB
            x += float((seed >> 8) & 0xffffff) / float(1 << 24)
        x *= sd
        if abs(x) <= clamp:
            break
    return x + 1.0

@numba.njit
def wave_eval(op, state, sat, st_idx, line_times, sd=0.0, seed=0):
    lut, z_idx, a_idx, b_idx = op
    overflows = int(0)

    _seed = (seed << 4) + (z_idx << 20) + (st_idx << 1)

    a_mem = sat[a_idx, 0]
    b_mem = sat[b_idx, 0]
    z_mem, z_cap, _ = sat[z_idx]

    a_cur = int(0)
    b_cur = int(0)
    z_cur = lut & 1
    if z_cur == 1:
        state[z_mem, st_idx] = TMIN

    a = state[a_mem, st_idx] + line_times[a_idx, 0, z_cur] * rand_gauss(_seed ^ a_mem ^ z_cur, sd)
    b = state[b_mem, st_idx] + line_times[b_idx, 0, z_cur] * rand_gauss(_seed ^ b_mem ^ z_cur, sd)

    previous_t = TMIN

    current_t = min(a, b)
    inputs = int(0)

    while current_t < TMAX:
        z_val = z_cur & 1
        if b < a:
            b_cur += 1
            b = state[b_mem + b_cur, st_idx]
            b += line_times[b_idx, 0, z_val ^ 1] * rand_gauss(_seed ^ b_mem ^ z_val ^ 1, sd)
            thresh = line_times[b_idx, 1, z_val] * rand_gauss(_seed ^ b_mem ^ z_val, sd)
            inputs ^= 2
            next_t = b
        else:
            a_cur += 1
            a = state[a_mem + a_cur, st_idx]
            a += line_times[a_idx, 0, z_val ^ 1] * rand_gauss(_seed ^ a_mem ^ z_val ^ 1, sd)
            thresh = line_times[a_idx, 1, z_val] * rand_gauss(_seed ^ a_mem ^ z_val, sd)
            inputs ^= 1
            next_t = a

        if (z_cur & 1) != ((lut >> inputs) & 1):
            # we generate a toggle in z_mem, if:
            #   ( it is the first toggle in z_mem OR
            #   following toggle is earlier OR
            #   pulse is wide enough ) AND enough space in z_mem.
            if z_cur == 0 or next_t < current_t or (current_t - previous_t) > thresh:
                if z_cur < (z_cap - 1):
                    state[z_mem + z_cur, st_idx] = current_t
                    previous_t = current_t
                    z_cur += 1
                else:
                    overflows += 1
                    previous_t = state[z_mem + z_cur - 1, st_idx]
                    z_cur -= 1
            else:
                z_cur -= 1
                if z_cur > 0:
                    previous_t = state[z_mem + z_cur - 1, st_idx]
                else:
                    previous_t = TMIN
        current_t = min(a, b)

    if overflows > 0:
        state[z_mem + z_cur, st_idx] = TMAX_OVL
    else:
        state[z_mem + z_cur, st_idx] = a if a > b else b  # propagate overflow flags by storing biggest TMAX from input

    return overflows

class WaveSimCuda(WaveSim):
    """A GPU-accelerated waveform-based combinational logic timing simulator.

    The API is the same as for :py:class:`WaveSim`.
    All internal memories are mirrored into GPU memory upon construction.
    Some operations like access to single waveforms can involve large communication overheads.
    """
    def __init__(self, circuit, timing, sims=8, wavecaps=16, strip_forks=False, keep_waveforms=True):
        super().__init__(circuit, timing, sims, wavecaps, strip_forks, keep_waveforms)

        self.tdata = np.zeros((len(self.interface), 3, (sims - 1) // 8 + 1), dtype='uint8')

        self.d_state = cuda.to_device(self.state)
        self.d_sat = cuda.to_device(self.sat)
        self.d_ops = cuda.to_device(self.ops)
        self.d_timing = cuda.to_device(self.timing)
        self.d_tdata = cuda.to_device(self.tdata)
        self.d_cdata = cuda.to_device(self.cdata)

        self._block_dim = (32, 16)

    def __repr__(self):
        total_mem = self.state.nbytes + self.sat.nbytes + self.ops.nbytes + self.timing.nbytes + \
                    self.tdata.nbytes + self.cdata.nbytes
        return f'<WaveSimCuda {self.circuit.name} sims={self.sims} ops={len(self.ops)} ' + \
               f'levels={len(self.level_starts)} mem={hr_bytes(total_mem)}>'

    def get_line_delay(self, line, polarity):
        return self.d_timing[line, 0, polarity]

    def set_line_delay(self, line, polarity, delay):
        self.d_timing[line, 0, polarity] = delay

    def assign(self, vectors, time=0.0, offset=0):
        assert (offset % 8) == 0
        byte_offset = offset // 8
        assert byte_offset < vectors.data.shape[-1]
        pdim = min(vectors.data.shape[-1] - byte_offset, self.tdata.shape[-1])

        self.tdata[..., 0:pdim] = vectors.data[..., byte_offset:pdim + byte_offset]
        if vectors.m == 2:
            self.tdata[:, 2, 0:pdim] = 0
        cuda.to_device(self.tdata, to=self.d_tdata)

        grid_dim = self._grid_dim(self.sims, len(self.interface))
        assign_kernel[grid_dim, self._block_dim](self.d_state, self.d_sat, self.ppi_offset,
                                                 len(self.interface), self.d_tdata, time)

    def _grid_dim(self, x, y):
        gx = math.ceil(x / self._block_dim[0])
        gy = math.ceil(y / self._block_dim[1])
        return gx, gy

    def propagate(self, sims=None, sd=0.0, seed=1):
        sims = min(sims or self.sims, self.sims)
        for op_start, op_stop in zip(self.level_starts, self.level_stops):
            grid_dim = self._grid_dim(sims, op_stop - op_start)
            wave_kernel[grid_dim, self._block_dim](self.d_ops, op_start, op_stop, self.d_state, self.sat, int(0),
                                                   sims, self.d_timing, sd, seed)
        cuda.synchronize()
        self.lst_eat_valid = False

    def wave(self, line, vector):
        if line < 0:
            return [TMAX]
        mem, wcap, _ = self.sat[line]
        if mem < 0:
            return [TMAX]
        return self.d_state[mem:mem + wcap, vector]

    def capture(self, time=TMAX, sd=0, seed=1, cdata=None, offset=0):
        grid_dim = self._grid_dim(self.sims, len(self.interface))
        capture_kernel[grid_dim, self._block_dim](self.d_state, self.d_sat, self.ppo_offset,
                                                  self.d_cdata, time, sd * math.sqrt(2), seed)
        self.cdata[...] = self.d_cdata
        if cdata is not None:
            assert offset < cdata.shape[1]
            cap_dim = min(cdata.shape[1] - offset, self.sims)
            cdata[:, offset:cap_dim + offset] = self.cdata[:, 0:cap_dim]
        self.lst_eat_valid = True
        return self.cdata

    def reassign(self, time=0.0):
        grid_dim = self._grid_dim(self.sims, len(self.interface))
        reassign_kernel[grid_dim, self._block_dim](self.d_state, self.d_sat, self.ppi_offset, self.ppo_offset,
                                                   self.d_cdata, time)
        cuda.synchronize()

    def wavecaps(self):
        gx = math.ceil(len(self.circuit.lines) / 512)
        wavecaps_kernel[gx, 512](self.d_state, self.d_sat, self.sims)
        self.sat[...] = self.d_sat
        return self.sat[..., 2]

@cuda.jit()
def wavecaps_kernel(state, sat, sims):
    idx = cuda.grid(1)
    if idx >= len(sat): return

    lidx, lcap, _ = sat[idx]
    if lidx < 0: return

    wcap = 0
    for sidx in range(sims):
        for tidx in range(lcap):
            t = state[lidx + tidx, sidx]
            if tidx > wcap:
                wcap = tidx
            if t >= TMAX: break

    sat[idx, 2] = wcap + 1

@cuda.jit()
def reassign_kernel(state, sat, ppi_offset, ppo_offset, cdata, ppi_time):
    vector, y = cuda.grid(2)
    if vector >= state.shape[-1]: return
    if ppo_offset + y >= len(sat): return

    ppo, _, _ = sat[ppo_offset + y]
    ppi, ppi_cap, _ = sat[ppi_offset + y]
    if ppo < 0: return
    if ppi < 0: return

    ppo_val = int(cdata[y, vector, 1])
    ppi_val = int(0)
    for tidx in range(ppi_cap):
        t = state[ppi + tidx, vector]
        if t >= TMAX: break
        ppi_val ^= 1

    # make new waveform at PPI
    toggle = 0
    if ppi_val:
        state[ppi + toggle, vector] = TMIN
        toggle += 1
    if ppi_val != ppo_val:
        state[ppi + toggle, vector] = ppi_time
        toggle += 1
    state[ppi + toggle, vector] = TMAX

@cuda.jit()
def capture_kernel(state, sat, ppo_offset, cdata, time, s_sqrt2, seed):
    x, y = cuda.grid(2)
    if ppo_offset + y >= len(sat): return
    line, tdim, _ = sat[ppo_offset + y]
    if line < 0: return
    if x >= state.shape[-1]: return
    vector = x
    m = 0.5
    acc = 0.0
    eat = TMAX
    lst = TMIN
    tog = 0
    ovl = 0
    val = int(0)
    final = int(0)
    for tidx in range(tdim):
        t = state[line + tidx, vector]
        if t >= TMAX:
            if t == TMAX_OVL:
                ovl = 1
            break
        m = -m
        final ^= 1
        if t < time:
            val ^= 1
        if t <= TMIN: continue
        if s_sqrt2 > 0:
            acc += m * (1 + math.erf((t - time) / s_sqrt2))
        eat = min(eat, t)
        lst = max(lst, t)
        tog += 1
    if s_sqrt2 > 0:
        if m < 0:
            acc += 1
        if acc >= 0.99:
            val = 1
        elif acc > 0.01:
            seed = (seed << 4) + (vector << 20) + (y << 1)
            seed = int(0xDEECE66D) * seed + 0xB
            seed = int(0xDEECE66D) * seed + 0xB
            rnd = float((seed >> 8) & 0xffffff) / float(1 << 24)
            val = rnd < acc
        else:
            val = 0
    else:
        acc = val

    cdata[y, vector, 0] = acc
    cdata[y, vector, 1] = val
    cdata[y, vector, 2] = final
    cdata[y, vector, 3] = (val != final)
    cdata[y, vector, 4] = eat
    cdata[y, vector, 5] = lst
    cdata[y, vector, 6] = ovl

@cuda.jit()
def assign_kernel(state, sat, ppi_offset, intf_len, tdata, time):
    x, y = cuda.grid(2)
    if y >= intf_len: return
    line = sat[ppi_offset + y, 0]
    if line < 0: return
    sdim = state.shape[-1]
    if x >= sdim: return
    vector = x
    a0 = tdata[y, 0, vector // 8]
    a1 = tdata[y, 1, vector // 8]
    a2 = tdata[y, 2, vector // 8]
    m = np.uint8(1 << (7 - (vector % 8)))
    toggle = 0
    if a1 & m:
        state[line + toggle, x] = TMIN
        toggle += 1
    if (a2 & m) and ((a0 & m) != (a1 & m)):
        state[line + toggle, x] = time
        toggle += 1
    state[line + toggle, x] = TMAX

@cuda.jit(device=True)
def rand_gauss_dev(seed, sd):
    clamp = 0.5
    if sd <= 0.0:
        return 1.0
    while True:
        x = -6.0
        for _ in range(12):
            seed = int(0xDEECE66D) * seed + 0xB
            x += float((seed >> 8) & 0xffffff) / float(1 << 24)
        x *= sd
        if abs(x) <= clamp:
            break
    return x + 1.0

@cuda.jit()
def wave_kernel(ops, op_start, op_stop, state, sat, st_start, st_stop, line_times, sd, seed):
    x, y = cuda.grid(2)
    st_idx = st_start + x
    op_idx = op_start + y
    if st_idx >= st_stop: return
    if op_idx >= op_stop: return
    lut = ops[op_idx, 0]
    z_idx = ops[op_idx, 1]
    a_idx = ops[op_idx, 2]
    b_idx = ops[op_idx, 3]
    overflows = int(0)

    _seed = (seed << 4) + (z_idx << 20) + (st_idx << 1)

    a_mem = sat[a_idx, 0]
    b_mem = sat[b_idx, 0]
    z_mem, z_cap, _ = sat[z_idx]

    a_cur = int(0)
    b_cur = int(0)
    z_cur = lut & 1
    if z_cur == 1:
        state[z_mem, st_idx] = TMIN

    a = state[a_mem, st_idx] + line_times[a_idx, 0, z_cur] * rand_gauss_dev(_seed ^ a_mem ^ z_cur, sd)
    b = state[b_mem, st_idx] + line_times[b_idx, 0, z_cur] * rand_gauss_dev(_seed ^ b_mem ^ z_cur, sd)

    previous_t = TMIN

    current_t = min(a, b)
    inputs = int(0)

    while current_t < TMAX:
        z_val = z_cur & 1
        if b < a:
            b_cur += 1
            b = state[b_mem + b_cur, st_idx]
            b += line_times[b_idx, 0, z_val ^ 1] * rand_gauss_dev(_seed ^ b_mem ^ z_val ^ 1, sd)
            thresh = line_times[b_idx, 1, z_val] * rand_gauss_dev(_seed ^ b_mem ^ z_val, sd)
            inputs ^= 2
            next_t = b
        else:
            a_cur += 1
            a = state[a_mem + a_cur, st_idx]
            a += line_times[a_idx, 0, z_val ^ 1] * rand_gauss_dev(_seed ^ a_mem ^ z_val ^ 1, sd)
            thresh = line_times[a_idx, 1, z_val] * rand_gauss_dev(_seed ^ a_mem ^ z_val, sd)
            inputs ^= 1
            next_t = a

        if (z_cur & 1) != ((lut >> inputs) & 1):
            # we generate a toggle in z_mem, if:
            #   ( it is the first toggle in z_mem OR
            #   following toggle is earlier OR
            #   pulse is wide enough ) AND enough space in z_mem.
            if z_cur == 0 or next_t < current_t or (current_t - previous_t) > thresh:
                if z_cur < (z_cap - 1):
                    state[z_mem + z_cur, st_idx] = current_t
                    previous_t = current_t
                    z_cur += 1
                else:
                    overflows += 1
                    previous_t = state[z_mem + z_cur - 1, st_idx]
                    z_cur -= 1
            else:
                z_cur -= 1
                if z_cur > 0:
                    previous_t = state[z_mem + z_cur - 1, st_idx]
                else:
                    previous_t = TMIN
        current_t = min(a, b)

    if overflows > 0:
        state[z_mem + z_cur, st_idx] = TMAX_OVL
    else:
        state[z_mem + z_cur, st_idx] = a if a > b else b  # propagate overflow flags by storing biggest TMAX from input


#!/usr/bin/env python
"""
Outputs the a time range's event counts, payloads, and schemas for each stream.

Usage:
  python scripts/summarize_streams.py --kronos-url http://your-kronos-instance:port
                                     [--csv myfile.csv]
                                     [--start 2003-09-25T10:49:41.5-03:00]
                                     [--end 2003-09-25T10:49:41.5-03:00]

If not supplied, --start defaults to 24 hours before --end, and --end defaults
to now.
"""
import argparse
import csv
import datetime
import numpy as np
import sys
import ujson
from dateutil.parser import parse
from dateutil.tz import tzlocal
from pykronos import KronosClient
from pykronos.common.time import datetime_to_kronos_time

output = """
Stream: %(stream)s

Event count:
  Total:      %(total_events)d
  Per day:    %(events_per_day)f
  Per second: %(events_per_sec)f

Payload (bytes):
  Total:  %(payload_total_bytes)d
  Avg:    %(payload_avg_bytes)f
  Median: %(payload_med_bytes)d
  95th:   %(payload_95_bytes)d
  99th:   %(payload_99_bytes)d

Schema:
  %(schema)s

"""

def indent(msg, indent):
  msg = msg.splitlines()
  for line in msg:
    sys.stdout.write('%s%s\n' % (' ' * indent, line))

def main(args):
  client = KronosClient(args.kronos_url)
  headers = [
    'stream',
    'total_events',
    'events_per_day',
    'events_per_sec',
    'payload_total_bytes',
    'payload_avg_bytes',
    'payload_med_bytes',
    'payload_95_bytes',
    'payload_99_bytes',
    'schema',
  ]
  if args.csv:
    csv_file = open(args.csv, 'w')
    writer = csv.DictWriter(csv_file, headers)
    writer.writeheader()
  else:
    print '-' * 79
  for stream in client.get_streams():
    total_events = 0
    payloads = []
    for event in client.get(stream, args.start, args.end):
      payloads.append(len(ujson.dumps(event)))
      total_events += 1
    if total_events == 0:
      indent('%s has no events' % stream, 2)
      print '-' * 79
      continue
    timeframe_sec = (args.end - args.start).total_seconds()
    schema = client.infer_schema(stream)['schema']
    context = dict(zip(headers, [
      stream,
      total_events,
      (float(total_events) / timeframe_sec) * 60 * 60 * 24,
      float(total_events) / timeframe_sec,
      np.sum(payloads),
      np.mean(payloads),
      np.median(payloads),
      np.percentile(payloads, 95),
      np.percentile(payloads, 99),
      schema,
    ]))
    if args.csv:
      writer.writerow(context)
    else:
      indent(output % context, 2)
      print '-' * 79

def process_args():
  parser = argparse.ArgumentParser()
  parser.add_argument('--kronos-url', required=True,
                      help='The Kronos server to retrieve data from')
  parser.add_argument('--csv', help='A CSV file for output')
  parser.add_argument('--start',
                      help=('When to start retrieving? (format: '
                            '2003-09-25T10:49:41.5-03:00)'))
  parser.add_argument('--end',
                      help=('When to end retrieving? (format: '
                            '2003-09-25T10:49:41.5-03:00)'))
  args = parser.parse_args()
  if args.end:
    args.end = parse(args.end)
  else:
    args.end = datetime.datetime.utcnow()
  if args.start:
    args.start = parse(args.start)
  else:
    args.start = args.end - datetime.timedelta(hours=24)
  if args.end < args.start:
    sys.stderr.write('ERROR: --start must be earlier than --end\n')
    exit(1)
  if (args.end - args.start).total_seconds() < 1:
    sys.stderr.write('ERROR: timeframe must be at least 1 second\n')
    exit(1)
  return args

if __name__ == '__main__':
  main(process_args())


# -*- coding: utf-8 -*-

import os
import gym
import numpy as np

import torch
import torch.nn as nn
import torch.optim as optim
from rllite import DDPG
from rllite.common import ReplayBuffer,NormalizedActions,GaussianExploration,QNet,PolicyNet,soft_update

from tensorboardX import SummaryWriter

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class TD3(DDPG):
    def __init__(
            self,
            env_name = 'Pendulum-v0',
            load_dir = './ckpt',
            log_dir = "./log",
            buffer_size = 1e6,
            seed = 1,
            max_episode_steps = None,
            noise_decay_steps = 1e5,
            batch_size = 64,
            discount = 0.99,
            train_freq = 100,
            policy_freq = 2,
            learning_starts = 500,
            tau = 0.005,
            save_eps_num = 100,
            external_env = None
            ):
        self.env_name = env_name
        self.load_dir = load_dir
        self.log_dir = log_dir
        self.seed = seed
        self.max_episode_steps = max_episode_steps
        self.buffer_size = buffer_size
        self.noise_decay_steps = noise_decay_steps
        self.batch_size = batch_size
        self.discount = discount
        self.policy_freq = policy_freq
        self.learning_starts = learning_starts
        self.tau = tau
        self.save_eps_num = save_eps_num

        torch.manual_seed(self.seed)
        np.random.seed(self.seed)
        self.writer = SummaryWriter(log_dir=self.log_dir)

        if external_env == None:
            env = gym.make(self.env_name)
        else:
            env = external_env
        if self.max_episode_steps != None:
            env._max_episode_steps = self.max_episode_steps
        else:
            self.max_episode_steps = env._max_episode_steps
        self.env = NormalizedActions(env)

        self.noise = GaussianExploration(self.env.action_space, decay_period=self.noise_decay_steps)

        state_dim  = self.env.observation_space.shape[0]
        action_dim = self.env.action_space.shape[0]
        hidden_dim = 256

        self.value_net1 = QNet(state_dim, action_dim, hidden_dim).to(device)
        self.value_net2 = QNet(state_dim, action_dim, hidden_dim).to(device)
        self.policy_net = PolicyNet(state_dim, action_dim, hidden_dim).to(device)

        self.target_value_net1 = QNet(state_dim, action_dim, hidden_dim).to(device)
        self.target_value_net2 = QNet(state_dim, action_dim, hidden_dim).to(device)
        self.target_policy_net = PolicyNet(state_dim, action_dim, hidden_dim).to(device)

        try:
            self.load(directory=self.load_dir, filename=self.env_name)
            print('Load model successfully !')
        except:
            print('WARNING: No model to load !')

        soft_update(self.value_net1, self.target_value_net1, soft_tau=1.0)
        soft_update(self.value_net2, self.target_value_net2, soft_tau=1.0)
        soft_update(self.policy_net, self.target_policy_net, soft_tau=1.0)

        self.value_criterion = nn.MSELoss()

        policy_lr = 3e-4
        value_lr  = 3e-4

        self.value_optimizer1 = optim.Adam(self.value_net1.parameters(), lr=value_lr)
        self.value_optimizer2 = optim.Adam(self.value_net2.parameters(), lr=value_lr)
        self.policy_optimizer = optim.Adam(self.policy_net.parameters(), lr=policy_lr)

        self.replay_buffer = ReplayBuffer(self.buffer_size)

        self.total_steps = 0
        self.episode_num = 0
        self.episode_timesteps = 0

    def save(self, directory, filename):
        if not os.path.exists(directory):
            os.makedirs(directory)

        torch.save(self.value_net1.state_dict(), '%s/%s_value_net1.pkl' % (directory, filename))
        torch.save(self.value_net2.state_dict(), '%s/%s_value_net2.pkl' % (directory, filename))
        torch.save(self.policy_net.state_dict(), '%s/%s_policy_net.pkl' % (directory, filename))

    def load(self, directory, filename):
        self.value_net1.load_state_dict(torch.load('%s/%s_value_net1.pkl' % (directory, filename)))
        self.value_net2.load_state_dict(torch.load('%s/%s_value_net2.pkl' % (directory, filename)))
        self.policy_net.load_state_dict(torch.load('%s/%s_policy_net.pkl' % (directory, filename)))

    def train_step(self,
           step,
           noise_std = 0.2,
           noise_clip=0.5
          ):

        state, action, reward, next_state, done = self.replay_buffer.sample(self.batch_size)

        state      = torch.FloatTensor(state).to(device)
        next_state = torch.FloatTensor(next_state).to(device)
        action     = torch.FloatTensor(action).to(device)
        reward     = torch.FloatTensor(reward).unsqueeze(1).to(device)
        done       = torch.FloatTensor(np.float32(done)).unsqueeze(1).to(device)

        next_action = self.target_policy_net(next_state)
        noise = torch.normal(torch.zeros(next_action.size()), noise_std).to(device)
        noise = torch.clamp(noise, -noise_clip, noise_clip)
        next_action += noise

        target_q_value1  = self.target_value_net1(next_state, next_action)
        target_q_value2  = self.target_value_net2(next_state, next_action)
        target_q_value   = torch.min(target_q_value1, target_q_value2)
        expected_q_value = reward + (1.0 - done) * self.discount * target_q_value

        q_value1 = self.value_net1(state, action)
        q_value2 = self.value_net2(state, action)

        value_loss1 = self.value_criterion(q_value1, expected_q_value.detach())
        value_loss2 = self.value_criterion(q_value2, expected_q_value.detach())

        self.value_optimizer1.zero_grad()
        value_loss1.backward()
        self.value_optimizer1.step()

        self.value_optimizer2.zero_grad()
        value_loss2.backward()
        self.value_optimizer2.step()

        if step % self.policy_freq == 0:
            policy_loss = self.value_net1(state, self.policy_net(state))
            policy_loss = -policy_loss.mean()

            self.policy_optimizer.zero_grad()
            policy_loss.backward()
            self.policy_optimizer.step()

            soft_update(self.value_net1, self.target_value_net1, soft_tau=self.tau)
            soft_update(self.value_net2, self.target_value_net2, soft_tau=self.tau)
            soft_update(self.policy_net, self.target_policy_net, soft_tau=self.tau)

    def learn(self, max_steps=1e7):
        while self.total_steps < max_steps:
            state = self.env.reset()
            self.episode_timesteps = 0
            episode_reward = 0

            for step in range(self.max_episode_steps):
                action = self.policy_net.get_action(state)
                action = self.noise.get_action(action, self.total_steps)
                next_state, reward, done, _ = self.env.step(action)

                self.replay_buffer.push(state, action, reward, next_state, done)

                state = next_state
                episode_reward += reward
                self.total_steps += 1
                self.episode_timesteps += 1

                if done or self.episode_timesteps == self.max_episode_steps:
                    if len(self.replay_buffer) > self.learning_starts:
                        for i in range(self.episode_timesteps):
                            self.train_step(i)

                    self.episode_num += 1
                    if self.episode_num > 0 and self.episode_num % self.save_eps_num == 0:
                        self.save(directory=self.load_dir, filename=self.env_name)

                    self.writer.add_scalar('episode_reward', episode_reward, self.episode_num)
                    break
        self.env.close()

if __name__ == '__main__':
    model = TD3()
    model.learn()