# Copyright (c) 2014 Evalf
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

"""
The matrix module defines an abstract :class:`Matrix` object and several
implementations.  Matrix objects support basic addition and subtraction
operations and provide a consistent insterface for solving linear systems.
Matrices can be converted into other forms suitable for external processing via
the ``export`` method.
"""

from . import numpy, numeric, warnings, cache, types, config, util
import abc, sys, ctypes, treelog as log

class MatrixError(Exception): pass

class Backend(metaclass=abc.ABCMeta):
  'backend base class'

  def __enter__(self):
    if hasattr(self, '_old_backend'):
      raise RuntimeError('This context manager is not reentrant.')
    global _current_backend
    self._old_backend = _current_backend
    _current_backend = self
    return self

  def __exit__(self, etype, value, tb):
    if not hasattr(self, '_old_backend'):
      raise RuntimeError('This context manager is not yet entered.')
    global _current_backend
    _current_backend = self._old_backend
    del self._old_backend

  @abc.abstractmethod
  def assemble(self, data, index, shape):
    '''Assemble a (sparse) tensor based on index-value pairs.

    .. Note:: This function is abstract.
    '''

class Matrix(metaclass=types.CacheMeta):
  'matrix base class'

  def __init__(self, shape):
    assert len(shape) == 2
    self.shape = shape

  @abc.abstractmethod
  def __add__(self, other):
    'add two matrices'

  @abc.abstractmethod
  def __mul__(self, other):
    'multiply matrix with a scalar'

  @abc.abstractmethod
  def __neg__(self, other):
    'negate matrix'

  def __sub__(self, other):
    return self.__add__(-other)

  def __rmul__(self, other):
    return self.__mul__(other)

  def __truediv__(self, other):
    return self.__mul__(1/other)

  @property
  @abc.abstractmethod
  def T(self):
    'transpose matrix'

  @property
  def size(self):
    return numpy.prod(self.shape)

  def rowsupp(self, tol=0):
    'return row indices with nonzero/non-small entries'

    data, (row, col) = self.export('coo')
    supp = numpy.zeros(self.shape[0], dtype=bool)
    supp[row[abs(data) > tol]] = True
    return supp

  @abc.abstractmethod
  def solve(self, rhs=None, *, lhs0=None, constrain=None, rconstrain=None, **solverargs):
    '''Solve system given right hand side vector and/or constraints.

    Args
    ----
    rhs : :class:`float` vector or :any:`None`
        Right hand side vector. `None` implies all zeros.
    lhs0 : class:`float` vector or :any:`None`
        Initial values. `None` implies all zeros.
    constrain : :class:`float` or :class:`bool` array, or :any:`None`
        Column constraints. For float values, a number signifies a constraint,
        NaN signifies a free dof. For boolean, a True value signifies a
        constraint to the value in `lhs0`, a False value signifies a free dof.
        `None` implies no constraints.
    rconstrain : :class:`bool` array or :any:`None`
        Row constrains. A True value signifies a constrains, a False value a free
        dof. `None` implies that the constraints follow those defined in
        `constrain` (by implication the matrix must be square).

    Returns
    -------
    :class:`numpy.ndarray`
        Left hand side vector.
    '''

  @abc.abstractmethod
  def submatrix(self, rows, cols):
    '''Create submatrix from selected rows, columns.

    Args
    ----
    rows : :class:`bool`/:class:`int` array selecting rows for keeping
    cols : :class:`bool`/:class:`int` array selecting columns for keeping

    Returns
    -------
    :class:`Matrix`
        Matrix instance of reduced dimensions
    '''

  def export(self, form):
    '''Export matrix data to any of supported forms.

    Args
    ----
    form : :class:`str`
      - "dense" : return matrix as a single dense array
      - "csr" : return matrix as 3-tuple of (data, indices, indptr)
      - "coo" : return matrix as 2-tuple of (data, (row, col))
    '''

    raise NotImplementedError('cannot export {} to {!r}'.format(self.__class__.__name__, form))

  def __repr__(self):
    return '{}<{}x{}>'.format(type(self).__qualname__, *self.shape)

def preparesolvearguments(wrapped):
  '''Make rhs optional, add lhs0, constrain, rconstrain arguments.

  See Matrix.solve.'''

  def solve(self, rhs=None, *, lhs0=None, constrain=None, rconstrain=None, **solverargs):
    nrows, ncols = self.shape
    if lhs0 is None:
      x = numpy.zeros(ncols)
    else:
      x = numpy.array(lhs0, dtype=float)
      assert x.shape == (ncols,)
    if constrain is None:
      J = numpy.ones(ncols, dtype=bool)
    else:
      assert constrain.shape == (ncols,)
      if constrain.dtype == bool:
        J = ~constrain
      else:
        J = numpy.isnan(constrain)
        x[~J] = constrain[~J]
    if rconstrain is None:
      assert nrows == ncols
      I = J
    else:
      assert rconstrain.shape == (nrows,) and constrain.dtype == bool
      I = ~rconstrain
    assert I.sum() == J.sum(), 'constrained matrix is not square: {}x{}'.format(I.sum(), J.sum())
    if rhs is None:
      rhs = 0.
    b = (rhs - self.matvec(x))[J]
    if b.any():
      x[J] += wrapped(self if I.all() and J.all() else self.submatrix(I, J), b, **solverargs)
      if not numpy.isfinite(x).all():
        raise MatrixError('solver returned non-finite left hand side')
      log.info('solver returned with residual {:.0e}'.format(numpy.linalg.norm((rhs - self.matvec(x))[J])))
    else:
      log.info('skipping solver because initial vector is exact')
    return x
  return log.withcontext(solve)

## NUMPY BACKEND

class Numpy(Backend):
  '''matrix backend based on numpy array'''

  def assemble(self, data, index, shape):
    array = numeric.accumulate(data, index, shape)
    return NumpyMatrix(array) if len(shape) == 2 else array

class NumpyMatrix(Matrix):
  '''matrix based on numpy array'''

  def __init__(self, core):
    assert numeric.isarray(core)
    self.core = core
    super().__init__(core.shape)

  def __add__(self, other):
    if not isinstance(other, NumpyMatrix) or self.shape != other.shape:
      return NotImplemented
    return NumpyMatrix(self.core + other.core)

  def __mul__(self, other):
    if not numeric.isnumber(other):
      return NotImplemented
    return NumpyMatrix(self.core * other)

  def __neg__(self):
    return NumpyMatrix(-self.core)

  @property
  def T(self):
    return NumpyMatrix(self.core.T)

  def matvec(self, vec):
    return numpy.dot(self.core, vec)

  def export(self, form):
    if form == 'dense':
      return self.core
    if form == 'coo':
      ij = self.core.nonzero()
      return self.core[ij], ij
    if form == 'csr':
      rows, cols = self.core.nonzero()
      return self.core[rows, cols], cols, rows.searchsorted(numpy.arange(self.shape[0]+1))
    raise NotImplementedError('cannot export NumpyMatrix to {!r}'.format(form))

  def rowsupp(self, tol=0):
    return numpy.greater(abs(self.core), tol).any(axis=1)

  @preparesolvearguments
  def solve(self, rhs):
    try:
      return numpy.linalg.solve(self.core, rhs)
    except numpy.linalg.LinAlgError as e:
      raise MatrixError(e) from e

  def submatrix(self, rows, cols):
    return NumpyMatrix(self.core[numpy.ix_(rows, cols)])

## SCIPY BACKEND

try:
  import scipy.sparse.linalg
except ImportError:
  pass
else:

  class Scipy(Backend):
    '''matrix backend based on scipy's sparse matrices'''

    def assemble(self, data, index, shape):
      if len(shape) < 2:
        return numeric.accumulate(data, index, shape)
      if len(shape) == 2:
        csr = scipy.sparse.csr_matrix((data, index), shape)
        return ScipyMatrix(csr)
      raise MatrixError('{}d data not supported by scipy backend'.format(len(shape)))

  class ScipyMatrix(Matrix):
    '''matrix based on any of scipy's sparse matrices'''

    def __init__(self, core):
      self.core = core
      super().__init__(core.shape)

    def __add__(self, other):
      if not isinstance(other, ScipyMatrix) or self.shape != other.shape:
        return NotImplemented
      return ScipyMatrix(self.core + other.core)

    def __sub__(self, other):
      if not isinstance(other, ScipyMatrix) or self.shape != other.shape:
        return NotImplemented
      return ScipyMatrix(self.core - other.core)

    def __mul__(self, other):
      if not numeric.isnumber(other):
        return NotImplemented
      return ScipyMatrix(self.core * other)

    def __neg__(self):
      return ScipyMatrix(-self.core)

    def matvec(self, vec):
      return self.core.dot(vec)

    def export(self, form):
      if form == 'dense':
        return self.core.toarray()
      if form == 'csr':
        csr = self.core.tocsr()
        return csr.data, csr.indices, csr.indptr
      if form == 'coo':
        coo = self.core.tocoo()
        return coo.data, (coo.row, coo.col)
      raise NotImplementedError('cannot export NumpyMatrix to {!r}'.format(form))

    @property
    def T(self):
      return ScipyMatrix(self.core.transpose())

    @preparesolvearguments
    def solve(self, rhs, atol=0, solver='spsolve', callback=None, precon=None, **solverargs):
      if solver == 'spsolve':
        log.info('solving system using sparse direct solver')
        return scipy.sparse.linalg.spsolve(self.core, rhs)
      assert atol, 'tolerance must be specified for iterative solver'
      rhsnorm = numpy.linalg.norm(rhs)
      if rhsnorm <= atol:
        return numpy.zeros(self.shape[1])
      log.info('solving system using {} iterative solver'.format(solver))
      solverfun = getattr(scipy.sparse.linalg, solver)
      myrhs = rhs / rhsnorm # normalize right hand side vector for best control over scipy's stopping criterion
      mytol = atol / rhsnorm
      niter = numpy.array(0)
      def mycallback(arg):
        niter[...] += 1
        # some solvers provide the residual, others the left hand side vector
        res = numpy.linalg.norm(myrhs - self.matvec(arg)) if numpy.ndim(arg) == 1 else float(arg)
        if callback:
          callback(res)
        with log.context('residual {:.2e} ({:.0f}%)'.format(res, 100. * numpy.log10(res) / numpy.log10(mytol) if res > 0 else 0)):
          pass
      M = self.getprecon(precon) if isinstance(precon, str) else precon(self.core) if callable(precon) else precon
      mylhs, status = solverfun(self.core, myrhs, M=M, tol=mytol, callback=mycallback, **solverargs)
      if status != 0:
        raise MatrixError('{} solver failed with status {}'.format(solver, status))
      log.info('solver converged in {} iterations'.format(niter))
      return mylhs * rhsnorm

    def getprecon(self, name):
      name = name.lower()
      assert self.shape[0] == self.shape[1], 'constrained matrix must be square'
      log.info('building {} preconditioner'.format(name))
      if name == 'splu':
        try:
          precon = scipy.sparse.linalg.splu(self.core.tocsc()).solve
        except RuntimeError as e:
          raise MatrixError(e) from e
      elif name == 'spilu':
        try:
          precon = scipy.sparse.linalg.spilu(self.core.tocsc(), drop_tol=1e-5, fill_factor=None, drop_rule=None, permc_spec=None, diag_pivot_thresh=None, relax=None, panel_size=None, options=None).solve
        except RuntimeError as e:
          raise MatrixError(e) from e
      elif name == 'diag':
        diag = self.core.diagonal()
        if not diag.all():
          raise MatrixError("building 'diag' preconditioner: diagonal has zero entries")
        precon = numpy.reciprocal(diag).__mul__
      else:
        raise MatrixError('invalid preconditioner {!r}'.format(name))
      return scipy.sparse.linalg.LinearOperator(self.shape, precon, dtype=float)

    def submatrix(self, rows, cols):
      return ScipyMatrix(self.core[rows,:][:,cols])

## INTEL MKL BACKEND

libmkl = util.loadlib(linux='libmkl_rt.so', darwin='libmkl_rt.dylib', win32='mkl_rt.dll')
if libmkl is not None:

  # typedefs
  c_int = types.c_array[numpy.int32]
  c_long = types.c_array[numpy.int64]
  c_double = types.c_array[numpy.float64]

  libtbb = util.loadlib(linux='libtbb.so.2', darwin='libtbb.dylib', win32='tbb.dll')

  class MKL(Backend):
    '''matrix backend based on Intel's Math Kernel Library'''

    def __enter__(self):
      super().__enter__()
      usethreads = config.nprocs > 1
      libmkl.mkl_set_threading_layer(c_long(4 if usethreads else 1)) # 1:SEQUENTIAL, 4:TBB
      if usethreads and libtbb:
        self.tbbhandle = ctypes.c_void_p()
        libtbb._ZN3tbb19task_scheduler_init10initializeEim(ctypes.byref(self.tbbhandle), ctypes.c_int(config.nprocs), ctypes.c_int(2))
      else:
        self.tbbhandle = None
      return self

    def __exit__(self, etype, value, tb):
      if self.tbbhandle:
        libtbb._ZN3tbb19task_scheduler_init9terminateEv(ctypes.byref(self.tbbhandle))
      super().__exit__(etype, value, tb)

    @staticmethod
    def assemble(data, index, shape):
      if len(shape) < 2:
        return numeric.accumulate(data, index, shape)
      if len(shape) == 2:
        return MKLMatrix(data, index, shape)
      raise MatrixError('{}d data not supported by scipy backend'.format(len(shape)))

  class Pardiso:
    '''simple wrapper for libmkl.pardiso

    https://software.intel.com/en-us/mkl-developer-reference-c-pardiso
    '''

    _pardiso = libmkl.pardiso
    _errorcodes = {
      -1: 'input inconsistent',
      -2: 'not enough memory',
      -3: 'reordering problem',
      -4: 'zero pivot, numerical factorization or iterative refinement problem',
      -5: 'unclassified (internal) error',
      -6: 'reordering failed (matrix types 11 and 13 only)',
      -7: 'diagonal matrix is singular',
      -8: '32-bit integer overflow problem',
      -9: 'not enough memory for OOC',
     -10: 'error opening OOC files',
     -11: 'read/write error with OOC files',
     -12: 'pardiso_64 called from 32-bit library',
    }

    def __init__(self):
      self.pt = numpy.zeros(64, numpy.int64) # handle to data structure

    @types.apply_annotations
    def __call__(self, *, phase:c_int, iparm:c_int, maxfct:c_int=1, mnum:c_int=1, mtype:c_int=0, n:c_int=0, a:c_double=None, ia:c_int=None, ja:c_int=None, perm:c_int=None, nrhs:c_int=0, msglvl:c_int=0, b:c_double=None, x:c_double=None):
      error = ctypes.c_int32(1)
      self._pardiso(self.pt.ctypes, maxfct, mnum, mtype, phase, n, a, ia, ja, perm, nrhs, iparm, msglvl, b, x, ctypes.byref(error))
      if error.value:
        raise MatrixError(self._errorcodes.get(error.value, 'unknown error {}'.format(error.value)))

    def __del__(self):
      if self.pt.any(): # release all internal memory for all matrices
        self(phase=-1, iparm=numpy.zeros(64, dtype=numpy.int32))
        assert not self.pt.any(), 'it appears that Pardiso failed to release its internal memory'

  class MKLMatrix(Matrix):
    '''matrix implementation based on sorted coo data'''

    __cache__ = 'indptr',

    _factors = False

    def __init__(self, data, index, shape):
      assert index.shape == (2, len(data))
      if len(data):
        # sort rows, columns
        reorder = numpy.lexsort(index[::-1])
        index = index[:,reorder]
        data = data[reorder]
        # sum duplicate entries
        keep = numpy.empty(len(reorder), dtype=bool)
        keep[0] = True
        numpy.not_equal(index[:,1:], index[:,:-1]).any(axis=0, out=keep[1:])
        if not keep.all():
          index = index[:,keep]
          data = numeric.accumulate(data, [keep.cumsum()-1], [index.shape[1]])
        if not data.all():
          nz = data.astype(bool)
          data = data[nz]
          index = index[:,nz]
      self.data = numpy.ascontiguousarray(data, dtype=numpy.float64)
      self.index = numpy.ascontiguousarray(index, dtype=numpy.int32)
      super().__init__(shape)

    @property
    def indptr(self):
      return self.index[0].searchsorted(numpy.arange(self.shape[0]+1)).astype(numpy.int32, copy=False)

    def __add__(self, other):
      if not isinstance(other, MKLMatrix) or self.shape != other.shape:
        return NotImplemented
      return MKLMatrix(numpy.concatenate([self.data, other.data]), numpy.concatenate([self.index, other.index], axis=1), self.shape)

    def __sub__(self, other):
      if not isinstance(other, MKLMatrix) or self.shape != other.shape:
        return NotImplemented
      return MKLMatrix(numpy.concatenate([self.data, -other.data]), numpy.concatenate([self.index, other.index], axis=1), self.shape)

    def __mul__(self, other):
      if not numeric.isnumber(other):
        return NotImplemented
      return MKLMatrix(self.data * other, self.index, self.shape)

    def __neg__(self):
      return MKLMatrix(-self.data, self.index, self.shape)

    @property
    def T(self):
      return MKLMatrix(self.data, self.index[::-1], self.shape[::-1])

    def matvec(self, vec):
      rows, cols = self.index
      return numeric.accumulate(self.data * vec[cols], [rows], [self.shape[0]])

    def export(self, form):
      if form == 'dense':
        return numeric.accumulate(self.data, self.index, self.shape)
      if form == 'csr':
        return self.data, self.index[1], self.indptr
      if form == 'coo':
        return self.data, self.index
      raise NotImplementedError('cannot export MKLMatrix to {!r}'.format(form))

    def submatrix(self, rows, cols):
      I, J = self.index
      keep = numpy.logical_and(rows[I], cols[J])
      csI = rows.cumsum()
      csJ = cols.cumsum()
      return MKLMatrix(self.data[keep], numpy.array([csI[I[keep]]-1, csJ[J[keep]]-1]), shape=(csI[-1], csJ[-1]))

    @preparesolvearguments
    def solve(self, rhs):
      log.info('solving {0}x{0} system using MKL Pardiso'.format(self.shape[0]))
      if self._factors:
        log.info('reusing existing factorization')
        pardiso, iparm, mtype = self._factors
        phase = 33 # solve, iterative refinement
      else:
        pardiso = Pardiso()
        iparm = numpy.zeros(64, dtype=numpy.int32) # https://software.intel.com/en-us/mkl-developer-reference-c-pardiso-iparm-parameter
        iparm[0] = 1 # supply all values in components iparm[1:64]
        iparm[1] = 2 # fill-in reducing ordering for the input matrix: nested dissection algorithm from the METIS package
        iparm[9] = 13 # pivoting perturbation threshold 1e-13 (default for nonsymmetric)
        iparm[10] = 1 # enable scaling vectors (default for nonsymmetric)
        iparm[12] = 1 # enable improved accuracy using (non-) symmetric weighted matching (default for nonsymmetric)
        iparm[34] = 1 # zero base indexing
        mtype = 11 # real and nonsymmetric
        phase = 13 # analysis, numerical factorization, solve, iterative refinement
        self._factors = pardiso, iparm, mtype
      lhs = numpy.empty(self.shape[1], dtype=numpy.float64)
      pardiso(phase=phase, mtype=mtype, iparm=iparm, n=self.shape[0], nrhs=1, b=rhs, x=lhs, a=self.data, ia=self.indptr, ja=self.index[1])
      return lhs

## MODULE METHODS

_current_backend = Numpy()

def backend(names):
  for name in names.lower().split(','):
    for cls in Backend.__subclasses__():
      if cls.__name__.lower() == name:
        return cls()
  raise RuntimeError('matrix backend {!r} is not available'.format(names))

def assemble(data, index, shape):
  return _current_backend.assemble(data, index, shape)

def diag(d):
  assert d.ndim == 1
  return assemble(d, index=numpy.arange(len(d))[numpy.newaxis].repeat(2, axis=0), shape=d.shape*2)

def eye(n):
  return diag(numpy.ones(n))

# vim:sw=2:sts=2:et


r"""Classes for MRI excitation simulations
"""
import copy
from typing import Optional
import inspect

import numpy as np
from scipy import interpolate
import torch
from torch import tensor, Tensor

from mrphy import γH, dt0, gmax0, smax0, rfmax0, T1G, T2G, π
from mrphy import utils, beffective, sims

__all__ = ['Pulse', 'SpinArray', 'SpinCube', 'Examples']

class Pulse(object):
    r"""Pulse object of RF and GR

    Usage:
        ``pulse = Pulse(rf, gr, *, dt, gmax, smax, rfmax, desc, device,``\
        `` dtype)``

    Inputs:
        - ``rf``: `(N,xy, nT,(nCoils))` "Gauss", ``xy`` for separating real \
          and imag part.
        - ``gr``: `(N,xyz,nT)`, "Gauss/cm"
        - ``dt``: `()` ⊻ `(N ⊻ 1,)`, "Sec", dwell time.
        - ``gmax``: `()` ⊻ `(N ⊻ 1, xyz ⊻ 1)`, "Gauss/cm", max \|gradient\|.
        - ``smax``: `()` ⊻ `(N ⊻ 1, xyz ⊻ 1)`, "Gauss/cm/Sec", max \
          \|slew rate\|.
        - ``rfmax``: `()` ⊻ `(N ⊻ 1,(nCoils))`, "Gauss", max \|RF\|.
        - ``desc``: str, an description of the pulse to be constructed.
        - ``device``: torch.device.
        - ``dtype``: torch.dtype.

    Properties:
        - ``device``
        - ``dtype``
        - ``is_cuda``
        - ``shape``: ``(N,1,nT)``
        - ``gmax``: `(N ⊻ 1, xyz)`, "Gauss/cm", max \|gradient\|.
        - ``smax``: `(N ⊻ 1, xyz)`, "Gauss/cm/Sec", max \|slew rate\|.
        - ``rfmax``: `(N ⊻ 1,(nCoils))`, "Gauss", max \|RF\|.
        - ``rf``: `(N,xy, nT,(nCoils))`, "Gauss", ``xy`` for separating real \
          and imag part.
        - ``gr``: `(N,xyz,nT)`, "Gauss/cm"
        - ``dt``: `(N ⊻ 1,)`, "Sec", dwell time.
        - ``desc``: str, an description of the pulse to be constructed.
    """

    _readonly = ('device', 'dtype', 'is_cuda', 'shape')
    _limits = ('gmax', 'smax', 'rfmax')
    __slots__ = set(_readonly + _limits + ('rf', 'gr', 'dt', 'desc'))

    def __init__(
        self,
        rf: Optional[Tensor] = None, gr: Optional[Tensor] = None, *,
        dt: Tensor = dt0,
        gmax: Tensor = gmax0, smax: Tensor = smax0, rfmax: Tensor = rfmax0,
        desc: str = "generic pulse",
        device: torch.device = torch.device('cpu'),
        dtype: torch.dtype = torch.float32
    ):
        assert(isinstance(device, torch.device) and
               isinstance(dtype, torch.dtype))

        # Defaults
        rf_miss, gr_miss = rf is None, gr is None
        assert (not(rf_miss and gr_miss)), "Missing both `rf` and `gr` inputs"

        object.__setattr__(self, 'device', device)
        object.__setattr__(self, 'dtype', dtype)
        object.__setattr__(self, 'is_cuda', self.device.type == 'cuda')

        kw = {'device': self.device, 'dtype': self.dtype}

        if rf_miss:
            N, nT = gr.shape[0], gr.shape[2]
            rf = torch.zeros((N, 2, nT), **kw)
        else:
            N, nT = rf.shape[0], rf.shape[2]
            if gr_miss:
                gr = torch.zeros((N, 3, nT), **kw)
            else:
                assert (N == gr.shape[0] and nT == gr.shape[2])
        object.__setattr__(self, 'shape', torch.Size((N, 1, nT)))

        self.rf, self.gr = rf.to(**kw), gr.to(**kw)
        self.dt, self.gmax, self.smax, self.rfmax = dt, gmax, smax, rfmax
        self.desc = desc
        return

    def __setattr__(self, k, v):
        if 'deepcopy' in (_.function for _ in inspect.stack()):
            # Hack, this enables `deepcopy()` w/o overriding `__deepcopy__()`.
            # Generator is faster than list comprehension.
            object.__setattr__(self, k, v)
            return

        if k in self._readonly:
            raise AttributeError(f"'Pulse' object attribute '{k}'"
                                 " is read-only")

        if k != 'desc':
            kw = {'device': self.device, 'dtype': self.dtype}
            v = (v.to(**kw) if isinstance(v, Tensor) else tensor(v, **kw))

        if k in ('rf', 'gr'):
            assert(v.shape[0] == self.shape[0] and v.shape[2] == self.shape[2])
        elif (k in ('gmax', 'smax')):  # -> (N ⊻ 1, xyz)
            v = v.expand((1 if v.ndim == 0 else v.shape[0], self.gr.shape[1]))
        elif k == 'rfmax':  # -> (N ⊻ 1, (nCoils))
            if v.ndim == 0:
                v = v[None]
            elif v.ndim == 2 and v.shape[1] == 1:
                v = v[:, 0]
        elif k == 'dt':
            if v.ndim == 0:
                v = v[None]
            assert(v.ndim == 1)

        object.__setattr__(self, k, v)
        return

    def asdict(self, *, toNumpy: bool = True) -> dict:
        r"""Convert mrphy.mobjs.Pulse object to dict

        Usage:
            ``d = pulse.asdict(*, toNumpy)``

        Inputs:
            - ``toNumpy``: [T/f], convert Tensor to Numpy arrays.
        Outputs:
            - ``d``: dict, dictionary with detached data identical to the \
              object.
        """
        _ = ('rf', 'gr', 'dt', 'gmax', 'smax', 'rfmax')
        fn_np = ((lambda x: x.detach().cpu().numpy()) if toNumpy else
                 (lambda x: x.detach()))

        d = {k: fn_np(getattr(self, k)) for k in _}
        d.update({k: getattr(self, k) for k in ('desc', 'device', 'dtype')})

        return d

    def beff(
            self, loc: Tensor, *,
            Δf: Optional[Tensor] = None, b1Map: Optional[Tensor] = None,
            γ: Tensor = γH
    ) -> Tensor:
        r"""Compute B-effective of provided location from the pulse

        Usage:
            ``beff = pulse.beff(loc, *, Δf, b1Map, γ)``
        Inputs:
            - ``loc``: `(N,*Nd,xyz)`, "cm", locations.
        Optionals:
            - ``Δf``: `(N,*Nd,)`, "Hz", off-resonance.
            - ``b1Map``: `(N,*Nd,xy,(nCoils))`, a.u., transmit sensitivity.
            - ``γ``: `(N,*Nd)`, "Hz/Gauss", gyro-ratio
        Outputs:
            - ``beff``: `(N,*Nd,xyz,nT)`
        """
        device = self.device
        loc = loc.to(device=device)
        fn = lambda x: None if x is None else x.to(device=device)  # noqa: E731
        Δf, b1Map, γ = (fn(x) for x in (Δf, b1Map, γ))

        return beffective.rfgr2beff(self.rf, self.gr, loc,
                                    Δf=Δf, b1Map=b1Map, γ=γ)

    def interpT(self, dt: Tensor, *, kind: str = 'linear') -> 'Pulse':
        r""" Interpolate pulse of `dt` by `kind`.

        Usage:
            ``new_pulse = pulse.interpT(dt, *, kind)``
        Inputs:
            - ``dt``: `(1,)`, "Sec", new simulation dwell time.
            - ``kind``: str, passed to scipy.interpolate.interp1d.
        Outputs:
            - ``new_pulse``: mrphy.mobjs.Pulse object.

        .. note::
            This method requires both `dt` and `self.dt` to be unique/global,
            i.e., of shape ``(1,)``, which ensures pulse length to be the same
            within a batch after interpolation.
        """
        assert(self.dt.numel() == dt.numel() == 1)

        dt_o_np, dt_n_np = self.dt.item(), dt.item()
        if dt_o_np == dt_n_np:
            return copy.deepcopy(self)

        axis = 2  # Along temporal dimension
        dkw = {'device': self.device, 'dtype': self.dtype}
        kw = {'axis': axis, 'kind': kind, 'copy': False, 'assume_sorted': True}

        f_np = lambda x: x.detach().cpu().numpy()  # noqa: E731
        f_0 = lambda x: np.dstack((np.zeros_like(x[:, :, [0]]),  # noqa: E731
                                   x))

        # convert to np array, then prepend 0's.
        rf_np, gr_np = f_0(f_np(self.rf)), f_0(f_np(self.gr))

        nT = rf_np.shape[axis]

        t_o = np.arange(0, nT)*dt_o_np  # (nT,)
        t_n = np.arange(1, t_o[-1]//dt_n_np + 1)*dt_n_np

        f_rf = interpolate.interp1d(t_o, rf_np, **kw)
        f_gr = interpolate.interp1d(t_o, gr_np, **kw)

        rf_n, gr_n = tensor(f_rf(t_n), **dkw), tensor(f_gr(t_n), **dkw)

        desc = f"{self.desc} + interpT\'ed: dt = {dt_n_np}"
        return Pulse(rf_n, gr_n, dt=dt, desc=desc, **dkw)

    def to(
        self, *,
        device: torch.device = torch.device('cpu'),
        dtype: torch.dtype = torch.float32
    ) -> 'Pulse':
        r"""Duplicate the object to the prescribed device with dtype

        Usage:
            ``new_pulse = pulse.to(*, device, dtype)``
        Inputs:
            - ``device``: torch.device
            - ``dtype``: torch.dtype
        Outputs:
            - ``new_pulse``: mrphy.mobjs.Pulse object.
        """
        if self.device == device and self.dtype == dtype:
            return self
        return Pulse(self.rf, self.gr, dt=self.dt, desc=self.desc,
                     device=device, dtype=dtype)

class SpinArray(object):
    r"""mrphy.mobjs.SpinArray object

    Usage:
        ``spinarray = SpinArray(shape, mask, *, T1_, T2_, γ_, M_, device,``\
        `` dtype)``
        ``spinarray = SpinArray(shape, mask, *, T1, T2, γ, M, device, dtype)``
    Inputs:
        - ``shape``: tuple, e.g., ``(N, nx, ny, nz)``.
    Optionals:
        - ``mask``: `(1, *Nd)`, where does compact attributes locate in `Nd`.
        - ``T1`` ⊻ ``T1_``: `(N, *Nd ⊻ nM)`, "Sec", T1 relaxation coeff.
        - ``T2`` ⊻ ``T2_``: `(N, *Nd ⊻ nM)`, "Sec", T2 relaxation coeff.
        - ``γ`` ⊻ ``γ_``: `(N, *Nd ⊻ nM)`,  "Hz/Gauss", gyro ratio.
        - ``M`` ⊻ ``M_``: `(N, *Nd ⊻ nM, xyz)`, spins, equilibrium ``[0 0 1]``.
        - ``device``: torch.device.
        - ``dtype``: torch.dtype

    Properties:
        - ``shape``: `(N, *Nd)`.
        - ``mask``: `(1, *Nd)`.
        - ``device``.
        - ``dtype``.
        - ``ndim``: ``len(shape)``
        - ``nM``: ``nM = torch.count_nonzero(mask).item()``.
        - ``T1_``: `(N, nM)`, "Sec", T1 relaxation coeff.
        - ``T2_``: `(N, nM)`, "Sec", T2 relaxation coeff.
        - ``γ_``: `(N, nM)`, "Hz/Gauss", gyro ratio.
        - ``M_``: `(N, nM, xyz)`, spins, equilibrium [0 0 1]

    .. warning::
        - Do NOT modify the ``mask`` of an object, e.g., \
          ``spinarray.mask[0] = True``.
        - Do NOT proceed indexed/masked assignments over any non-compact \
          attribute, e.g., ``spinarray.T1[0] = T1G`` or \
          ``spinarray.T1[mask] = T1G``.
          The underlying compact attributes will **NOT** be updated, since \
          they do not share memory.
          The only exception is when ``torch.all(mask == True)`` and the \
          underlying compact is **contiguous**, where the non-compact is just \
          a ``view((N, *Nd, ...))``.
          Checkout :func:`~mrphy.mobjs.SpinArray.crds_` and \
          :func:`~mrphy.mobjs.SpinArray.mask_` for indexed/masked access to \
          compacts.

    .. tip::
        - ``mask`` is GLOBAL for a batch, in other words, one cannot specify \
          distinct masks w/in a batch. \
          This design is to reduce storage/computations in, e.g., \
          ``applypulse`` (``blochsim``), avoiding extra allocations. \
          For DNN applications where an in-batch variation of ``mask`` may \
          seemingly be of interest, having ``torch.all(mask == True)`` and \
          postponing the variations to eventual losses evaluation can be a \
          better design, which allows reuse of ``M_``, etc., avoiding \
          repetitive allocations.
    """

    _readonly = ('shape', 'mask', 'device', 'dtype', 'is_cuda', 'ndim', 'nM')
    _compact = ('T1_', 'T2_', 'γ_', 'M_')
    __slots__ = set(_readonly + _compact)

    def __init__(
        self, shape: tuple, mask: Optional[Tensor] = None, *,
        T1: Optional[Tensor] = None, T1_: Optional[Tensor] = None,
        T2: Optional[Tensor] = None, T2_: Optional[Tensor] = None,
        γ: Optional[Tensor] = None,  γ_: Optional[Tensor] = None,
        M: Optional[Tensor] = None,  M_: Optional[Tensor] = None,
        device: torch.device = torch.device('cpu'),
        dtype: torch.dtype = torch.float32
    ):

        mask = (torch.ones((1,)+shape[1:], dtype=torch.bool, device=device)
                if mask is None else mask.to(device=device))

        assert(isinstance(device, torch.device) and
               isinstance(dtype, torch.dtype) and
               mask.dtype == torch.bool and
               mask.shape == (1,)+shape[1:])

        object.__setattr__(self, 'shape', shape)
        object.__setattr__(self, 'mask', mask)
        object.__setattr__(self, 'ndim', len(shape))
        object.__setattr__(self, 'nM', torch.count_nonzero(mask).item())
        object.__setattr__(self, 'device', device)
        object.__setattr__(self, 'dtype', dtype)
        object.__setattr__(self, 'is_cuda', self.device.type == 'cuda')

        assert((T1 is None) or (T1_ is None))
        if T1 is None:
            self.T1_ = (T1G if T1_ is None else T1_)
        else:
            self.T1 = T1

        assert((T2 is None) or (T2_ is None))
        if T2 is None:
            self.T2_ = (T2G if T2_ is None else T2_)
        else:
            self.T2 = T2

        assert((γ is None) or (γ_ is None))
        if γ is None:
            self.γ_ = (γH if γ_ is None else γ_)
        else:
            self.γ = γ

        assert((M is None) or (M_ is None))
        if M is None:
            self.M_ = (tensor([0., 0., 1.]) if M_ is None else M_)
        else:
            self.M = M

        return

    def __getattr__(self, k):  # provoked only when `__getattribute__` failed
        if k+'_' not in self._compact:
            raise AttributeError(f"'SpinArray' has no attribute '{k}'")

        v_ = getattr(self, k+'_')
        return (self.embed(v_) if self.nM != np.prod(self.shape[1:]) else
                v_.reshape(self.shape+v_.shape[2:]))  # ``mask`` is all True

    def __setattr__(self, k_, v_):
        if 'deepcopy' in (_.function for _ in inspect.stack()):
            # Hack, this enables `deepcopy()` w/o overriding `__deepcopy__()`.
            # Generator is faster than list comprehension.
            object.__setattr__(self, k_, v_)
            return

        if k_ in self._readonly:
            raise AttributeError(f"'SpinArray' object attribute '{k_}'"
                                 " is read-only")

        # Transfer ``v_`` to ``kw`` before ``extract`
        kw = {'device': self.device, 'dtype': self.dtype}
        v_ = (v_.to(**kw) if isinstance(v_, Tensor) else tensor(v_, **kw))

        shape = self.shape
        if k_+'_' in self._compact:  # enable non-compact assignment
            k_ = k_+'_'
            v_ = self.extract(v_.expand(shape+(3,) if k_ == 'M_' else shape))

        # `tensor.expand(size)` needs `tensor.shape` broadcastable with `size`
        if k_ == 'M_':
            if v_.shape != shape[:1]+(self.nM, 3):  # (N, nM, xyz)
                v_ = v_.expand(shape[:1]+(self.nM, 3)).clone()
        elif k_ in self._compact:  # (T1_, T2_, γ_)
            v_ = v_.expand((self.shape[0], self.nM))  # (N, nM)

        object.__setattr__(self, k_, v_)
        return

    def applypulse(
        self, pulse: Pulse, *,
        doEmbed: bool = False, doRelax: bool = True, doUpdate: bool = False,
        loc: Optional[Tensor] = None, loc_: Optional[Tensor] = None,
        Δf: Optional[Tensor] = None, Δf_: Optional[Tensor] = None,
        b1Map: Optional[Tensor] = None, b1Map_: Optional[Tensor] = None
    ) -> Tensor:
        r"""Apply a pulse to the spinarray object

        Typical usage:
            ``M = spinarray.applypulse(pulse, *, loc, doEmbed=True, doRelax,``\
            `` doUpdate, Δf, b1Map)``
            ``M_ = spinarray.applypulse(pulse, *, loc_, doEmbed=False, `` \
            ``doRelax, doUpdate, Δf_, b1Map_)``
        Inputs:
            - ``pulse``: mrphy.mobjs.Pulse.
            - ``loc`` ⊻ ``loc_``: `(N,*Nd ⊻ nM,xyz)`, "cm", locations.
        Optionals:
            - ``doEmbed``: [t/F], return ``M`` or ``M_``
            - ``doRelax``: [T/f], do relaxation during Bloch simulation.
            - ``doUpdate``: [t/F], update ``self.M_``
            - ``Δf``⊻ ``Δf_``: `(N,*Nd ⊻ nM)`, "Hz", off-resonance.
            - ``b1Map`` ⊻ ``b1Map_``: `(N,*Nd ⊻ nM,xy,(nCoils))`, transmit \
              sensitivity.
        Outputs:
            - ``M`` ⊻ ``M_``: `(N,*Nd ⊻ nM,xyz)`

        .. note::
            When ``doUpdate == True and doEmbed == False``, the output compact
            magnetization Tensor is a reference to ``self.M_``, and needs
            caution when being accessed.
        """
        assert ((loc_ is None) != (loc is None))  # XOR
        loc_ = (loc_ if loc is None else self.extract(loc))

        assert ((Δf_ is None) or (Δf is None))
        Δf_ = (Δf_ if Δf is None else self.extract(Δf))

        assert ((b1Map_ is None) or (b1Map is None))
        b1Map_ = (b1Map_ if b1Map is None else self.extract(b1Map))

        beff_ = self.pulse2beff(pulse, loc_=loc_,
                                Δf_=Δf_, b1Map_=b1Map_, doEmbed=False)

        if doRelax:
            kw_bsim = {'T1': self.T1_, 'T2': self.T2_}
        else:
            kw_bsim = {'T1': None, 'T2': None}

        kw_bsim['γ'] = self.γ_
        kw_bsim['dt'] = pulse.dt

        M_ = sims.blochsim(self.M_, beff_, **kw_bsim)
        if doUpdate:
            self.M_ = M_
        M_ = (self.embed(M_) if doEmbed else M_)
        return M_

    def asdict(self, *, toNumpy: bool = True, doEmbed: bool = True) -> dict:
        r"""Convert mrphy.mobjs.SpinArray object to dict

        Usage:
            ``d = spinarray.asdict(*, toNumpy, doEmbed)``

        Inputs:
            - ``toNumpy``: [T/f], convert ``Tensor`` to Numpy arrays.
            - ``doEmbed``: [T/f], embed compactly stored (nM) data to the \
              mask (\*Nd).
        Outputs:
            - ``d``: dict, dictionary with detached data identical to the \
              object.
        """
        fn_np = ((lambda x: x.detach().cpu().numpy()) if toNumpy else
                 (lambda x: x.detach()))

        _ = (('T1', 'T2', 'γ', 'M') if doEmbed else ('T1_', 'T2_', 'γ_', 'M_'))
        d = {k: fn_np(getattr(self, k)) for k in _}
        d['mask'] = fn_np(getattr(self, 'mask'))

        d.update({k: getattr(self, k) for k in ('shape', 'device', 'dtype')})
        return d

    def crds_(self, crds: list) -> list:
        r"""Compute crds for compact attributes

        Data in a SpinArray object is stored compactly, such that only those
        correspond to ``1`` on the ``spinarray.mask`` is kept.
        This function is provided to facilitate indexing the compact data from
        regular indices, by computing (ix, iy, iz) -> iM

        Usage:
            ``crds_ = spinarray.crds_(crds)``
        Inputs:
            - ``crds``: indices for indexing non-compact attributes.
        Outputs:
            - ``crds_``: list, ``len(crds_) == 2+len(crds)-self.ndim``.

        ``v_[crds_] == v[crds]``, when ``v_[crds_]=new_value`` is effective.
        """
        mask, ndim, nM = self.mask, self.ndim, self.nM
        assert (len(crds) >= ndim)
        crds_ = [crds[i] for i in (0,)+tuple(range(ndim, len(crds)))]
        m = torch.zeros(mask.shape, dtype=tensor(mask.numel()).dtype)-1
        m[mask] = torch.arange(nM)
        inds_ = [ind_ for ind_ in m[[[0]]+crds[1:ndim]].tolist() if ind_ != -1]

        crds_.insert(1, inds_)

        return crds_

    def dim(self) -> int:
        r"""Nd of the spinarray object, syntax sugar for len(spinarray.shape)

        Usage:
            ``Nd = spinarray.dim()``
        """
        return len(self.shape)

    def embed(self, v_: Tensor, *, out: Optional[Tensor] = None) -> Tensor:
        """Embed compact data into the spinarray.mask

        Usage:
            ``out = spinarray.embed(v_, *, out)``
        Inputs:
            - ``v_``: `(N, nM, ...)`, must be contiguous.
        Optionals:
            - ``out``: `(N, *Nd, ...)`, in-place holder.
        Outputs:
            - ``out``: `(N, *Nd, ...)`.
        """
        oshape = self.shape+v_.shape[2:]
        out = (v_.new_full(oshape, float('NaN')) if out is None else out)
        mask = self.mask.expand(self.shape)
        out[mask] = v_.view((-1,)+v_.shape[2:])
        # `v.reshape()` has intermediate alloc, leaving `out` pointless.
        # out[mask] = v_.reshape((-1,)+v_.shape[2:])
        return out

    def extract(self, v: Tensor, *, out_: Optional[Tensor] = None) -> Tensor:
        r"""Extract data with the spinarray.mask, making it compact

        Usage:
            ``out_ = spinarray.extract(v, *, out_)``
        Inputs:
            - ``v``: `(N, *Nd, ...)`.
        Optionals:
            - ``out_``: `(N, nM, ...)`, in-place holder, must be contiguous.
        Outputs:
            - ``out_``: `(N, nM, ...)`.
        """
        oshape = (self.shape[0], self.nM)+v.shape[self.ndim:]
        out_ = (v.new_empty(oshape) if out_ is None else out_)
        mask = self.mask.expand(self.shape)
        # ! do NOT use ``out_.reshape()`; It creats new tensor when should
        # fail instead.
        out_.view((-1,)+v.shape[self.ndim:]).copy_(v[mask])
        # ``v[mask].reshape()`` has intermediate alloc, leaving ``out_``
        # pointless.
        # out_.copy_(v[mask].reshape((-1,)+v.shape[self.ndim:]))
        return out_

    def freeprec(
        self, dur: Tensor, *,
        doEmbed: bool = False, doRelax: bool = True, doUpdate: bool = False,
        Δf: Optional[Tensor] = None, Δf_: Optional[Tensor] = None
    ) -> Tensor:
        r"""Free precession of duration ``dur``

        Typical usage:
            ``M = obj.freeprec(dur, doEmbed=True, doRelax, doUpdate, Δf)``
            ``M_ = obj.applypulse(dur, doEmbed=False, doRelax, doUpdate, Δf_)``
        Inputs:
            - ``dur``: `()` ⊻ `(N ⊻ 1,)`, "Sec", duration of free-precession.
        Optionals:
            - ``doEmbed``: [t/F], return ``M`` or ``M_``
            - ``doRelax``: [T/f], do relaxation during free precession.
            - ``doUpdate``: [t/F], update ``self.M_``
            - ``Δf``⊻ ``Δf_``: `(N ⊻ 1,*Nd ⊻ nM)`, "Hz", off-resonance.
        Outputs:
            - ``M`` ⊻ ``M_``: `(N,*Nd ⊻ nM,xyz)`

        .. note::
            When ``doUpdate == True and doEmbed == False``, the output compact
            magnetization Tensor is a reference to ``self.M_``, and needs
            caution when being accessed.
        """
        assert ((Δf_ is None) or (Δf is None))
        Δf_ = (Δf_ if Δf is None else self.extract(Δf))

        if doRelax:
            kw_bsim = {'T1': self.T1_, 'T2': self.T2_}
        else:
            kw_bsim = {'T1': None, 'T2': None}

        M_ = sims.freeprec(self.M_, dur, **kw_bsim, Δf=Δf_)
        if doUpdate:
            self.M_ = M_
        M_ = (self.embed(M_) if doEmbed else M_)
        return M_

    def mask_(self, *, mask: Tensor) -> Tensor:
        r"""Extract the compact region of an input external ``mask``.

        Usage:
            ``mask_ = spinarray.mask_(mask)``
        Inputs:
            - ``mask``: `(1, *Nd)`.
        Outputs:
            - ``mask_``: `(1, nM)`, ``mask_`` can be used on compact \
              attributes.
        """
        mask_ = mask(self.mask).reshape((1, -1))
        return mask_

    def numel(self) -> int:
        r"""Number of spins for the spinarray object, incompact.

        Syntax sugar of ``spinarray.mask.numel()``, effectively
        ``prod(spinarray.size())``.

        Usage:
            ``res = spinarray.numel()``
        """
        return self.mask.numel()

    def pulse2beff(
        self, pulse: Pulse, *, doEmbed: bool = False,
        loc: Optional[Tensor] = None, loc_: Optional[Tensor] = None,
        Δf: Optional[Tensor] = None, Δf_: Optional[Tensor] = None,
        b1Map: Optional[Tensor] = None, b1Map_: Optional[Tensor] = None
    ) -> Tensor:
        r"""Compute B-effective of ``pulse`` with the spinarray's parameters

        Typical usage:
            ``beff = spinarray.pulse2beff(pulse, *, loc, doEmbed=True, Δf, ``\
            ``b1Map)``
            ``beff_ = spinarray.pulse2beff(pulse, *, loc_, doEmbed=False, ``\
            ``Δf_, b1Map_)``
        Inputs:
            - ``pulse``: mrphy.mobjs.Pulse.
            - ``loc`` ⊻ ``loc_``: `(N,*Nd ⊻ nM,xyz)`, "cm", locations.
        Optionals:
            - ``doEmbed``: [t/F], return ``beff`` or ``beff_``
            - ``Δf`` ⊻ ``Δf_``: `(N,*Nd ⊻ nM)`, "Hz", off-resonance.
            - ``b1Map`` ⊻ ``b1Map_``: `(N,*Nd ⊻ nM,xy,(nCoils))`, transmit \
              sensitivity.
        Outputs:
            - ``beff`` ⊻ ``beff_``: `(N,*Nd ⊻ nM,xyz,nT)`.
        """
        assert ((loc_ is None) != (loc is None))  # XOR
        loc_ = (loc_ if loc is None else self.extract(loc))

        assert ((Δf_ is None) or (Δf is None))
        Δf_ = (Δf_ if Δf is None else self.extract(Δf))

        assert ((b1Map_ is None) or (b1Map is None))
        b1Map_ = (b1Map_ if b1Map is None else self.extract(b1Map))

        pulse = pulse.to(device=self.device, dtype=self.dtype)
        beff_ = pulse.beff(loc_, γ=self.γ_, Δf=Δf_, b1Map=b1Map_)
        beff_ = (self.embed(beff_) if doEmbed else beff_)
        return beff_

    def size(self) -> tuple:
        r"""Size of the spinarray object.

        Syntax sugar of ``spinarray.shape``.

        Usage:
            ``sz = spinarray.size()``
        """
        return self.shape

    def to(
        self, *,
        device: torch.device = torch.device('cpu'),
        dtype: torch.dtype = torch.float32
    ) -> 'SpinArray':
        r"""Duplicate the object to the prescribed device with dtype

        Usage:
            ``new_spinarray = spinarray.to(*, device, dtype)``
        Inputs:
            - ``device``: torch.device
            - ``dtype``: torch.dtype
        Outputs:
            - ``new_spinarray``: mrphy.mobjs.SpinArray object
        """
        if self.device == device and self.dtype == dtype:
            return self
        return SpinArray(self.shape, self.mask, T1_=self.T1_, T2_=self.T2_,
                         γ_=self.γ_, M_=self.M_, device=device, dtype=dtype)

class SpinCube(SpinArray):
    r"""mrphy.mobjs.SpinCube object

    Usage:
        ``SpinCube(shape, fov, mask, *, ofst, Δf_, T1_, T2_, γ_, M_, device,``\
        `` dtype)``
        ``SpinCube(shape, fov, mask, *, ofst, Δf, T1, T2, γ, M, device,``\
        '' dtype)``
    Inputs:
        - ``shape``: tuple, e.g., ``(N, nx, ny, nz)``.
        - ``fov``: `(N, xyz)`, "cm", field of view.
    Optionals:
        - ``mask``: `(1, *Nd)`, where does compact attributes locate in `Nd`.
        - ``ofst``: `(N, xyz)`, Tensor "cm", fov offset from iso-center.
        - ``Δf`` ⊻ ``Δf_``: `(N, *Nd ⊻ nM)`, "Hz", off-resonance map.
        - ``T1`` ⊻ ``T1_``: `(N, *Nd ⊻ nM)`, "Sec", T1 relaxation coeff.
        - ``T2`` ⊻ ``T2_``: `(N, *Nd ⊻ nM)`, "Sec", T2 relaxation coeff.
        - ``γ`` ⊻ ``γ_``: `(N, *Nd ⊻ nM)`,  "Hz/Gauss", gyro ratio.
        - ``M`` ⊻ ``M_``: `(N, *Nd ⊻ nM, xyz)`, spins, equilibrium ``[0 0 1]``.
        - ``device``: torch.device.
        - ``dtype``: torch.dtype

    Properties:
        - ``spinarray``: SpinArray object.
        - ``Δf_``: `(N, nM)`, "Hz", off-resonance map.
        - ``loc_``: `(N, nM, xyz)`, "cm", location of spins.
        - ``fov``: `(N, xyz)`, "cm", field of view.
        - ``ofst``: `(N, xyz)`, "cm", fov offset from iso-center.
    """

    _readonly = ('spinarray', 'loc_')
    _compact = ('Δf_', 'loc_')  # `loc_` depends on `shape`, `fov` and `ofst`
    __slots__ = set(_readonly+_compact+('fov', 'ofst'))

    def __init__(
        self, shape: tuple, fov: Tensor, *, mask: Optional[Tensor] = None,
        ofst: Tensor = tensor([[0., 0., 0.]]),
        Δf: Optional[Tensor] = None, Δf_: Optional[Tensor] = None,
        T1: Optional[Tensor] = None, T1_: Optional[Tensor] = None,
        T2: Optional[Tensor] = None, T2_: Optional[Tensor] = None,
        γ: Optional[Tensor] = None,  γ_: Optional[Tensor] = None,
        M: Optional[Tensor] = None,  M_: Optional[Tensor] = None,
        device: torch.device = torch.device('cpu'),
        dtype: torch.dtype = torch.float32
    ):
        # 1) `SpinCube` is subclassed to `SpinArray`.
        # 2) Attribute `spinarray` is included to enable extracting the
        # `SpinArray` part of a `SpinCube()` instance.
        # 3) `SpinCube.__getattr__` is tweaked for immitating super class
        # access to `SpinCube().spinarray`'s attributes.
        # To have these three features, we cannot have `super().__init__()` in
        # `SpinCube`, which will disable the `__getattr__` tweak.
        sp = SpinArray(shape, mask, T1=T1, T1_=T1_, T2=T2, T2_=T2_, γ=γ, γ_=γ_,
                       M=M, M_=M_, device=device, dtype=dtype)
        object.__setattr__(self, 'spinarray', sp)

        kw = {'device': sp.device, 'dtype': sp.dtype}
        # setattr(self, k, v), avoid computing `loc_` w/ `fov` & `ofst` not set
        object.__setattr__(self, 'fov', fov.to(**kw))
        object.__setattr__(self, 'ofst', ofst.to(**kw))
        # Initialize ``loc_`` in memory, reuse it.
        object.__setattr__(self, 'loc_',
                           torch.zeros((sp.shape[0], sp.nM, 3), **kw))
        self._update_loc_()  # compute ``loc_`` from set ``fov`` & ``ofst`

        assert((Δf is None) or (Δf_ is None))
        if Δf is None:
            self.Δf_ = (tensor(0.) if Δf_ is None else Δf_)
        else:
            self.Δf = Δf

        return

    def __getattr__(self, k):  # provoked only when `__getattribute__` failed
        if k+'_' not in self._compact:  # k not in ('Δf_', 'loc')
            # Cannot do `self.spinarray` or `getattr(self, 'spinarray')` here.
            # They cause infinite recursion when `SpinCube().__getattr__()` is
            # queried with `spinarray`, which may happen during `deepcopy` or
            # `pickle`.
            # Therefore, call `object.__getattribute__()` here, and let it fail
            # when it should.
            spinarray = object.__getattribute__(self, 'spinarray')
            try:
                return getattr(spinarray, k)
            except AttributeError:
                raise AttributeError(f"'SpinCube' has no attribute '{k}'")

        v_, sp = getattr(self, k+'_'), self.spinarray
        return (sp.embed(v_) if sp.nM != np.prod(sp.shape[1:]) else
                v_.reshape(sp.shape+v_.shape[2:]))  # `mask` is all True

    def __setattr__(self, k_, v_):
        if 'deepcopy' in (_.function for _ in inspect.stack()):
            # Hack, this enables `deepcopy()` w/o overriding `__deepcopy__()`.
            # Generator is faster than list comprehension.
            object.__setattr__(self, k_, v_)
            return

        if (k_ in self._readonly) or (k_+'_' in self._readonly):
            raise AttributeError(f"'SpinCube' object attribute '{k_}'"
                                 " is read-only")

        sp = self.spinarray
        if k_ in SpinArray.__slots__ or k_+'_' in SpinArray.__slots__:
            setattr(sp, k_, v_)
            return

        kw = {'device': sp.device, 'dtype': sp.dtype}
        v_ = (v_.to(**kw) if isinstance(v_, Tensor) else tensor(v_, **kw))

        shape = sp.shape
        if k_+'_' in self._compact:  # `loc_` excluded by beginning assert
            k_ = k_+'_'
            v_ = self.extract(v_.expand(shape+(3,) if k_ == 'loc_' else shape))

        if k_ == 'Δf_':
            v_ = v_.expand((shape[0], sp.nM))  # (N, nM)
        elif k_ in ('fov', 'ofst'):
            assert(v_.ndim == 2)

        object.__setattr__(self, k_, v_)

        # update `loc_` when needed
        if k_ in ('fov', 'ofst'):
            self._update_loc_()
        return

    def _update_loc_(self):
        r"""Update ``spincube.loc_`` using FOV and offset

        The ``spincube``'s spin locations are computed internally from set FOV
        and offset.

        Usage:
            ``loc_ = spincube._update_loc_()``
        """
        loc_, fov, ofst = self.loc_, self.fov, self.ofst
        sp = self.spinarray
        kw = {'device': sp.device, 'dtype': sp.dtype}

        # locn (1, prod(Nd), xyz)  normalized locations, [-0.5, 0.5)
        shape, mask = sp.shape, sp.mask
        crdn = ((torch.arange(x, **kw)-utils.ctrsub(x))/x for x in shape[1:])
        _locn = torch.meshgrid(*crdn)  # ((*Nd,), (*Nd), (*Nd))

        for i in range(3):  # xyz, (N, nM)
            # According to `memory_profiler`, this does not provoke allocs.
            # `torch.addr`'s `vec2`, _locn[i][mask[0, ...]], provokes alloc.
            loc_[..., i] = (fov[:, None, i]*_locn[i][mask[0, ...]][None, ...]
                            + ofst[:, None, i])

        return

    def applypulse(
        self, pulse: Pulse, *,
        doEmbed: bool = False, doRelax: bool = True, doUpdate: bool = False,
        b1Map: Optional[Tensor] = None, b1Map_: Optional[Tensor] = None
    ) -> Tensor:
        r"""Apply a pulse to the spincube object

        Usage:
            ``M = spincube.applypulse(pulse, *, doEmbed=True, doRelax, b1Map)``
            ``M_ = spincube.applypulse(pulse, *, doEmbed=False, doRelax,``\
            `` b1Map_)``

        Inputs:
            - ``pulse``: mobjs.Pulse object.
        Optionals:
            - ``doEmbed``: [t/F], return ``M`` or ``M_``.
            - ``doRelax``: [T/f], do relaxation during Bloch simulation.
            - ``b1Map`` ⊻ ``b1Map_``: `(N,*Nd ⊻ nM,xy,(nCoils))`, transmit \
              sensitivity.
        Outputs:
            - ``M`` ⊻ ``M_``: `(N,*Nd ⊻ nM,xyz)`.
        """
        assert ((b1Map_ is None) or (b1Map is None))
        b1Map_ = (b1Map_ if b1Map is None else self.extract(b1Map))

        return self.spinarray.applypulse(pulse, doEmbed=doEmbed,
                                         doRelax=doRelax, doUpdate=doUpdate,
                                         Δf_=self.Δf_, loc_=self.loc_,
                                         b1Map_=b1Map_)

    def freeprec(
        self, dur: Tensor, *,
        doEmbed: bool = False, doRelax: bool = True, doUpdate: bool = False
    ) -> Tensor:
        r"""Free precession of duration ``dur``

        Typical usage:
            ``M = obj.freeprec(dur, doEmbed=True, doRelax, doUpdate)``
            ``M_ = obj.applypulse(dur, doEmbed=False, doRelax, doUpdate)``
        Inputs:
            - ``dur``: `()` ⊻ `(N ⊻ 1,)`, "Sec", duration of free-precession.
        Optionals:
            - ``doEmbed``: [t/F], return ``M`` or ``M_``
            - ``doRelax``: [T/f], do relaxation during free precession.
            - ``doUpdate``: [t/F], update ``self.M_``
        Outputs:
            - ``M`` ⊻ ``M_``: `(N,*Nd ⊻ nM,xyz)`

        .. note::
            When ``doUpdate == True and doEmbed == False``, the output compact
            magnetization Tensor is a reference to ``self.M_``, and needs
            caution when being accessed.
        """

        return self.spinarray.freeprec(dur, Δf_=self.Δf_, doEmbed=doEmbed,
                                       doRelax=doRelax, doUpdate=doUpdate)

    def asdict(self, *, toNumpy: bool = True, doEmbed: bool = True) -> dict:
        r"""Convert mrphy.mobjs.SpinCube object to dict

        Usage:
            ``d = spincube.asdict(*, toNumpy, doEmbed)``

        Inputs:
            - ``toNumpy``: [T/f], convert ``Tensor`` to Numpy arrays.
            - ``doEmbed``: [T/f], embed compactly stored (nM) data to the \
              mask `(*Nd)`.
        Outputs:
            - ``d``: dict, dictionary with detached data identical to the \
              object.
        """
        fn_np = ((lambda x: x.detach().cpu().numpy()) if toNumpy else
                 (lambda x: x.detach()))

        _ = (('loc', 'Δf') if doEmbed else ('loc', 'Δf'))
        d = {k: fn_np(getattr(self, k)) for k in _}

        d.update({k: getattr(self, k) for k in ('fov', 'ofst')})

        d.update(self.spinarray.asdict(toNumpy=toNumpy, doEmbed=doEmbed))
        return d

    def pulse2beff(
        self, pulse: Pulse, *,
        doEmbed: bool = False,
        b1Map: Optional[Tensor] = None, b1Map_: Optional[Tensor] = None
    ) -> Tensor:
        r"""Compute B-effective of ``pulse`` with the spincube's parameters

        Typical usage:
            ``beff = spincube.pulse2beff(pulse, *, doEmbed=True, b1Map)``
            ``beff_ = spincube.pulse2beff(pulse, *, doEmbed=False, b1Map_)``
        Inputs:
            - ``pulse``: mrphy.mobjs.Pulse.
        Optionals:
            - ``doEmbed``: [t/F], return ``beff`` or ``beff_``.
            - ``b1Map`` ⊻ ``b1Map_``: `(N,*Nd ⊻ nM,xy,(nCoils))`, transmit \
              sensitivity.
        Outputs:
            - ``beff`` ⊻ ``beff_``: `(N,*Nd ⊻ nM,xyz,nT)`.
        """
        return self.spinarray.pulse2beff(pulse, self.loc_, doEmbed=doEmbed,
                                         Δf_=self.Δf_,
                                         b1Map=b1Map, b1Map_=b1Map_)

    def to(
        self, *,
        device: torch.device = torch.device('cpu'),
        dtype: torch.dtype = torch.float32
    ) -> 'SpinCube':
        r"""Duplicate the object to the prescribed device with dtype

        Usage:
            ``new_spincube = spincube.to(*, device, dtype)``
        Inputs:
            - ``device``: torch.device.
            - ``dtype``: torch.dtype.
        Outputs:
            - ``new_spincube``: mrphy.mobjs.SpinCube object.
        """
        if self.device == device and self.dtype == dtype:
            return self
        return SpinCube(self.shape, self.fov, ofst=self.ofst, Δf_=self.Δf_,
                        T1_=self.T1_, T2_=self.T2_, γ_=self.γ_, M_=self.M_,
                        device=device, dtype=dtype)

class SpinBolus(SpinArray):
    def __init__(
        self
    ):
        pass
    pass

class Examples(object):
    r"""Class for quickly creating exemplary instances to play around with.
    """
    @staticmethod
    def pulse() -> Pulse:
        r"""Create a mrphy.mobjs.Pulse object.
        """
        device = torch.device('cpu')
        dtype = torch.float32

        kw = {'dtype': dtype, 'device': device}
        N, nT, dt = 1, 512, dt0

        # pulse: Sec; Gauss; Gauss/cm.
        pulse_size = (N, 1, nT)
        t = torch.arange(0, nT, **kw).reshape(pulse_size)
        rf = 10*torch.cat([torch.cos(t/nT*2*π),                # (1,xy, nT)
                           torch.sin(t/nT*2*π)], 1)
        gr = torch.cat([torch.ones(pulse_size, **kw),
                        torch.ones(pulse_size, **kw),
                        10*torch.atan(t - round(nT/2))/π], 1)  # (1,xyz,nT)

        # Pulse
        p = Pulse(rf=rf, gr=gr, dt=dt, **kw)
        return p

    @staticmethod
    def spinarray() -> SpinArray:
        r"""Create a mrphy.mobjs.SpinArray object.
        """
        device = torch.device('cpu')
        dtype = torch.float32
        kw = {'dtype': dtype, 'device': device}

        N, Nd, γ_ = 1, (3, 3, 3), γH
        shape = (N, *Nd)
        mask = torch.zeros((1,)+Nd, device=device, dtype=torch.bool)
        mask[0, :, 1, :], mask[0, 1, :, :] = True, True
        T1_, T2_ = tensor([[1.]], **kw), tensor([[4e-2]], **kw)

        array = SpinArray(shape, mask=mask, T1_=T1_, T2_=T2_, γ_=γ_, **kw)
        return array

    @staticmethod
    def spincube() -> SpinCube:
        r"""Create a mrphy.mobjs.SpinCube object.
        """
        device = torch.device('cpu')
        dtype = torch.float32
        kw = {'dtype': dtype, 'device': device}

        N, Nd, γ_ = 1, (3, 3, 3), γH
        shape = (N, *Nd)
        mask = torch.zeros((1,)+Nd, device=device, dtype=torch.bool)
        mask[0, :, 1, :], mask[0, 1, :, :] = True, True
        fov, ofst = tensor([[3., 3., 3.]], **kw), tensor([[0., 0., 1.]], **kw)
        T1_, T2_ = tensor([[1.]], **kw), tensor([[4e-2]], **kw)

        cube = SpinCube(shape, fov, mask=mask, ofst=ofst,
                        T1_=T1_, T2_=T2_, γ_=γ_, **kw)

        cube.Δf = torch.sum(-cube.loc[0:1, :, :, :, 0:2], dim=-1) * γ_
        return cube


"""
An Implementation of Euler Integration for reaction path following.

References:
[1] C. Gonzalez and H. B. Schlegel, J. Chem. Phys. 90(4):2154 (An Improved Algorithm for Reaction-Path Following)
"""

__authors__  = "Wallace D. Derricotte"
__credits__  = ["Wallace D. Derricotte"]

__copyright__  = "(c) 2018-2019, Derricotte Research Group"
__license__    = "MIT License"
__date__       = "2019-01-02"

import numpy as np
import psi4
import os
import sys
import json
from pyscf import gto, scf, dft, grad, solvent, mp

########################
## Gradient Functions ##
########################

def energy_calc(params, current_geom, mol):
    energy = 0.0
    if(params.qm_program=='pyscf'):
        pymol = gto.Mole()
        pymol.verbose = 0
        geom_vec = []
        for i in range(params.natoms):
            atom = [params.symbols[i],]
            atom_coords = []
            for j in range(3):
                atom_coords.append(current_geom[i][j])
            atom_coords = tuple(atom_coords)
            atom.append(atom_coords)
            geom_vec.append(atom)
        #print(geom_vec)
        pymol.atom = geom_vec
        pymol.unit = 'Bohr'
        pymol.basis = params.basis
        pymol.charge = params.charge
        pymol.spin = params.mult - 1
        pymol.build()
        if(params.method == "scf"):
            scf_obj = scf.RHF(pymol)
        #if(params.method == "mp2"): #TODO Doesn't work yet. Few things left to figure out.
        #    scf_obj = scf.RHF(pymol).run()
        #    scf_obj = mp.MP2(scf_obj).run()
        #if(params.method == "dft"):
        #    scf_obj = dft.RKS(mol)
        #    scf_obj.xc = params.xc_functional
        if(params.do_solvent):
            solv_obj = solvent.ddCOSMO(scf_obj)
            solv_obj.with_solvent.eps = params.eps
            solv_obj.run()
            energy = solv_obj.kernel()
            print(energy)
        else:
            energy = scf_obj.scf()

    if(params.qm_program=='psi4'):
        mol.set_geometry(psi4.core.Matrix.from_array(current_geom))
        grad_method = "%s/%s" %(params.method,params.basis)
        psi4.core.set_output_file("psi4_out.dat", False)
        psi4.set_options(params.keywords)
        psi4.set_num_threads(params.nthreads)
        energy = psi4.energy(grad_method)
    return energy

def grad_calc(params,current_geom, mol):
    """
        Uses Psi4 to calculate the energy gradient and returns the mass-weighted
        gradient and energy. Here any of the keywords the user provides in the
        .json input are used to set the options for the energy calculation.

        Parameters:
        ----------
            params(self) -- contains initialized shared parameters.
            current_geom(np array) -- Matrix of size natoms x 3 containing the geometry.
            mol(psi4.Molecule) -- Psi4 molecule object containing the current molecule.
        Returns:
        -------
            grad_mw(np array) -- Mass weighted gradient matrix of size natoms x 3.
            E(float) -- single-point energy from Psi4 calculation.
    """
    if(params.qm_program=='pyscf'):
        pymol = gto.Mole()
        pymol.verbose = 0
        geom_vec = []
        for i in range(params.natoms):
            atom = [params.symbols[i],]
            atom_coords = []
            for j in range(3):
                atom_coords.append(current_geom[i][j])
            atom_coords = tuple(atom_coords)
            atom.append(atom_coords)
            geom_vec.append(atom)
        #print(geom_vec)
        pymol.atom = geom_vec
        pymol.unit = 'Bohr'
        pymol.basis = params.basis
        pymol.charge = params.charge
        pymol.spin = params.mult - 1
        pymol.build()
        if(params.method == "scf"):
            scf_obj = scf.RHF(pymol)
        #if(params.method == "mp2"): #TODO Doesn't work yet. Few things left to figure out.
        #    mf = scf.RHF(pymol).run()
        #    scf_obj = mp.MP2(mf).run()
        if(params.do_solvent):
            solv_obj = solvent.ddCOSMO(scf_obj)
            solv_obj.with_solvent.eps = params.eps
            solv_obj.run()
            E = solv_obj.kernel()
            grad = solv_obj.nuc_grad_method().kernel()
        else:
            E = scf_obj.kernel()
            grad = scf_obj.nuc_grad_method().kernel()
        #print(grad)
        grad_mw = mass_weight(params.natoms, grad, mol)
    if(params.qm_program=='psi4'):
        mol.set_geometry(psi4.core.Matrix.from_array(current_geom))
        mol.fix_orientation(True)
        mol.fix_com(True)
        mol.reset_point_group('c1')
        grad_method = "%s/%s" %(params.method,params.basis)
        psi4.core.set_output_file("psi4_out.dat", False)
        psi4.set_options(params.keywords)
        psi4.set_num_threads(params.nthreads)
        E, wfn = psi4.energy(grad_method,return_wfn=True)
        psi4.set_num_threads(params.nthreads)
        grad = np.asarray(psi4.gradient(grad_method,ref_wfn=wfn))
        #print(grad)
        grad_mw = mass_weight(params.natoms, grad, mol)
    return grad_mw, E

def parabolic_fit(xs, ys):
    fit = np.polyfit(xs, ys, deg=2)
    fit = np.poly1d(fit)
    minima = fit.deriv().r
    real_minima = minima[minima.imag==0].real
    return real_minima

def ishida_morokuma(params,output_file):
    """
        This function runs the Ishida-Morokuma irc procedure
    """
    max_steps = 1000
    #params = Params()
    line_step_size = 0.3333*params.step_size
    #line_step_size = 0.025*params.step_size
    current_geom = params.geometry
    mol = psi4.geometry(params.geometry)
    starting_vec = np.asarray(params.ts_vec)
    grad_method = "%s/%s" %(params.method,params.basis)
    steps = 0
    E = 0.0
    previous_E = 0.0
    del_E = 0.0
    energies = []
    current_geom = np.asarray(mol.geometry())
    #print(current_geom)
    output = open(output_file, "a")
    output.write('\n\n--Intrinsic Reaction Coordinate (%s)--\n' %(params.direction))
    output.write('\n--------------------------------------------------------------------------------------\n')
    output.write('\n{:>20} {:>20} {:>20} {:>20}\n'.format('Coordinate', 'E', 'Delta E', 'Gradient Norm'))
    output.write('-------------------------------------------------------------------------------------\n')
    output.close()
    last_energy = None
    while (steps <= max_steps):
        if(steps==0):
            grad_0 = mass_weight(params.natoms, starting_vec, mol)
            E_0 = energy_calc(params, current_geom, mol)
        else:
            grad_0, E_0  = grad_calc(params, current_geom, mol)

        if(last_energy):
            del_E = E_0 - last_energy
        if((last_energy and (E_0 > last_energy) and steps > params.grace_period)):
            output = open(output_file, "a")
            print("\nIRC Energy Has Increased! You're Likely Near a Minimum!\n")
            output.close()
            break
        if(last_energy and np.abs(del_E)<params.e_conv and steps > params.grace_period):
            output = open(output_file, "a")
            print("\nIRC Has Converged!\n")
            output.close()
            break
        mol.save_xyz_file('imk_step_'+str(steps)+'.xyz',False)
        coords_1 = euler_step(params.natoms, current_geom, grad_0,params.step_size,mol)
        current_geom = coords_1
        #mol.set_geometry(psi4.core.Matrix.from_array(current_geom))
        grad_1, E_1 = grad_calc(params, current_geom, mol)
        grad_0_norm = np.linalg.norm(grad_0)
        grad_1_norm = np.linalg.norm(grad_1)

        # Calculate Bisector (Eq. 6)
        D = grad_0/grad_0_norm - grad_1/grad_1_norm
        D_normed = D/np.linalg.norm(D)

        line_xs = [0,]
        line_energies = [E_1,]

        line_step_size_thresh = 1.5*line_step_size
        #line_step_size_thresh = 2.0*line_step_size

        # Find useful point by projecting grad_1 on D
        grad_1_normed = grad_1/grad_1_norm
        step_D1 = grad_1*D_normed*D_normed*line_step_size
        step_D1_norm = np.linalg.norm(step_D1)
       # if step_D1_norm < line_step_size_thresh:
       #     coords_1 = mass_weight_geom(params.natoms, coords_1, mol)
       #     current_geom = coords_1 + step_D1
       #     current_geom = un_mass_weight_geom(params.natoms, current_geom, mol)
       #     mol.set_geometry(psi4.core.Matrix.from_array(current_geom))
       #     step_D1_E = psi4.energy(grad_method)
       #     line_xs.append(step_D1_norm)
       #     line_energies.append(step_D1_E)
        # Otherwise take a step along D
       # else:
        step_D2 = line_step_size*D_normed
        step_D2_norm = np.linalg.norm(step_D2)
        coords_1 = mass_weight_geom(params.natoms, coords_1, mol)
        current_geom = coords_1 + step_D2
        current_geom = un_mass_weight_geom(params.natoms, coords_1, mol)
        #mol.set_geometry(psi4.core.Matrix.from_array(current_geom))
        #step_D2_E = psi4.energy(grad_method)
        step_D2_E = energy_calc(params, current_geom, mol)
        #line_xs.append(step_D2_norm)
        line_xs.append(step_D2_norm)
        line_energies.append(step_D2_E)

        # Calculate 3rd point by taking a half step size
        if(line_energies[1] >= line_energies[0]):
            step_D3 = 0.5*line_step_size*D_normed # Half Step Size
            #new_del = 0.5*line_step_size
        else:
            step_D3 = 2.0*line_step_size*D_normed #Double Step Size
            #new_del = 2.0*line_step_size

        step_D3_norm = np.linalg.norm(step_D3)
        #coords_1 = mass_weight_geom(params.natoms, coords_1, mol)
        current_geom = coords_1 + step_D3
        current_geom = un_mass_weight_geom(params.natoms, current_geom, mol)
        #mol.set_geometry(psi4.core.Matrix.from_array(current_geom))
        #step_D3_E = psi4.energy(grad_method)
        step_D3_E = energy_calc(params, current_geom, mol)
        line_xs.append(step_D3_norm)
        line_energies.append(step_D3_E)

        real_minimum = parabolic_fit(line_xs, line_energies)

        current_geom = coords_1 + (real_minimum*D_normed)
        current_geom = un_mass_weight_geom(params.natoms, current_geom, mol)

        mol.set_geometry(psi4.core.Matrix.from_array(current_geom))

        last_energy = E_0

        #params.damp -= 0.5

        if(params.direction=="backward"):
            coord = -1*steps*params.step_size
        else:
            coord = steps*params.step_size
        print_step(output_file,coord, E_0, del_E, grad_0_norm)
        steps = steps+1
    output = open(output_file, "a")
    output.write('-------------------------------------------------------------------------------------\n')
    output.close()
    with open('irc_%s.xyz' %params.direction,'w') as outfile:
        for i in range(steps):
            with open('imk_step_'+str(i)+'.xyz')as infile:
                outfile.write(infile.read())
        os.system('rm imk_step*')

def mass_weight(natoms,grad, mol):
    """
        Mass weights the given gradient

        Parameters:
        ----------
            natoms(int) -- number of atoms in the molecule.
            grad(np array) -- matrix of size natoms x 3 containing the gradients.
        Returns:
        -------
            grad_mw(np array) -- Mass weighted gradient matrix of size natoms x 3.
    """
    grad_mw = np.zeros((natoms, 3))
    for i in range(natoms):
        for j in range(3):
            grad_mw[i][j] = grad[i][j]/np.sqrt(mol.mass(i))
    return grad_mw

def mass_weight_geom(natoms,geom,mol):
    """
        Mass weights the given coordinates

        Parameters:
        ----------
            natoms(int) -- number of atoms in the molecule.
            geom(np array) -- matrix of size natoms x 3 containing the cartesian geometry (bohr).
        Returns:
        -------
            coord_mw(np array) -- Mass weighted coordinate matrix of size natoms x 3 in units
            amu^(1/2)*bohr.
    """
    coord_mw = np.zeros((natoms, 3))
    for i in range(natoms):
        for j in range(3):
            coord_mw[i][j] = geom[i][j]*np.sqrt(mol.mass(i))
    return coord_mw

def un_mass_weight_geom(natoms,geom_mw,mol):
    """
        Un-Mass weights the given coordinates. Displacements are done along the mass weighted
        coordinate and then un-mass-weighted prior to gradient calculations.

        Parameters:
        ----------
            natoms(int) -- number of atoms in the molecule.
            geom_mw(np array) -- matrix of size natoms x 3 containing the cartesian geometry
            in units amu^(1/2)*bohr
        Returns:
        -------
            coord(np array) -- Coordinate matrix of size natoms x 3 in units bohr.
    """
    coord = np.zeros((natoms, 3))
    for i in range(natoms):
        for j in range(3):
            coord[i][j] = geom_mw[i][j]/np.sqrt(mol.mass(i))
    return coord

def euler_step(natoms,current_geom,grad,step_size,mol):
    """
        Take a single Euler step along the gradient. The coordinates are first mass weighted prior
        to the Euler step and then un-mass-weighted for printing/gradient calculation. This
        procedure is adapted from Equation 2 in Ref [1].

        Parameters:
        ----------
            natoms(int) -- number of atoms in the molecule.
            current_geom(np array) -- matrix of size natoms x 3 containing the cartesian
            geometry in units bohr.
            grad(np array) -- mass-weighted gradient matrix of size natoms x 3 containing
            the gradients in units amu^(1/2)*bohr.
            step_size(float) -- user provided IRC step size in units amu^(1/2)*bohr.
    """
    current_geom = mass_weight_geom(natoms, current_geom,mol)
    grad_norm = np.asarray(np.linalg.norm(grad))
    current_geom -= step_size*(grad/grad_norm)
    #real_step = step_size/grad_norm
    #current_geom -= step_size*(grad)
    #current_geom -= real_step*(grad)
    current_geom = un_mass_weight_geom(natoms, current_geom,mol)
    return current_geom

def irc(output_file):
    """
        This function runs the irc procedure
    """
    max_steps = 1000
    params = Params()

    mol = psi4.geometry(params.geometry)
    print(mol)
    starting_vec = np.asarray(params.ts_vec)

    steps = 0
    E = 0.0
    previous_E = 0.0
    energies = []
    current_geom = np.asarray(mol.geometry())
    #print(current_geom)
    output = open(output_file, "a")
    output.write('\n\n--Intrinsic Reaction Coordinate (%s)--\n' %(params.direction))
    output.write('\n-------------------------------------------------------------------------------------')
    output.write('\n{:>20} {:>20} {:>20} {:>20}\n'.format('Coordinate', 'E', 'Delta E', 'Gradient Norm'))
    output.write('-------------------------------------------------------------------------------------\n')
    output.close()
    while (steps <= max_steps):
        mol.save_xyz_file('euler_step_'+str(steps)+'.xyz',False)
        if(steps==0):
            grad = mass_weight(params.natoms, starting_vec, mol)
        else:
            grad, E  = grad_calc(params, current_geom, mol)

        current_geom = euler_step(params.natoms, current_geom, grad,params.step_size,mol)
        if(steps > 20):
            if(E>previous_E):
                #print("pyREX: New energy is greater! Likely near a minimum!")
                break
        steps = steps + 1
        del_E = E - previous_E
        previous_E = E
        if(params.direction=="backward"):
            coord = -1*steps*params.step_size
        else:
            coord = steps*params.step_size
        print_step(output_file,coord, E, del_E, grad)
    output = open(output_file, "a")
    output.write('-------------------------------------------------------------------------------------\n')
    output.close()
    with open('irc_%s.xyz' %params.direction,'w') as outfile:
        for i in range(steps):
            with open('euler_step_'+str(i)+'.xyz')as infile:
                outfile.write(infile.read())
        os.system('rm euler_step*')

def print_step(output_file, coord, energy, del_E, grad):
    output = open(output_file, "a")
    grad_norm = np.linalg.norm(grad)
    output.write('\n{:>20.4f} {:>20.7f} {:>20.10f} {:>20.10f}\n'.format(coord, energy, del_E, grad_norm))
    output.close()


__author__ = "Luke Liu"
#encoding="utf-8"
# -*- coding: utf-8 -*-

# 引入第三方库
import tensorflow as tf
import numpy as np
import urllib
import tarfile
import os
import matplotlib.pyplot as plt
from imageio import imread, imsave, mimsave
from scipy.misc import imresize
import glob

"""
准备工作
"""
# database's path
filename = 'D:/BaiduYunDownload/python_exe/dataset/lfw.tgz'
directory = 'lfw_imgs'
new_dir = 'lfw_new_imgs'
# # 解压文件
# tar = tarfile.open(filename, 'r:gz')
# tar.extractall(path=directory)
# tar.close()
# #统计照片的个数，并且把照片存入new_dir
# count = 0
# for dir_, _, files in os.walk(directory):
#     for file_ in files:
#         img = imread(os.path.join(dir_, file_))
#         imsave(os.path.join(new_dir, '%d.png' % count), img)
#         count += 1
# print(count)
#指定dataset
dataset = 'D:\BaiduYunDownload\python_exe\dataset\scut_faces\AF' # LFW
# dataset = 'celeba' # CelebA
images = glob.glob(os.path.join(dataset, '*.*'))
print(len(images))
# 定义一个输出sample的file
OUTPUT_DIR = 'samples_'
if not os.path.exists(OUTPUT_DIR):
    os.mkdir(OUTPUT_DIR)

"""
建立参数

"""
#
batch_size = 100
# noise dim is 100
z_dim = 100
WIDTH = 64
HEIGHT = 64

# 指定输入与是否训练
X = tf.placeholder(dtype=tf.float32, shape=[None, HEIGHT, WIDTH, 3], name='X')
noise = tf.placeholder(dtype=tf.float32, shape=[None, z_dim], name='noise')
is_training = tf.placeholder(dtype=tf.bool, name='is_training')

def lrelu(x, leak=0.2):
    return tf.maximum(x, leak * x)

#先用sigmoid处理到0-1，然后进过cross-entropy
def sigmoid_cross_entropy_with_logits(x, y):
    return tf.nn.sigmoid_cross_entropy_with_logits(logits=x, labels=y)

# 判别器部分, 2 return ,
def discriminator(image, reuse=None, is_training=is_training):
    momentum = 0.9
    with tf.variable_scope('discriminator', reuse=reuse):
        #Conv1_: 64 filter is 5, stride is 2, activation is Lekrelu
        h0 = lrelu(tf.layers.conv2d(image, kernel_size=5, filters=64, strides=2, padding='same'))
        # Conv2_: 128 filter is 5 ,stride is 2,activation is lkeRelu
        h1 = tf.layers.conv2d(h0, kernel_size=5, filters=128, strides=2, padding='same')
        # BN
        h1 = lrelu(tf.contrib.layers.batch_norm(h1, is_training=is_training, decay=momentum))
        #

        h2 = tf.layers.conv2d(h1, kernel_size=5, filters=256, strides=2, padding='same')
        h2 = lrelu(tf.contrib.layers.batch_norm(h2, is_training=is_training, decay=momentum))

        h3 = tf.layers.conv2d(h2, kernel_size=5, filters=512, strides=2, padding='same')
        h3 = lrelu(tf.contrib.layers.batch_norm(h3, is_training=is_training, decay=momentum))
       # 展开，但不要连全连接层
        h4 = tf.contrib.layers.flatten(h3)
        # 直接交sigmoid
        h4 = tf.layers.dense(h4, units=1)
        return tf.nn.sigmoid(h4), h4

def generator(z, is_training=is_training):
    momentum = 0.9
    with tf.variable_scope('generator', reuse=None):
# d is the deep of the image
        d = 4
        h0 = tf.layers.dense(z, units=d * d * 512)
# 增加一个维度
        h0 = tf.reshape(h0, shape=[-1, d, d, 512])
# 前面的卷积转置都不加激活函数
        h0 = tf.nn.relu(tf.contrib.layers.batch_norm(h0, is_training=is_training, decay=momentum))

        h1 = tf.layers.conv2d_transpose(h0, kernel_size=5, filters=256, strides=2, padding='same')
        h1 = tf.nn.relu(tf.contrib.layers.batch_norm(h1, is_training=is_training, decay=momentum))

        h2 = tf.layers.conv2d_transpose(h1, kernel_size=5, filters=128, strides=2, padding='same')
        h2 = tf.nn.relu(tf.contrib.layers.batch_norm(h2, is_training=is_training, decay=momentum))

        h3 = tf.layers.conv2d_transpose(h2, kernel_size=5, filters=64, strides=2, padding='same')
        h3 = tf.nn.relu(tf.contrib.layers.batch_norm(h3, is_training=is_training, decay=momentum))

# 最后一个卷积装置的输出要经过tanh激活函数
        h4 = tf.layers.conv2d_transpose(h3, kernel_size=5, filters=3, strides=2, padding='same', activation=tf.nn.tanh,
                                        name='g')
        return h4

# 定义损失函数
# 产生噪声
g = generator(noise)
d_real, d_real_logits = discriminator(X)
d_fake, d_fake_logits = discriminator(g, reuse=True)
# generator的可训练参数以及discriminator的可训练参数
vars_g = [var for var in tf.trainable_variables() if var.name.startswith('generator')]
vars_d = [var for var in tf.trainable_variables() if var.name.startswith('discriminator')]
# generator 的loss以及 discriminator 的loss
loss_d_real = tf.reduce_mean(sigmoid_cross_entropy_with_logits(d_real_logits, tf.ones_like(d_real)))
loss_d_fake = tf.reduce_mean(sigmoid_cross_entropy_with_logits(d_fake_logits, tf.zeros_like(d_fake)))
loss_g = tf.reduce_mean(sigmoid_cross_entropy_with_logits(d_fake_logits, tf.ones_like(d_fake)))
loss_d = loss_d_real + loss_d_fake

# 优化参数
# 关于在batch_norm中，即为更新mean和variance的操作，因此需要update_ops
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
# 首先优化dis
    optimizer_d = tf.train.AdamOptimizer(learning_rate=0.0002, beta1=0.5).minimize(loss_d, var_list=vars_d)
# 然后优化gen
    optimizer_g = tf.train.AdamOptimizer(learning_rate=0.0002, beta1=0.5).minimize(loss_g, var_list=vars_g)

def read_image(path, height, width):
    image = imread(path)
    h = image.shape[0]
    w = image.shape[1]

    if h > w:
        image = image[h // 2 - w // 2: h // 2 + w // 2, :, :]
    else:
        image = image[:, w // 2 - h // 2: w // 2 + h // 2, :]

    image = imresize(image, (height, width))
    return image / 255.
# 合成多张图片
def montage(images):
    if isinstance(images, list):
        images = np.array(images)
    img_h = images.shape[1]
    img_w = images.shape[2]
    n_plots = int(np.ceil(np.sqrt(images.shape[0])))
    if len(images.shape) == 4 and images.shape[3] == 3:
        m = np.ones(
            (images.shape[1] * n_plots + n_plots + 1,
             images.shape[2] * n_plots + n_plots + 1, 3)) * 0.5
    elif len(images.shape) == 4 and images.shape[3] == 1:
        m = np.ones(
            (images.shape[1] * n_plots + n_plots + 1,
             images.shape[2] * n_plots + n_plots + 1, 1)) * 0.5
    elif len(images.shape) == 3:
        m = np.ones(
            (images.shape[1] * n_plots + n_plots + 1,
             images.shape[2] * n_plots + n_plots + 1)) * 0.5
    else:
        raise ValueError('Could not parse image shape of {}'.format(images.shape))
    for i in range(n_plots):
        for j in range(n_plots):
            this_filter = i * n_plots + j
            if this_filter < images.shape[0]:
                this_img = images[this_filter]
                m[1 + i + i * img_h:1 + i + (i + 1) * img_h,
                  1 + j + j * img_w:1 + j + (j + 1) * img_w] = this_img
    return m

# 进行训练
sess = tf.Session()
sess.run(tf.global_variables_initializer())
z_samples = np.random.uniform(-1.0, 1.0, [batch_size, z_dim]).astype(np.float32)
samples = []
loss = {'d': [], 'g': []}

saver = tf.train.Saver()
offset = 0
for i in range(60000):
    n = np.random.uniform(-1.0, 1.0, [batch_size, z_dim]).astype(np.float32)
    offset = (offset + batch_size) % len(images)
    batch = np.array([read_image(img, HEIGHT, WIDTH) for img in images[offset: offset + batch_size]])
    batch = (batch - 0.5) * 2
    d_ls, g_ls = sess.run([loss_d, loss_g], feed_dict={X: batch, noise: n, is_training: True})
    loss['d'].append(d_ls)
    loss['g'].append(g_ls)
     #每优化一次discrimintor,优化两次generator
    sess.run(optimizer_d, feed_dict={X: batch, noise: n, is_training: True})
    sess.run(optimizer_g, feed_dict={X: batch, noise: n, is_training: True})
    sess.run(optimizer_g, feed_dict={X: batch, noise: n, is_training: True})
    print("now is iteratio {}".format(i))

    if i % 50 == 0:
        print("iteration {}, the d_ls is {}, and the g_ls is {}".format(i,d_ls,g_ls))
        #生成图片
        gen_imgs = sess.run(g, feed_dict={noise: z_samples, is_training: False})
        gen_imgs = (gen_imgs + 1) / 2
        imgs = [img[:, :, :] for img in gen_imgs]
        gen_imgs = montage(imgs)
        imsave(os.path.join(OUTPUT_DIR, 'sample_%d.jpg' % i), gen_imgs)
        samples.append(gen_imgs)

plt.plot(loss['d'], label='Discriminator')
plt.plot(loss['g'], label='Generator')
plt.legend(loc='upper right')
plt.savefig(os.path.join(OUTPUT_DIR, 'Loss.png'))
plt.show()
mimsave(os.path.join(OUTPUT_DIR, 'samples.gif'), samples, fps=10)
# save the data checkpoint
saver.save(sess, os.path.join(OUTPUT_DIR, 'dcgan_' + dataset), global_step=60000)


#!/usr/bin/env python3
import numpy as np

#we can define these centrality bins for all STAR observables
STAR_cent_bins = np.array( [ [0,5],[5,10],[10,20],[20,30],[30,40],[40,50],[50,60],[60,70], [70,80] ] ) # 9 bins
#more central bins to use for parameter estimation to avoid including empty events
central_STAR_cent_bins = np.array( [ [0,5],[5,10],[10,20],[20,30],[30,40],[40,50] ] ) # 6 bins
#more central bins for some PHENIX measurements
central_PHENIX_cent_bins = np.array( [ [0,5],[5,10],[10,15],[15,20],[20,30],[30,40],[40,50] ] ) # 7 bins
#these bins are common to many ALICE observables
ALICE_cent_bins = np.array( [ [0,5],[5,10],[10,20],[20,30],[30,40],[40,50],[50,60],[60,70] ] ) # 8 bins
#Tmunu_cents = ALICE_cent_bins

#the observables which will be used for parameter estimation
obs_cent_list = {

'Pb-Pb-2760': {
	'dNch_deta' : ALICE_cent_bins,
	'dET_deta' : np.array([[0, 2.5], [2.5, 5], [5, 7.5], [7.5, 10],
		                   [10, 12.5], [12.5, 15], [15, 17.5], [17.5, 20],
		                   [20, 22.5], [22.5, 25], [25, 27.5], [27.5, 30],
		                   [30, 32.5], [32.5, 35], [35, 37.5], [37.5, 40],
		                   [40, 45], [45, 50], [50, 55], [55, 60],
		                   [60, 65], [65, 70]]), # 22 bins
	'dN_dy_pion'   : ALICE_cent_bins,
	'dN_dy_kaon'   : ALICE_cent_bins,
	'dN_dy_proton' : ALICE_cent_bins,
	'dN_dy_Lambda' : np.array([[0,5],[5,10],[10,20],[20,40],[40,60]]), # 5 bins
	'dN_dy_Omega'  : np.array([[0,10],[10,20],[20,40],[40,60]]), # 4 bins
	'dN_dy_Xi'     : np.array([[0,10],[10,20],[20,40],[40,60]]), # 4 bins
	#'dN_dy_d'      : np.array([[0,10],[10,20],[20,40],[40,60]]), # 4 bins
	'mean_pT_pion'   : ALICE_cent_bins,
	'mean_pT_kaon'   : ALICE_cent_bins,
	'mean_pT_proton' : ALICE_cent_bins,
	#'mean_pT_d'      : np.array([[0,10],[10,20],[20,40],[40,60]]), # 4 bins
	'pT_fluct' : np.array([[0,5],[5,10],[10,15],[15,20], [20,25],[25,30],[30,35],[35,40], [40,45],[45,50],[50,55],[55,60]]), #12 bins
	'v22' : ALICE_cent_bins,
	'v32' : np.array([[0,5],[5,10],[10,20],[20,30], [30,40],[40,50]]), # 6 bins
	'v42' : np.array([[0,5],[5,10],[10,20],[20,30], [30,40],[40,50]]), # 6 bins

	#'Tmunu0' : Tmunu_cents,
    #'Tmunu1' : Tmunu_cents,
	#'Tmunu2' : Tmunu_cents,
	#'Tmunu3' : Tmunu_cents,
	#'Tmunu4' : Tmunu_cents,
	#'Tmunu5' : Tmunu_cents,
	#'Tmunu6' : Tmunu_cents,
	#'Tmunu7' : Tmunu_cents,
	#'Tmunu8' : Tmunu_cents,
	#'Tmunu9' : Tmunu_cents
    },

'Pb-Pb-5020': {
	'dNch_deta' : np.array( [ [0,2.5],[2.5,5],[5,7.5],[7.5,10],[10,20],[20,30],[30,40],[40,50],[50,60],[60,70] ] ),
	'dN_dy_pion'   : ALICE_cent_bins,
	'dN_dy_kaon'   : ALICE_cent_bins,
	'dN_dy_proton' : ALICE_cent_bins,
	#'dN_dy_d'      : np.array([[0,10],[10,20],[20,40],[40,60]]), # 4 bins
	'mean_pT_pion'   : ALICE_cent_bins,
	'mean_pT_kaon'   : ALICE_cent_bins,
	'mean_pT_proton' : ALICE_cent_bins,
	#'mean_pT_d'      : np.array([[0,10],[10,20],[20,40],[40,60]]), # 4 bins
	'v22' : ALICE_cent_bins,
	'v32' : np.array([[0,5],[5,10],[10,20],[20,30],[30,40],[40,50]]), # 6 bins
	'v42' : np.array([[0,5],[5,10],[10,20],[20,30],[30,40],[40,50]]), # 6 bins
    },

'Xe-Xe-5440': {
	'dNch_deta' : np.array( [ [0,2.5],[2.5,5],[5,7.5],[7.5,10],[10,20],[20,30],[30,40],[40,50],[50,60],[60,70] ] ),
	'v22' : ALICE_cent_bins,
	'v32' : ALICE_cent_bins,
    },

'Au-Au-200': {
	'dN_dy_pion'   : central_STAR_cent_bins,
	'dN_dy_kaon'   : central_STAR_cent_bins,
	#current calculations use STAR centrality bins
	#NOTE that the model calculations need to be re-averaged using the PHENIX cent bins if we want to include proton
	'dN_dy_proton' : central_STAR_cent_bins,
	'mean_pT_pion'   : central_STAR_cent_bins,
	'mean_pT_kaon'   : central_STAR_cent_bins,
	'mean_pT_proton' : central_STAR_cent_bins,
	'v22' : central_STAR_cent_bins,
	'v32' : central_STAR_cent_bins,
    },

}

#these just define some 'reasonable' ranges for plotting purposes
obs_range_list = {
    'Pb-Pb-2760': {
		'dNch_deta': [0,2000],
		'dET_deta': [0,2200],
		'dN_dy_pion': [0,1700],
		'dN_dy_kaon': [0,400],
		'dN_dy_proton': [0,120],
		'dN_dy_Lambda': [0,40],
		'dN_dy_Omega': [0,2],
		'dN_dy_Xi': [0,10],
		'mean_pT_pion': [0,1],
		'mean_pT_kaon': [0,1.5],
		'mean_pT_proton': [0,2],
		'pT_fluct': [0,0.05],
		'v22': [0,0.16],
		'v32': [0,0.1],
		'v42': [0,0.1]
    },
	'Au-Au-200': {
		'dNch_deta': [0,1000],
		'dET_deta': [0,1200],
		'dN_dy_pion': [0,800],
		'dN_dy_kaon': [0,120],
		'dN_dy_proton': [0,40],
		'dN_dy_Lambda': [0,40],
		'dN_dy_Omega': [0,2],
		'dN_dy_Xi': [0,10],
		'mean_pT_pion': [0,1],
		'mean_pT_kaon': [0,1.5],
		'mean_pT_proton': [0,2],
		'pT_fluct': [0,0.05],
		'v22': [0,0.16],
		'v32': [0,0.1],
		'v42': [0,0.1]
    },
}


"""
MIT License

Copyright (c) 2019 Simon Olofsson

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
"""

import numpy as np
import math
from scipy import special as sp

class StateConstraint:
    def __init__ (self, bounds, **kwargs):
        self.bounds = np.asarray( bounds )

        if self.bounds.ndim == 2:
            # Constant bounds
            assert self.bounds.shape[1] == 2
            self.num_states = self.bounds.shape[0]
        elif self.bounds.ndim == 3:
            # Time-dependent bounds
            assert self.bounds.shape[2] == 2
            self.num_steps  = self.bounds.shape[0]
            self.num_states = self.bounds.shape[1]

        if 'conf' in kwargs:
            # Constant bounds
            self.conf = kwargs['conf']

    @property
    def has_added_variables (self):
        return False

    def num_constraints (self):
        # Number of individual constraints constructed by class
        raise NotImplementedError

    def __call__ (self, M, S, step=None, grad=False):
        """
        Input:
        M   [ num_meas ]              Matrix of predictive means
        S   [ num_meas x num_meas ]   Matrix of predictive covariances
        grad                          Return gradients wrt M and S

        Output:
        c      Constraint score for M, S
        dcdM   Gradient d c / d M
        dcdS   Gradient d c / d S
        """
        raise NotImplementedError

class ConstantMeanStateConstraint (StateConstraint):
    r"""
    Mean constraint:
        bounds[i,0] <= \mu_i(t) <= bounds[i,1]
    """
    def __init__ (self, bounds):
        super().__init__ (bounds)
        assert self.bounds.ndim == 2

    def num_constraints (self):
        return 2 * self.num_states

    def __call__ (self, M, S, step=None, grad=False):
        C = np.zeros( 2 * self.num_states )
        #print( 'mean' )
        if grad:
            dCdM = np.zeros( C.shape + M.shape )
            dCdS = np.zeros( C.shape + S.shape )

        for i in range(self.num_states):
            C[2*i]   = M[i] - self.bounds[i,0]
            #print('lower: %d'%i)
            #print(self.bounds[i,0])
            C[2*i+1] = self.bounds[i,1] - M[i]
            #print('upper: %d'%i)
            #print(self.bounds[i,1])
            if grad:
                dCdM[2*i,  i] = 1.
                dCdM[2*i+1,i] = -1.
        return C if not grad else (C, dCdM, dCdS)

class MovingMeanStateConstraint (StateConstraint):
    r"""
    Mean constraint:
        bounds[t,i,0] <= \mu_i(t) <= bounds[t,i,1]
    """
    def __init__ (self, bounds):
        super().__init__ (bounds)
        assert self.bounds.ndim == 3

    def num_constraints (self):
        return 2 * self.num_states

    def __call__ (self, M, S, step, grad=False):
        step = np.min(( self.num_steps-1, step ))
        C    = np.zeros( 2 * self.num_states )
        if grad:
            dCdM = np.zeros( C.shape + M.shape )
            dCdS = np.zeros( C.shape + S.shape )

        for i in range(self.num_states):
            C[2*i]   = M[i] - self.bounds[step,i,0]
            C[2*i+1] = self.bounds[step,i,1] - M[i]
            if grad:
                dCdM[2*i,  i] = 1.
                dCdM[2*i+1,i] = -1.
        return C if not grad else (C, dCdM, dCdS)

class SingleChanceStateConstraint (StateConstraint):
    r"""
    Chance constraint:
        conf = probability
        r = distance from the mean value for a defined tolerance
        P(\mu_i(t)+r*S <= bounds[i,1]) > 1-conf
        P(\mu_i(t)-r*S >= bounds[i,1]) > 1-conf
    """

    def __init__ (self, bounds, **kwargs):
        super().__init__ (bounds,**kwargs)
        assert self.bounds.ndim == 2

    def num_constraints (self):
        return 2 * self.num_states

    def __call__ (self, M, S, step=None, grad=False):
        C = np.zeros( 2 * self.num_states )
        r = math.pow(2,0.5)*sp.erfinv(self.conf)
        if grad:
            dCdM = np.zeros( C.shape + M.shape )
            dCdS = np.zeros( C.shape + S.shape )

        for i in range(self.num_states):
            C[2*i]   = M[i] - r*np.sqrt(S[i,i]) - self.bounds[i,0];
            C[2*i+1] = self.bounds[i,1]- r*np.sqrt(S[i,i]) - M[i];

            if grad:
                dCdM[2*i,  i] = 1.
                dCdM[2*i+1,i] = -1.
                dCdS[2*i,  i] = -r/(2*np.sqrt(S[i,i]))
                dCdS[2*i+1,i] = -r/(2*np.sqrt(S[i,i]))
        return C if not grad else (C, dCdM, dCdS)

class JointTimeChanceStateConstraint (StateConstraint):
    r"""
    Chance constraint:
        P1(\mu_i(t)+r*S <= bounds[i,1]) > 1-eps
        P2(\mu_i(t)-r*S >= bounds[i,1]) > 1-eps
    """

    def __init__ (self, bounds,**kwargs):
        super().__init__ (bounds,**kwargs)
        assert self.bounds.ndim == 2

    def num_constraints (self):
        return 2 * self.num_states

    def __call__ (self, M, S, step=None, grad=False):
        C = np.zeros( 2 * self.num_states )

        if grad:
            dCdM = np.zeros( C.shape + M.shape )
            dCdS = np.zeros( C.shape + S.shape )

        for i in range(self.num_states):
            r1 = (M[i] - self.bounds[i,0])/np.sqrt(S[i,i])
            r2 = (self.bounds[i,1] - M[i])/np.sqrt(S[i,i])
            P1 = 1-(0.5+0.5*sp.erf(r1/(math.pow(2,0.5))))
            P2 = 1-(0.5+0.5*sp.erf(r2/(math.pow(2,0.5))))

            C[2*i]   = P1
            C[2*i+1] = P2

            if grad:
                dCdM[2*i,  i] = -math.exp(-math.pow(M[i]- self.bounds[i,0],2)/(2*S[i,i]))/(np.sqrt(2*math.pi*S[i,i]))
                dCdM[2*i+1,i] = -math.exp(-math.pow(self.bounds[i,1] - M[i],2)/(2*S[i,i]))/(np.sqrt(2*math.pi*S[i,i]))
                dCdS[2*i,  i] = (M[i]-self.bounds[i,0])*math.exp(-math.pow(M[i]- self.bounds[i,0],2)/(2*S[i,i]))/(np.sqrt(2*math.pi)*S[i,i])
                dCdS[2*i+1,i] = (self.bounds[i,1] - M[i])*math.exp(-math.pow(self.bounds[i,1] - M[i],2)/(2*S[i,i]))/(np.sqrt(2*math.pi)*S[i,i])
        return C if not grad else (C, dCdM, dCdS)