# Copyright (c) 2017, Lehrstuhl fuer Angewandte Mechanik, Technische
# Universitaet Muenchen.
#
# Distributed under BSD-3-Clause License. See LICENSE-File for more information
#
"""
Module handling the whole mechanical system, no matter if it's a finite element
system, defined by certain parameters or a multibody system.
"""

import time
import os
import copy

import h5py
import numpy as np

from .mesh import Mesh
from .assembly import Assembly
from .boundary import DirichletBoundary

__all__ = ['MechanicalSystem',
           'ReducedSystem',
           'reduce_mechanical_system'
           ]

class MechanicalSystem():
    '''
    Master class for mechanical systems with the goal to black-box the routines
    of assembly and element selection.

    Attributes
    ----------
    mesh_class : instance of Mesh()
        Class handling the mesh.
    assembly_class : instance of Assembly()
        Class handling the assembly.
    dirichlet_class : instance of DirichletBoundary
        Class handling the Dirichlet boundary conditions.
    T_output : list of floats
        List of timesteps saved.
    u_output : list of ndarrays
        List of unconstrained displacement arrays corresponding to the
        timesteps in T_output.
    S_output : list of ndarrays
        List of stress arrays corresponding to the timesteps in T_output.
    E_output : list of ndarrays
        List of strain arrays corresponding to the timesteps in T_output.
    stress : ndarray
        Array of nodal stress of the last assembly run. Shape is
        (no_of_nodes, 6).
    strain : ndarray
        Array of nodal strain of the last assembly run. Shape is
        (no_of_nodes, 6).
    stress_recovery : bool
        Flag for option stress_recovery.
    iteration_info : ndarray
        array containing the information of an iterative solution procedure.
        iteration_info[:,0] is the time information,
        iteration_info[:,1] is the number of iterations,
        iteration_info[:,3] is the residual.
    M_constr : ?
        Mass matrix
    D_constr : ?
        Damping matrix
    '''

    def __init__(self, stress_recovery=False):
        '''
        Parameters
        ----------
        stress_recovery : bool, optional
            Flag, for setting stress recovery option. Default is False.

        '''
        self.stress_recovery = stress_recovery
        self.T_output = []
        self.u_output = []
        self.S_output = []
        self.E_output = []
        self.stress = None
        self.strain = None
        self.iteration_info = np.array([])

        # instantiate the important classes needed for the system:
        self.mesh_class = Mesh()
        self.assembly_class = Assembly(self.mesh_class)
        self.dirichlet_class = DirichletBoundary(np.nan)

        # make syntax a little bit leaner
        # !Christian Meyer: ! careful: This prohibits to easily change dirichlet_class instance, because old instance
        # still will be referenced!
        self.unconstrain_vec = self.dirichlet_class.unconstrain_vec
        self.constrain_vec = self.dirichlet_class.constrain_vec
        self.constrain_matrix = self.dirichlet_class.constrain_matrix

        # initializations to be overwritten by loading functions
        self.M_constr = None
        self.D_constr = None
        self.no_of_dofs_per_node = None

        # external force to be overwritten by user-defined external forces
        # self._f_ext_unconstr = lambda t: np.zeros(self.mesh_class.no_of_dofs)

    def load_mesh_from_gmsh(self, msh_file, phys_group, material,
                            scale_factor=1):
        '''
        Load the mesh from a msh-file generated by gmsh.

        Parameters
        ----------
        msh_file : str
            file name to an existing .msh file
        phys_group : int
            integer key of the physical group which is considered as the
            mesh part
        material : amfe.Material
            Material associated with the physical group to be computed
        scale_factor : float, optional
            scale factor for the mesh to adjust the units. The default value is
            1, i.e. no scaling is done.

        Returns
        -------
        None
        '''
        self.mesh_class.import_msh(msh_file, scale_factor=scale_factor)
        self.mesh_class.load_group_to_mesh(phys_group, material)
        self.no_of_dofs_per_node = self.mesh_class.no_of_dofs_per_node

        self.assembly_class.preallocate_csr()
        self.dirichlet_class.no_of_unconstrained_dofs = self.mesh_class.no_of_dofs
        self.dirichlet_class.update()

    def deflate_mesh(self):
        '''
        Remove free floating nodes not connected to a selected element from
        the mesh.

        Parameters
        ----------
        None

        Returns
        -------
        None
        '''
        self.mesh_class.deflate_mesh()
        self.assembly_class.preallocate_csr()
        self.dirichlet_class.no_of_unconstrained_dofs = self.mesh_class.no_of_dofs
        self.dirichlet_class.update()

    def load_mesh_from_csv(self, node_list_csv, element_list_csv,
                           no_of_dofs_per_node=2,
                           explicit_node_numbering=False,
                           ele_type=False):
        '''
        Loads the mesh from two csv-files containing the node and the element list.

        Parameters
        ----------
        node_list_csv: str
            filename of the csv-file containing the coordinates of the nodes (x, y, z)
        element_list_csv: str
            filename of the csv-file containing the nodes which belong to one element
        no_of_dofs_per_node: int, optional
            degree of freedom per node as saved in the csv-file
        explicit_node_numbering : bool, optional
            flag stating, if the node numbers are explcitly numbered in the
            csv file, i.e. if the first column gives the numbers of the nodes.
        ele_type: str
            Spezifiy elements type of the mesh (e.g. for a Tri-Mesh different
            elements types as Tri3, Tri4, Tri6 can be used)
            If not spezified value is set to 'False'

        Returns
        -------
        None

        Examples
        --------
        todo

        '''
        self.mesh_class.import_csv(node_list_csv, element_list_csv,
                                   explicit_node_numbering=explicit_node_numbering,
                                   ele_type=ele_type)
        self.no_of_dofs_per_node = no_of_dofs_per_node
        self.assembly_class.preallocate_csr()
        return

    def tie_mesh(self, master_key, slave_key, master_prop='phys_group',
                 slave_prop='phys_group', tying_type='fixed', verbose=False,
                 conform_slave_mesh=False, fix_mesh_dist=1E-3):
        '''
        Tie nonconforming meshes for a given master and slave side.

        Parameters
        ----------
        master_key : int or string
            mesh key of the master face mesh. The master face mesh has to be at
            least the size of the slave mesh. It is better, when the master
            mesh is larger than the slave mesh.
        slave_key : int or string
            mesh key of the slave face mesh or point cloud
        master_prop : string, optional
            mesh property for which master_key is specified.
            Default value: 'phys_group'
        slave_prop : string, optional
            mesh property for which slave_key is specified.
            Default value: 'phys_group'
        tying_type : string {'fixed', 'slide'}
            Mesh tying type. 'fixed' glues the meshes together while 'slide'
            allows for a sliding motion between the meshes.

        Returns
        -------
        None

        Notes
        -----
        The master mesh has to embrace the full slave mesh. If this is not the
        case, the routine will fail, a slave point outside the master mesh
        cannot be addressed to a specific element.

        '''

        vals = self.mesh_class.tie_mesh(master_key=master_key,
                                        slave_key=slave_key,
                                        master_prop=master_prop,
                                        slave_prop=slave_prop,
                                        tying_type=tying_type,
                                        verbose=verbose,
                                        fix_mesh_dist=fix_mesh_dist)

        self.dirichlet_class.add_constraints(*vals)
        self.dirichlet_class.update()
        return

    def apply_dirichlet_boundaries(self, key, coord, mesh_prop='phys_group'):
        '''
        Apply dirichlet-boundaries to the system.

        Parameters
        ----------
        key : int
            Key for mesh property which is to be chosen. Matches the group given
            in the gmsh file. For help, the function mesh_information or
            boundary_information gives the groups
        coord : str {'x', 'y', 'z', 'xy', 'xz', 'yz', 'xyz'}
            coordinates which should be fixed
        mesh_prop : str {'phys_group', 'geom_entity', 'el_type'}, optional
            label of which the element should be chosen from. Default is
            'phys_group'.

        Returns
        -------
        None
        '''
        self.mesh_class.set_dirichlet_bc(key, coord, mesh_prop)
        self.dirichlet_class.constrain_dofs(self.mesh_class.dofs_dirichlet)
        return

    def apply_neumann_boundaries(self, key, val, direct, time_func=None,
                                 shadow_area=False, mesh_prop='phys_group'):
        '''
        Apply neumann boundaries to the system via skin elements.

        Parameters
        ----------
        key : int
            Key of the physical domain to be chosen for the neumann bc
        val : float
            value for the pressure/traction onto the element
        direct : ndarray or str 'normal'
            Direction, in which force should act at. If
        time_func : function object
            Function object returning a value between -1 and 1 given the
            input t:

            >>> val = time_func(t)

        shadow_area : bool, optional
            flag, if force should be proportional to shadow area of surface
            with respect to direction. Default: False.
        mesh_prop : str {'phys_group', 'geom_entity', 'el_type'}, optional
            label of which the element should be chosen from. Default is
            phys_group.

        Returns
        -------
        None
        '''
        self.mesh_class.set_neumann_bc(key=key, val=val, direct=direct,
                                       time_func=time_func,
                                       shadow_area=shadow_area,
                                       mesh_prop=mesh_prop)
        self.assembly_class.compute_element_indices()
        return

    def export_paraview(self, filename, field_list=None):
        '''
        Export the system with the given information to paraview.

        Parameters
        ----------
        filename : str
            filename to which the xdmf file and the hdf5 file will be saved.
        field_list : list, optional
            list of tuples containing a field to be exported as well as a
            dictionary with the attribute information of the hdf5 file.
            Example:

                >>> # example field list with reduced displacement not to export
                >>> # ParaView and strain epsilon to be exported to ParaView
                >>> field_list = [(q_red, {'ParaView':False, 'Name':'q_red'}),
                                  (eps, {'ParaView':True,
                                         'Name':'epsilon',
                                         'AttributeType':'Tensor6',
                                         'Center':'Node',
                                         'NoOfComponents':6})]

        Returns
        -------
        None
        '''
        if field_list is None:
            field_list = []
        t1 = time.time()
        if len(self.T_output) is 0:
            self.T_output.append(0)
            self.u_output.append(np.zeros(self.mesh_class.no_of_dofs))
            if self.stress_recovery:
                self.S_output.append(np.zeros((self.mesh_class.no_of_nodes, 6)))
                self.E_output.append(np.zeros((self.mesh_class.no_of_nodes, 6)))
        print('Start exporting mesh for paraview to:\n    ', filename)

        if self.stress_recovery and len(self.S_output) > 0 \
           and len(self.E_output) > 0:
            no_of_timesteps = len(self.T_output)
            S_array = np.array(self.S_output).reshape((no_of_timesteps, -1))
            E_array = np.array(self.E_output).reshape((no_of_timesteps, -1))
            S_export = (S_array.T, {'ParaView':True,
                                  'Name':'stress',
                                  'AttributeType':'Tensor6',
                                  'Center':'Node',
                                  'NoOfComponents':6})
            E_export = (E_array.T, {'ParaView':True,
                                  'Name':'strain',
                                  'AttributeType':'Tensor6',
                                  'Center':'Node',
                                  'NoOfComponents':6})
            field_list.append(S_export)
            field_list.append(E_export)

        bmat = self.dirichlet_class.b_matrix()
        self.mesh_class.save_mesh_xdmf(filename, field_list, bmat, u=self.u_output, timesteps=self.T_output)
        t2 = time.time()
        print('Mesh for paraview successfully exported in ' +
              '{0:4.2f} seconds.'.format(t2 - t1))
        return

    def M(self, u=None, t=0):
        '''
        Compute the Mass matrix of the dynamical system.

        Parameters
        ----------
        u : ndarray, optional
            array of the displacement
        t : float
            time

        Returns
        -------
        M : sp.sparse.sparse_matrix
            Mass matrix with applied constraints in sparse csr-format
        '''
        if u is not None:
            u_unconstr = self.unconstrain_vec(u)
        else:
            u_unconstr = None

        M_unconstr = self.assembly_class.assemble_m(u_unconstr, t)
        self.M_constr = self.constrain_matrix(M_unconstr)
        return self.M_constr

    def K(self, u=None, t=0):
        '''
        Compute the stiffness matrix of the mechanical system

        Parameters
        ----------
        u : ndarray, optional
            Displacement field in voigt notation
        t : float, optional
            Time

        Returns
        -------
        K : sp.sparse.sparse_matrix
            Stiffness matrix with applied constraints in sparse csr-format
        '''
        if u is None:
            u = np.zeros(self.dirichlet_class.no_of_constrained_dofs)

        K_unconstr = \
            self.assembly_class.assemble_k_and_f(self.unconstrain_vec(u), t)[0]

        return self.constrain_matrix(K_unconstr)

    def D(self, u=None, t=0):
        '''
        Return the damping matrix of the mechanical system

        Parameters
        ----------
        u : ndarray, optional
            Displacement field in voigt notation
        t : float, optional
            Time

        Returns
        -------
        D : sp.sparse.sparse_matrix
            Damping matrix with applied constraints in sparse csr-format
        '''
        if self.D_constr is None:
            return self.K()*0
        else:
            return self.D_constr

    def f_int(self, u, t=0):
        '''Return the elastic restoring force of the system '''
        f_unconstr = \
            self.assembly_class.assemble_k_and_f(self.unconstrain_vec(u), t)[1]
        return self.constrain_vec(f_unconstr)

    def _f_ext_unconstr(self, u, t):
        '''
        Return the unconstrained external force coming from the Neumann BCs.

        This function is just a placeholder if you want to change the behavior of f_ext:
        This function may be monkeypatched if necessary, for instance, when a
        global external force, e.g. gravity, should be applied.
        '''
        f_unconstr = \
            self.assembly_class.assemble_k_and_f_neumann(self.unconstrain_vec(u), t)[1]
        return f_unconstr

    def f_ext(self, u, du, t):
        '''
        Return the nonlinear external force of the right hand side
        of the equation, i.e. the excitation.
        '''
        if u is None:
            u = np.zeros(self.dirichlet_class.no_of_constrained_dofs)
        return self.constrain_vec(self._f_ext_unconstr(u, t))

    def K_and_f(self, u=None, t=0):
        '''
        Compute tangential stiffness matrix and nonlinear force vector
        in one assembly run.
        '''
        if u is None:
            u = np.zeros(self.dirichlet_class.no_of_constrained_dofs)
        if self.stress_recovery: # make sure, that current stress / strain is exported
            K_unconstr, f_unconstr, self.stress, self.strain = \
                self.assembly_class.assemble_k_f_S_E(self.unconstrain_vec(u), t)
        else:
            K_unconstr, f_unconstr = \
                self.assembly_class.assemble_k_and_f(self.unconstrain_vec(u), t)
        K = self.constrain_matrix(K_unconstr)
        f = self.constrain_vec(f_unconstr)
        return K, f

    def S_and_res(self, u, du, ddu, dt, t, beta, gamma):
        r'''
        Compute jacobian and residual for implicit Newmark time integration.

        Parameters
        ----------
        u : ndarray
            displacement; dimension (ndof,)
        du : ndarray
            velocity; dimension (ndof,)
        ddu : ndarray
            acceleration; dimension (ndof,)
        dt : float
            time step width
        t : float
            time of current time step (for time dependent loads)
        beta : float
            weighting factor for position in generalized-:math:`\alpha` scheme
        gamma : float
            weighting factor for velocity in generalized-:math:`\alpha` scheme

        Returns
        -------
        S : ndarray
            jacobian matrix of residual; dimension (ndof, ndof)
        res : ndarray
            residual; dimension (ndof,)

        Notes
        -----
        Time integration scheme: The iteration matrix is composed using the
        Newmark scheme:

        .. math:: \mathbf S \Delta u = -\mathbf{res}

        .. math:: \mathbf S = K + \frac{\gamma}{\beta h} \mathbf{D}
                    + \frac{1}{\beta h^2} \mathbf{M}

        which bases on the time discretization of the velocity and the
        displacement:

        .. math:: \mathbf{\dot{q}}_{n+1} & = \mathbf{\dot{q}}_{n} + (1-\gamma)h
                  \mathbf{\ddot{q}}_{n} + \gamma h \mathbf{\ddot{q}}_{n+1}

        #.. math:: \mathbf{q}_{n+1} & = \mathbf{q}_n + h \mathbf{\dot{q}}_n +
        #          \left(\frac{1}{2} - \beta\right)h^2\mathbf{\ddot{q}}_n +
        #          h^2\beta\mathbf{\ddot{q}}_{n+1}
        .. math:: \mathbf{\ddot{q}}_{n+1} &= \frac{1}{\beta h^2} \left(
                    \mathbf{\Delta q} - h \mathbf{q}_n - (\frac{1}{2} - \beta) h^2 \mathbf{\ddot q}_n \right)

        This method is using the variables/methods

            - self.M()
            - self.M_constr
            - self.K_and_f()
            - self.f_ext()

        If these methods are implemented correctly in a daughter class, the
        time integration interface should work properly.

        '''
        # compute mass matrix only once if it hasnt's been computed yet
        if self.M_constr is None:
            self.M()

        K, f = self.K_and_f(u, t)
        f_ext = self.f_ext(u, du, t)
        if self.D_constr is None:
            S = K + 1/(beta*dt**2)*self.M_constr
            res = f - f_ext + self.M_constr @ ddu
        else: # damping
            S = K \
                + gamma/(beta*dt) * self.D_constr \
                + 1/(beta*dt**2) * self.M_constr
            res = f - f_ext + self.M_constr @ ddu + self.D_constr @ du
        return S, res, f_ext

    def gen_alpha(self, q, dq, ddq, q_old, dq_old, ddq_old,
                  f_ext_old, dt, t, alpha_m, alpha_f, beta, gamma):
        '''
        Computation of Jacobian and residual for generalized-alpha time
        integration scheme.

        TODO

        '''
        # compute mass matrix only if it has not been computed yet
        if self.M_constr is None:
            self.M()

        ddq_m = (1-alpha_m)*ddq + alpha_m*ddq_old
        dq_f = (1-alpha_f)*dq + alpha_f*dq_old
        q_f = (1-alpha_f)*q + alpha_f*q_old

        K_f, f_f = self.K_and_f(q_f, t)
        f_ext = self.f_ext(q, dq, t)
        f_ext_f = (1-alpha_f) * f_ext + alpha_f * f_ext_old

        if self.D_constr is None:
            Jac = (1-alpha_f) * K_f + (1-alpha_m)/(beta*dt**2) * self.M_constr
            res = f_f - f_ext_f + self.M_constr @ ddq_m
        else: # damping
            Jac =   (1-alpha_f) * K_f \
                  + (1-alpha_f)*gamma/(beta*dt) * self.D_constr \
                  + (1-alpha_m)/(beta*dt**2) * self.M_constr
            res = f_f - f_ext_f + self.D_constr @ dq_f + self.M_constr @ ddq_m

        return Jac, res, f_ext

    def apply_rayleigh_damping(self, alpha, beta):
        '''
        Apply Rayleigh damping to the system.

        The damping matrix D is defined as

        D = alpha*M + beta*K(0)

        Thus, it is Rayleigh Damping applied to the linearized system
        around zero deformation.

        Parameters
        ----------
        alpha : float
            damping coefficient for the mass matrix
        beta : float
            damping coefficient for the stiffness matrix

        '''
        if self.M_constr is None:
            self.M()
        self.D_constr = alpha*self.M_constr + beta*self.K()
        return

    def write_timestep(self, t, u):
        '''
        write the timestep to the mechanical_system class
        '''
        self.T_output.append(t)
        self.u_output.append(self.unconstrain_vec(u))
        # Check both, if stress recovery and if stress and strain is there
        if self.stress_recovery:
            # catch the case when no stress was computed, for instance in time
            # integration
            if self.stress is None and self.strain is None:
                self.stress = np.zeros((self.mesh_class.no_of_nodes,6))
                self.strain = np.zeros((self.mesh_class.no_of_nodes,6))
            self.S_output.append(self.stress.copy())
            self.E_output.append(self.strain.copy())

    def clear_timesteps(self):
        '''
        Clear the timesteps gathered internally
        '''
        self.T_output = []
        self.u_output = []
        self.S_output = []
        self.E_output = []
        self.stress = None
        self.strain = None
        self.iteration_info = np.array([])
        return

class ReducedSystem(MechanicalSystem):
    '''
    Class for reduced systems.
    It is directly inherited from MechanicalSystem.
    Provides the interface for an integration scheme and so on where a basis
    vector is to be chosen...

    Notes
    -----
    The Basis V is a Matrix with x = V*q mapping the reduced set of coordinates
    q onto the physical coordinates x. The coordinates x are constrained, i.e.
    the x denotes the full system in the sense of the problem set and not of
    the pure finite element set.

    The system runs without providing a V_basis when constructing the method
    only for the unreduced routines.

    Attributes
    ----------
    V : ?
        Set of basis vectors the system has been reduced with u_constr = V * q
    V_unconstr : ?
        Extended reduction basis that is extended by the displacement coordinates of the constrained degrees of
        freedom
    u_red_output : ?
        Stores the timeseries of the generalized coordinates (similar to u_output)
    assembly_type : {'indirect', 'direct'}
        Stores the type of assembly method how the reduced system is computed

    Examples
    --------

    my_system = amfe.MechanicalSystem()
    V = vibration_modes(my_system, n=20)
    my_reduced_system = amfe.reduce_mechanical_system(my_system, V)

    '''

    def __init__(self, V_basis=None, assembly='indirect', **kwargs):
        '''
        Parameters
        ----------
        V_basis : ndarray, optional
            Basis onto which the problem will be projected with an
            Galerkin-Projection.
        assembly : str {'direct', 'indirect'}
            flag setting, if direct or indirect assembly is done. For larger
            reduction bases, the indirect method is much faster.
        **kwargs : dict, optional
            Keyword arguments to be passed to the mother class MechanicalSystem.

        Returns
        -------
        None
        '''
        MechanicalSystem.__init__(self, **kwargs)
        self.V = V_basis
        self.u_red_output = []
        self.V_unconstr = self.dirichlet_class.unconstrain_vec(V_basis)
        self.assembly_type = assembly

    def K_and_f(self, u=None, t=0):
        if u is None:
            u = np.zeros(self.V.shape[1])
        if self.assembly_type == 'direct':
            # this is really slow! So this is why the assembly is done diretly
            K, f_int = self.assembly_class.assemble_k_and_f_red(self.V_unconstr,
                                                                u, t)
        elif self.assembly_type == 'indirect':
            K_raw, f_raw = self.assembly_class.assemble_k_and_f(self.V_unconstr @ u,
                                                                t)
            K = self.V_unconstr.T @ K_raw @ self.V_unconstr
            f_int = self.V_unconstr.T @ f_raw
        else:
            raise ValueError('The given assembly type for a reduced system '
                             + 'is not valid.')
        return K, f_int

    def K(self, u=None, t=0):
        if u is None:
            u = np.zeros(self.V.shape[1])

        if self.assembly_type == 'direct':
            # this is really slow! So this is why the assembly is done diretly
            K, f_int = self.assembly_class.assemble_k_and_f_red(self.V_unconstr,
                                                                u, t)
        elif self.assembly_type == 'indirect':
            K_raw, f_raw = self.assembly_class.assemble_k_and_f(self.V_unconstr @ u,
                                                                t)
            K = self.V_unconstr.T @ K_raw @ self.V_unconstr
        else:
            raise ValueError('The given assembly type for a reduced system '
                             + 'is not valid.')
        return K

    def f_ext(self, u, du, t):
        return self.V.T @ MechanicalSystem.f_ext(self, self.V @ u, du, t)

    def f_int(self, u, t=0):

        if self.assembly_type == 'direct':
            # this is really slow! So this is why the assembly is done diretly
            K, f_int = self.assembly_class.assemble_k_and_f_red(self.V_unconstr,
                                                                u, t)
        elif self.assembly_type == 'indirect':
            K_raw, f_raw = self.assembly_class.assemble_k_and_f(self.V_unconstr @ u,
                                                                t)
            f_int = self.V_unconstr.T @ f_raw
        else:
            raise ValueError('The given assembly type for a reduced system '
                             + 'is not valid.')

        return f_int

    def D(self, u=None, t=0):

        if self.assembly_type == 'direct':
            raise NotImplementedError('The direct method is note implemented yet for damping matrices')
        elif self.assembly_type == 'indirect':
            self.D_constr = self.V.T @ MechanicalSystem.D(self, self.V @ u, t) @ self.V
        else:
            raise ValueError('The given assembly type for a reduced system '
                             + 'is not valid.')

        return self.D_constr

    def M(self, u=None, t=0):
        # Just a plain projection
        # not so well but works...
        self.M_constr = self.V.T @ MechanicalSystem.M(self, u, t) @ self.V
        return self.M_constr

    def write_timestep(self, t, u):
        MechanicalSystem.write_timestep(self, t, self.V @ u)
        self.u_red_output.append(u.copy())

    def K_unreduced(self, u=None, t=0):
        '''
        Unreduced Stiffness Matrix.

        Parameters
        ----------
        u : ndarray, optional
            Displacement of constrained system. Default is zero vector.
        t : float, optionial
            Time. Default is 0.

        Returns
        -------
        K : sparse csr matrix
            Stiffness matrix

        '''
        return MechanicalSystem.K(self, u, t)

    def f_int_unreduced(self, u, t=0):
        '''
        Internal nonlinear force of the unreduced system.

        Parameters
        ----------
        u : ndarray
            displacement of unreduces system.
        t : float, optional
            time, default value: 0.

        Returns
        -------
        f_nl : ndarray
            nonlinear force of unreduced system.

        '''
        return MechanicalSystem.f_int(self, u, t)

    def M_unreduced(self):
        '''
        Unreduced mass matrix.
        '''
        return MechanicalSystem.M(self)

    def export_paraview(self, filename, field_list=None):
        '''
        Export the produced results to ParaView via XDMF format.
        '''
        u_red_export = np.array(self.u_red_output).T
        u_red_dict = {'ParaView':'False', 'Name':'q_red'}

        if field_list is None:
            new_field_list = []
        else:
            new_field_list = field_list.copy()

        new_field_list.append((u_red_export, u_red_dict))

        MechanicalSystem.export_paraview(self, filename, new_field_list)

        # add V and Theta to the hdf5 file
        filename_no_ext, _ = os.path.splitext(filename)
        with h5py.File(filename_no_ext + '.hdf5', 'r+') as f:
            f.create_dataset('reduction/V', data=self.V)

        return

    def clear_timesteps(self):
        MechanicalSystem.clear_timesteps(self)
        self.u_red_output = []

def reduce_mechanical_system(mechanical_system, V, overwrite=False,
                             assembly='indirect'):
    '''
    Reduce the given mechanical system with the linear basis V.

    Parameters
    ----------
    mechanical_system : instance of MechanicalSystem
        Mechanical system which will be transformed to a ReducedSystem.
    V : ndarray
        Reduction Basis for the reduced system
    overwrite : bool, optional
        switch, if mechanical system should be overwritten (is less memory
        intensive for large systems) or not.
    assembly : str {'direct', 'indirect'}
            flag setting, if direct or indirect assembly is done. For larger
            reduction bases, the indirect method is much faster.

    Returns
    -------
    reduced_system : instance of ReducedSystem
        Reduced system with same properties of the mechanical system and
        reduction basis V

    Example
    -------

    '''

    if overwrite:
        reduced_sys = mechanical_system
    else:
        reduced_sys = copy.deepcopy(mechanical_system)
    reduced_sys.__class__ = ReducedSystem
    reduced_sys.V = V.copy()
    reduced_sys.V_unconstr = reduced_sys.dirichlet_class.unconstrain_vec(V)
    reduced_sys.u_red_output = []
    reduced_sys.M_constr = None
    # reduce Rayleigh damping matrix
    if reduced_sys.D_constr is not None:
        reduced_sys.D_constr = V.T @ reduced_sys.D_constr @ V
    reduced_sys.assembly_type = assembly
    return reduced_sys

# This class is not integrated in AMfe yet.
# It should apply linear combinations of external forces
# f_ext = B(x) * F(t)
# This function is not integrated in the Neumann Boundary conditions.
#
# class ExternalForce:
#     '''
#     Class for mimicking the external forces based on a force basis and time
#     values. The force values are linearly interpolated.
#
#     '''
#     def __init__(self, force_basis, force_series, t_series):
#         '''
#         Parameters
#         ----------
#         force_basis : ndarray, shape(ndim, n_dofs)
#             force basis for the force series
#         force_seris : ndarray, shape(n_timesteps, n_dofs)
#             array containing the force dofs corresponding to the time values
#             given in t_series
#         t_series : ndarray, shape(n_timesteps)
#             array containing the time values
#
#         '''
#         self.force_basis = force_basis
#         self.force_series= force_series
#         self.T = t_series
#         return
#
#     def f_ext(self, u, du, t):
#         '''
#         Mimicked external force for the given force time series
#         '''
#         # Catch the case that t is larger than the data set
#         if t >= self.T[-1]:
#             return self.force_basis @ self.force_series[-1]
#
#         t2_idx = np.where(self.T > t)[0][0]
#
#         # if t is smaller than lowest value of T, pick the first value in the
#         # force series
#         if t2_idx == 0:
#             force_amplitudes = self.force_series[0]
#         else:
#             t1_idx = t2_idx - 1
#             t2 = self.T[t2_idx]
#             t1 = self.T[t1_idx]
#
#             force_amplitudes = ( (t2-t)*self.force_series[t1_idx]
#                                + (t-t1)*self.force_series[t2_idx]) / (t2-t1)
#         return self.force_basis @ force_amplitudes


import numpy as np

# -----------------------------------------------------------------------------------------
#
#

class Ion:
    def __init__(self, protons, neutrons, electrons, energy, x=0., dx=0., y=0., dy=0., dl=0., dp=0.):
        self.protons = protons
        self.neutrons = neutrons
        self.electrons = electrons
        self.energy = energy
        self.x = x
        self.dx = dx
        self.y = y
        self.dy = dy
        self.dl = dl
        self.dp = dp

    def __repr__(self):
        r = "Ion(mass=" + str(self.mass) + "u, "
        r += "chargeState=" + str(self.charge) + "e, "
        r += "energy=" + str(self.energy) + "eV, "
        r += "x=" + str(self.x) + "m, "
        r += "dx=" + str(self.dx) + "rad, "
        r += "y=" + str(self.y) + "m, "
        r += "dy=" + str(self.dy) + "rad, "
        r += "dl=" + str(self.dl) + "m, "
        r += "dp=" + str(self.dp) + ")"
        return r

    def __str__(self):
        if self.charge>0:
            return str(self.mass_number) + self.nomenclature + "+" + str(self.charge)
        else:
            return str(self.mass_number) + self.nomenclature + str(self.charge)

    @property
    def charge(self):
        return self.protons-self.electrons

    @property
    def mass_number(self):
        return self.protons+self.neutrons

    @property
    def nomenclature(self):
        symbols = ["H", "He", "Li", "Be", "B", "C", "N", "O", "F", "Ne", "Na",
         "Mg", "Al", "Si", "P", "S", "Cl", "Ar", "K", "Ca", "Sc", "Ti",
         "V", "Cr", "Mn", "Fe", "Co", "Ni", "Cu", "Zn", "Ga", "Ge", "As",
         "Se", "Br", "Kr", "Rb", "Sr", "Y", "Zr", "Nb", "Mo", "Tc", "Ru",
         "Rh", "Pd", "Ag", "Cd", "In", "Sn", "Sb", "Te", "I", "Xe", "Cs",
         "Ba", "La", "Ce", "Pr", "Nd", "Pm", "Sm", "Eu", "Gd", "Tb", "Dy",
         "Ho", "Er", "Tm", "Yb", "Lu", "Hf", "Ta", "W", "Re", "Os", "Ir",
         "Pt", "Au", "Hg", "TI", "Pb", "Bi", "Po", "At", "Rn", "Fr", "Ra",
         "Ac", "Th", "Pa", "U", "Np", "Pu", "Am", "Cm", "Bk", "Cf", "Es",
         "Fm", "Md", "No", "Lr", "Rf", "Db", "Sg", "Bh", "Hs", "Mt", "Ds",
         "Rg", "Cn", "Uut", "Fl", "Uup", "Lv"]
        if self.protons > 117:
            return "??"
        elif self.protons < 1:
            raise Exception("Number of protons < 1")
        else:
            return symbols[self.protons-1]


import numpy as np
import math
import ctypes
import util.utilities as util
import gadget_lib.gadget as gadget

def checklen(x):
    return len(np.array(x,ndmin=1));

def int_round(x):
    return np.int(np.round(x));

def ok_scan(input,xmax=1.0e10,pos=0):
    if (pos==1):
        return (np.isnan(input)==False) & (np.fabs(input)<=xmax) & (input > 0.);
    else:
        return (np.isnan(input)==False) & (np.fabs(input)<=xmax);

def fcor(x):
    return np.array(x,dtype='f',ndmin=1)
def vfloat(x):
    return x.ctypes.data_as(ctypes.POINTER(ctypes.c_float));

def calculate_zoom_center(sdir,snum,cen=[0.,0.,0.],clip_size=2.e10):
    rgrid=np.array([1.0e10,1000.,700.,500.,300.,200.,100.,70.,50.,30.,20.,10.,5.,2.5,1.]);
    rgrid=rgrid[rgrid <= clip_size];
    Ps=gadget.readsnap(sdir,snum,4,cosmological=1);
    n_new=Ps['m'].shape[0];
    if (n_new > 1):
        pos=Ps['p']; x0s=pos[:,0]; y0s=pos[:,1]; z0s=pos[:,2];
    Pg=gadget.readsnap(sdir,snum,0,cosmological=1);
    rho=np.array(Pg['rho'])*407.5;
    if (rho.shape[0] > 0):
        pos=Pg['p']; x0g=pos[:,0]; y0g=pos[:,1]; z0g=pos[:,2];
    rho_cut=1.0e-5;
    cen=np.array(cen);

    for i_rcut in range(len(rgrid)):
        for j_looper in range(5):
            if (n_new > 1000):
                x=x0s; y=y0s; z=z0s;
            else:
                ok=(rho > rho_cut);
                x=x0g[ok]; y=y0g[ok]; z=z0g[ok];
            x=x-cen[0]; y=y-cen[1]; z=z-cen[2];
            r = np.sqrt(x*x + y*y + z*z);
            ok = (r < rgrid[i_rcut]);
            if (len(r[ok]) > 1000):
                x=x[ok]; y=y[ok]; z=z[ok];
                if (i_rcut <= len(rgrid)-5):
                    cen+=np.array([np.median(x),np.median(y),np.median(z)]);
                else:
                    cen+=np.array([np.mean(x),np.mean(y),np.mean(z)]);
            else:
                if (len(r[ok]) > 200):
                    x=x[ok]; y=y[ok]; z=z[ok];
                    cen+=np.array([np.mean(x),np.mean(y),np.mean(z)]);

    return cen;

def test():
    import gadget
    sdir='../zooms/m10_dw/'
    snum=440
    cen=calculate_zoom_center(sdir,snum);
    cen=np.array([ 3306.55575601,  3029.35099489,  3235.4757624 ]);
    print cen;

    PPPs=gadget.readsnap(sdir,snum,4,cosmological=1);
    PPP=gadget.readsnap(sdir,snum,0,cosmological=1);
    source_pos=(np.random.rand(3,5)-0.5)*100.
    source_pos[2,:] *= 0.02
    for i in [0,1,2]: source_pos[i,:]+=cen[i];
    source_pos=PPPs['p']

    mm=PPP['m'];
    rho=PPP['rho'];

    #mm *= PPP['nh'];
    #rho *= PPP['nh'];
    #print np.min(PPP['nh']), np.median(PPP['nh']), np.max(PPP['nh'])

    xlen = 10.
    los_nh,los_nh_hot,los_z = \
      return_columns_to_sources(source_pos,\
      PPP['p'],PPP['u'],rho,PPP['h'],PPP['ne'],PPP['nh'],PPP['z'],mm,\
      xrange=cen[0]+[-xlen,xlen],yrange=cen[1]+[-xlen,xlen],zrange=cen[2]+[-xlen,xlen])

    print los_nh
    print los_nh_hot
    print los_z

##
## return: los_NH_allgas, los_NH_hotphase, los_gas_metallicity
##
def return_columns_to_sources( source_pos, gas_pos, \
    gas_u, gas_rho, gas_hsml, gas_numh, gas_nume, gas_metallicity, gas_mass, \
    xrange=0, yrange=0, zrange=0, \
    MIN_CELL_SIZE=0.01, OUTER_RANGE_OF_INT=1200., \
    TRIM_PARTICLES=1 ):

    ## check the ordering of the position matrices:
    if ((checklen(gas_pos[0,:])==3) & (checklen(gas_pos[:,0]) !=3)): gas_pos=np.transpose(gas_pos);
    if ((checklen(source_pos[0,:])==3) & (checklen(source_pos[:,0]) !=3)): source_pos=np.transpose(source_pos);
    ## and that metallicities are a vector, not a matrix
    if (len(gas_metallicity.shape)>1): gas_metallicity=gas_metallicity[:,0]

    if ((checklen(gas_pos[:,0]) != 3) | (checklen(gas_pos[0,:]) <= 1)):
        print 'ERROR WILL OCCUR :: need pos to be (3,N)'

    x=source_pos[0,:] ; y=source_pos[1,:] ; z=source_pos[2,:]
    if(checklen(xrange)<=1): xrange=[np.min(x),np.max(x)];
    if(checklen(yrange)<=1): yrange=[np.min(y),np.max(y)];
    xr=xrange; yr=yrange;
    if(checklen(zrange)<=1):
        zrr=np.sqrt((xr[1]-xr[0])**2.+(yr[1]-yr[0])**2.)/np.sqrt(2.);
        zmin=np.median(z)-zrr; zmax=np.median(z)+zrr;
        if (np.min(z) > zmin): zmin=np.min(z);
        zrange=[zmin,zmax]; print 'z_range (calc) == ',zrange
    zr=zrange;
    x00=0.5*(xr[1]+xr[0]); y00=0.5*(yr[1]+yr[0]); z00=0.5*(zr[1]+zr[0]);
    tolfac = 1.0e10;
    if (TRIM_PARTICLES==1):
        tolfac = 0.05;
        #tolfac = -0.01;
        ## trim down the incoming list to only whats in the range plotted
        ##   (saves a ton of time and memory overflow crashes)

    dx=(0.5+tolfac)*(xr[1]-xr[0]); dy=(0.5+tolfac)*(yr[1]-yr[0]); dz=(0.5+tolfac)*(zr[1]-zr[0]);
    ok_sources=ok_scan(x-x00,xmax=dx) & ok_scan(y-y00,xmax=dy) & ok_scan(z-z00,xmax=dz);
    x=gas_pos[0,:] ; y=gas_pos[1,:] ; z=gas_pos[2,:]
    gw=gas_rho ; gh=gas_hsml ; gz=gas_metallicity ; gm=gas_mass
    ok_gas=ok_scan(x-x00,xmax=dx) & ok_scan(y-y00,xmax=dy) & ok_scan(z-z00,xmax=dz) & \
        ok_scan(gw,pos=1) & ok_scan(gh,pos=1) & ok_scan(gz,pos=1) & ok_scan(gm,pos=1,xmax=1.0e40);

    Ngas = checklen(gas_mass[ok_gas]);
    Nstars = checklen(source_pos[0,ok_sources]);
    if (Nstars<=1) or (Ngas<=1):
        print ' UH-OH: EXPECT ERROR NOW, there are no valid source/gas particles to send!'
        print 'Ngas=',Ngas,'Nstars=',Nstars,'dx=',dx,'dy=',dy,'dz=',dz,'x00=',x00,'y00=',y00,'z00=',z00
        return -1,-1,-1;

    dzmax=np.max(gas_pos[2,ok_gas])-z00;
    if(dzmax<OUTER_RANGE_OF_INT): OUTER_RANGE_OF_INT=dzmax;
    print 'PASSING: N_gas=',Ngas,'N_sources=',Nstars,'MaxDist=',OUTER_RANGE_OF_INT,'MinCell=',MIN_CELL_SIZE;
    Nbh=0; theta=1.0e-4; phi=1.0e-4;

    ## load the routine we need
    exec_call=util.return_python_routines_cdir()+'/LOS_column_singlePOV/getnh.so'
    NH_routine=ctypes.cdll[exec_call];
    ## cast the variables to store the results
    nh_out_cast=ctypes.c_float*Nstars;
    los_NH_out=nh_out_cast(); los_NH_hot_out=nh_out_cast(); los_Z_out=nh_out_cast();

    ## ok this is a bit arcane but the routine will read appropriately this block order
    Coord = np.zeros((Ngas+Nstars,10),dtype='f');
    Coord[0:Ngas,0] = gas_pos[0,ok_gas]-x00;
    Coord[0:Ngas,1] = gas_pos[1,ok_gas]-y00;
    Coord[0:Ngas,2] = gas_pos[2,ok_gas]-z00;
    Coord[0:Ngas,3] = gas_u[ok_gas]
    Coord[0:Ngas,4] = gas_rho[ok_gas]
    Coord[0:Ngas,5] = gas_hsml[ok_gas]
    Coord[0:Ngas,6] = gas_numh[ok_gas]
    Coord[0:Ngas,7] = gas_nume[ok_gas]
    Coord[0:Ngas,8] = gas_metallicity[ok_gas]
    Coord[0:Ngas,9] = gas_mass[ok_gas]
    Coord[Ngas:Nstars+Ngas,0] = source_pos[0,ok_sources]-x00;
    Coord[Ngas:Nstars+Ngas,1] = source_pos[1,ok_sources]-y00;
    Coord[Ngas:Nstars+Ngas,2] = source_pos[2,ok_sources]-z00;
    Coord=np.copy(np.transpose(Coord));

    ## main call to the NH-calculation routine
    NH_routine.getnh( ctypes.c_int(Ngas), ctypes.c_int(Nstars), ctypes.c_int(Nbh), \
        ctypes.c_float(theta), ctypes.c_float(phi), \
        vfloat(Coord), \
        ctypes.byref(los_NH_out),  ctypes.byref(los_NH_hot_out),  ctypes.byref(los_Z_out), \
        ctypes.c_float(OUTER_RANGE_OF_INT), ctypes.c_float(MIN_CELL_SIZE) );
    ## now put the output arrays into a useful format
    print type(los_NH_out), los_NH_out
    los_NH = np.ctypeslib.as_array(np.copy(los_NH_out));
    los_NH_hot = np.ctypeslib.as_array(np.copy(los_NH_hot_out));
    los_Z = np.ctypeslib.as_array(np.copy(los_Z_out));

    # trap for really low NH value and zero metallicity (make it small instead)
    low_NH = 1.0e10;
    los_NH[los_NH<low_NH]=low_NH; los_NH_hot[los_NH_hot<low_NH]=low_NH;
    los_Z[los_Z<=1.0e-5]=1.0e-5;

    ## assign strong attenuation to all 'off-grid' sources, then fill in calc. vals
    Nstarstot=checklen(source_pos[0,:]);
    los_NH_allgas=np.zeros(Nstarstot,dtype='f')+1.0e23;
    los_NH_hotgas=np.zeros(Nstarstot,dtype='f')+1.0e23;
    los_gas_metallicity=np.zeros(Nstarstot,dtype='f')+0.02;
    nok=checklen(los_NH_allgas[ok_sources])
    los_NH_allgas[ok_sources]=fcor(los_NH[0:Nstars]);
    los_NH_hotgas[ok_sources]=fcor(los_NH_hot[0:Nstars]);
    los_gas_metallicity[ok_sources]=fcor(los_Z[0:Nstars]);

    return los_NH_allgas, los_NH_hotgas, los_gas_metallicity;


"""Module for solving different master equations."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import numpy as np
import scipy as sp
import scipy.sparse.linalg
import copy

from qmeq.mytypes import doublenp
from qmeq.mytypes import complexnp

from qmeq import Builder
from qmeq import Approach
from qmeq import StateIndexingDM
from qmeq import StateIndexingDMc
from qmeq import QuantumDot
from qmeq import LeadsTunneling
from qmeq import FunctionProperties
from .baths import PhononBaths

#-----------------------------------------------------------
# Python modules

from .approach.pauli import Approach_pyPauli
from .approach.lindblad import Approach_pyLindblad
from .approach.neumann1 import Approach_py1vN
from .approach.redfield import Approach_pyRedfield
from qmeq.approach.neumann2 import Approach_py2vN

# Cython compiled modules

from .approach.c_pauli import Approach_Pauli
from .approach.c_lindblad import Approach_Lindblad
from .approach.c_redfield import Approach_Redfield
from .approach.c_neumann1 import Approach_1vN
from qmeq.approach.c_neumann2 import Approach_2vN
#-----------------------------------------------------------

def check_parameters(indexing, symmetry, itype, kerntype):
    if indexing is 'n':
        if symmetry is 'spin' and kerntype not in {'py2vN', '2vN'}:
            indexing = 'ssq'
        else:
            indexing = 'charge'

    if not indexing in {'Lin', 'charge', 'sz', 'ssq'}:
        print("WARNING: Allowed indexing values are: \'Lin\', \'charge\', \'sz\', \'ssq\'. "+
              "Using default indexing=\'charge\'.")
        indexing = 'charge'

    if not itype in {0,1,2,3}:
        print("WARNING: itype needs to be 0, 1, 2, or 3. Using default itype=0.")
        itype = 0

    if isinstance(kerntype, str):
        if not kerntype in {'Pauli', 'Lindblad', 'Redfield', '1vN', '2vN',
                            'pyPauli', 'pyLindblad', 'pyRedfield', 'py1vN', 'py2vN'}:
            print("WARNING: Allowed kerntype values are: "+
                  "\'Pauli\', \'Lindblad\', \'Redfield\', \'1vN\', \'2vN\', "+
                  "\'pyPauli\', \'pyLindblad\', \'pyRedfield\', \'py1vN\', \'py2vN\'. "+
                  "Using default kerntype=\'Pauli\'.")
            kerntype = 'Pauli'

    if not indexing in {'Lin', 'charge'} and kerntype in {'py2vN', '2vN'}:
        print("WARNING: For 2vN approach indexing needs to be \'Lin\' or \'charge\'. "+
              "Using indexing=\'charge\' as a default.")
        indexing = 'charge'

    return indexing, itype, kerntype

# Inherit from qmeq.Builder
class Builder_elph(Builder):

    def __init__(self, nsingle=0, hsingle={}, coulomb={},
                       nleads=0, tleads={}, mulst={}, tlst={}, dband={},
                       nbaths=0, velph={}, tlst_ph={}, dband_ph={},
                       indexing='n', kpnt=None,
                       kerntype='Pauli', symq=True, norm_row=0, solmethod='n',
                       itype=0, itype_ph=0, dqawc_limit=10000,
                       mfreeq=False, phi0_init=None,
                       mtype_qd=complex, mtype_leads=complex,
                       symmetry='n', herm_hs=True, herm_c=False, m_less_n=True,
                       bath_func=None, eps_elph=1.0e-6):
        '''
        `nbaths', `velph', `tlst_ph', `dband_ph''
        are new parameters for Electron-Phonon coupling
        '''

        indexing, itype, kerntype = check_parameters(indexing, symmetry,
                                                     itype, kerntype)

        if not itype_ph in {0,2}:
            print("WARNING: itype needs to be 0, or 2. Using default itype=0.")
            itype_ph = 0

        if isinstance(kerntype, str):
            self.Approach = globals()['Approach_'+kerntype]
        elif issubclass(kerntype, Approach):
            self.Approach = kerntype
            kerntype = self.Approach.kerntype

        # Make copies of initialized parameters.
        hsingle = copy.deepcopy(hsingle)
        coulomb = copy.deepcopy(coulomb)
        tleads = copy.deepcopy(tleads)
        mulst = copy.deepcopy(mulst)
        tlst = copy.deepcopy(tlst)
        dband = copy.deepcopy(dband)
        phi0_init = copy.deepcopy(phi0_init)
        #
        velph = copy.deepcopy(velph)
        tlst_ph = copy.deepcopy(tlst_ph)
        dband_ph = copy.deepcopy(dband_ph)

        self.funcp = FunctionProperties(symq=symq, norm_row=norm_row, solmethod=solmethod,
                                        itype=itype, dqawc_limit=dqawc_limit,
                                        mfreeq=mfreeq, phi0_init=phi0_init,
                                        mtype_qd=mtype_qd, mtype_leads=mtype_leads,
                                        kpnt=kpnt, dband=dband)
        self.funcp.itype_ph = itype_ph
        self.funcp.eps_elph = eps_elph

        icn = self.Approach.indexing_class_name
        self.si = globals()[icn](nsingle, indexing, symmetry)
        self.qd = QuantumDot(hsingle, coulomb, self.si, herm_hs, herm_c, m_less_n, mtype_qd)
        self.leads = LeadsTunneling(nleads, tleads, self.si, mulst, tlst, dband, mtype_leads)
        self.baths = PhononBaths(nbaths, velph, self.si, tlst_ph, dband_ph)
        self.baths.bath_func = bath_func

        self.appr = self.Approach(self)
        self.create_si_elph()

    def create_si_elph(self):
        si = self.si
        si_elph = StateIndexingDMc(si.nsingle, si.indexing,
                                   si.symmetry, si.nleads)
        si_elph.nbaths = si.nbaths
        self.si_elph = si_elph
        self.appr.si_elph = si_elph

    def remove_states(self, dE):
        Builder.remove_states(self, dE)
        self.si_elph.set_statesdm(si.statesdm)

    # kerntype
    def get_kerntype(self):
        return self.appr.kerntype
    def set_kerntype(self, value):
        if isinstance(value, str):
            if self.appr.kerntype != value:
                self.Approach = globals()['Approach_'+value]
                self.change_si()
                self.appr = self.Approach(self)
                self.create_si_elph()
        else:
            if issubclass(value, Approach):
                self.Approach = value
                self.change_si()
                self.appr = self.Approach(self)
                self.create_si_elph()
    kerntype = property(get_kerntype, set_kerntype)


"""
The ``block`` class represents a block of material in a simulation. Blocks are not meant to be
interacted with directly -- when using the ``problem`` class to set up a simulation, blocks are
automatically added or deleted as needed using the provided interfaces. This documentation
is provided for completeness, but should not need to be used in regular use of the code.

The class contains information on the number of grid points in the block, the location of the block
within the simulation, the material properties using the ``material`` class, the types of boundary
conditions at the block edges, and any complex geometries specified through the ``surface``
and ``curve`` classes.
"""

from __future__ import division, print_function
from os.path import join
from .surface import surface, curve
from .material import material

import numpy as np

class block(object):
    '''
    Class representing a block in a simulation

    A block contains the following internal variables:

    :ivar ndim: Number of dimensions (2 or 3)
    :type ndim: int
    :ivar mode: Rupture mode (2 or 3, relevant only for 2D problems)
    :type mode: int
    :ivar nx: Number of grid points (tuple of 3 positive integers)
    :type nx: tuple
    :ivar xm: Coordinates of lower left corner in simulation (tuple of 3 nonnegative integers)
    :type xm: tuple
    :ivar coords: Location of block within simulation domain (tuple of 3 nonnegative integers)
    :type coords: tuple
    :ivar lx: Block length in each spatial dimension (tuple of 3 positive floats, can be overridden
                by setting a curve or surface to one of the edges)
    :type lx: tuple
    :param m: Material properties (see ``material`` class)
    :type m: material
    :param bounds: List of boundary conditions. Position indicates boundary location (0 = left,
                             1 = right, 2 = front, 3 = back, 4 = bottom, 5 = top). Possible strings for
                             boundary condition include ``'absorbing'`` (no incoming wave), ``'free'``
                             (traction free surface), ``'rigid'`` (no displacement), or ``'none'`` (boundary
                             conditions determined by interface conditions)
    :type bounds: list
    :param surfs: List of bounding surfaces. Position indicates boundary location (0 = left,
                             1 = right, 2 = front, 3 = back, 4 = bottom, 5 = top). For 2D problems,
                             you can only populate the list with curves, and 3D problems require
                             surfaces. If the default rectangular surface is to be used, use ``None``
                             for a particular surface.
    :type surfs: list
    '''
    def __init__(self, ndim, mode, nx, mat):
        '''
        Initialize a new ``block`` instance

        Creates a new block instance with the given dimensionality, rupture mode, number of
        grid points, and material type. If the problem is 2d, the number of z grid points will
        be automatically set to one. By default, the block length is unity in each direction,
        the lower left coordinate is ``(0., 0., 0.)``,  all boundary conditions are set to ``'none'``,
        material properties take on their default values, and there are no irregular edge shapes.
        All of these default properties can be modified using the provided interfaces.

        :param ndim: Number of spatial dimensions (must be 2 or 3)
        :type ndim: int
        :param mode: Slip mode (2 or 3, only relevant if the problem is in 2D)
        :type mode: int
        :param nx: Tuple of length 3 with number of grid points ``(nx, ny, nz)``
        :type nx: tuple or list
        :param mat: Block material type (string, must be ``'elastic'`` or ``'plastic'``). The
                            method initializes a default set of material properties based on this type.
        :type mat: str
        :returns: New block instance
        :rtype: block
        '''
        assert(ndim == 2 or ndim == 3), "ndim must be 2 or 3"
        assert(mode == 2 or mode == 3), "mode must be 2 or 3"
        assert len(nx) == 3, "nx must be a list or tuple of positive integers"
        assert (nx[0] > 0 and nx[1] > 0 and nx[2] >0),  "nx must be a list or tuple of positive integers"
        assert (mat == "elastic" or mat == "plastic"), "material type must be elastic or plastic"
        self.ndim = int(ndim)
        self.mode = int(mode)
        self.coords = (0, 0, 0)
        if (self.ndim == 2):
            self.nx = (int(nx[0]), int(nx[1]), 1)
        else:
            self.nx = (int(nx[0]), int(nx[1]), int(nx[2]))
        self.xm = (0., 0., 0.)
        self.lx = (1., 1., 1.)
        if (self.ndim == 2):
            self.lx = (1., 1., 0.)
        self.m = material(mat)
        self.bounds = 2*self.ndim*["none"]
        self.surfs = 2*self.ndim*[None]

    def get_mode(self):
        """
        Returns rupture mode (2 or 3), only valid for 2D problems (stored at domain level)

        :returns: Rupture mode
        :rtype: int
        """
        return self.mode

    def set_mode(self,mode):
        """
        Sets rupture mode

        Rupture mode is only valid for 2D problems, and is either 2 or 3 (other values will
        cause an error, and non-integer values will be converted to integers). For 3D problems,
        entering a different value of the rupture mode will alter the rupture mode cosmetically
        but will have no effect on the simulation.

        :param mode: New value of rupture mode
        :type mode: int
        :returns: None
        """
        assert(mode == 2 or mode == 3), "Rupture mode must be 2 or 3"
        self.mode = int(mode)

    def get_ndim(self):
        """
        Returns Number of spatial dimensions

        :returns: Number of spatial dimensions
        :rtype: int
        """
        return self.ndim

    def set_ndim(self,ndim):
        """
        Sets number of dimensions

        The new number of spatial dimensions must be an integer, either 2 or 3. If a different
        value is given, the code will raise an error. If a non-integer value is given that is acceptable,
        the code will convert it to an integer.

        **Note:** Converting a 3D problem into a 2D problem will automatically collapse the
        number of grid points and the number of blocks in the $z$ direction to be 1. Any
        modifications to these quantities that were done previously will be lost.

        :param ndim: New value for ndim (must be 2 or 3)
        :type ndim: int
        :returns: None
        """
        assert(ndim == 2 or ndim == 3), "Number of dimensions must be 2 or 3"
        self.ndim = int(ndim)
        if self.ndim == 2:
            self.nx = (self.nx[0], self.nx[1], 1)
            self.xm = (self.xm[0], self.xm[1], 0.)
            self.lx = (self.lx[0], self.lx[1], 0.)
            self.bounds = self.bounds[0:4]
            self.surfs = self.surfs[0:4]
        else:
            if len(self.bounds) == 4:
                self.bounds += 2*["none"]
            if len(self.surfs) == 4:
                self.surfs += 2*[None]

    def get_nx(self):
        """
        Returns number of grid points in (nx, ny, nz) format for the given block

        :returns: Number of grid points (tuple of three integers)
        :rtype: tuple
        """
        return self.nx

    def set_nx(self,nx):
        """
        Sets number of grid points

        Changes the number of grid points to the specified tuple/list of 3 nonnegative integers.
        Bad values of ``nx`` will raise an error.

        :param nx: New value of number of grid points (tuple of 3 positive integers)
        :type nx: tuple or list
        :returns: None
        """
        assert len(nx) == 3, "nx must be a list or tuple of length 3 of positive integers"
        for i in range(3):
            assert nx[i] >= 0, "nx must be a list or tuple of length 3 of positive integers"
        if (self.ndim == 2):
            self.nx = (int(nx[0]), int(nx[1]), 1)
        else:
            self.nx = (int(nx[0]), int(nx[1]), int(nx[2]))

    def get_xm(self):
        """
        Returns starting index (zero-indexed) of block (tuple of 3 integers)

        :returns: Coordinates of lower left corner (tuple of 3 integers)
        :rtype: tuple
        """
        return self.xm

    def set_xm(self,xm):
        """
        Sets block lower left coordinate

        Changes lower left coordinate of a block to the provided tuple/list of integers.

        :param xm: New value of lower left coordinate (list/tuple of integers)
        :type xm: tuple or list
        :returns: None
        """
        assert len(xm) == 3 or (self.ndim == 2 and len(xm) == 2), "xm must be a list or tuple of length 3 of floats"

        if self.ndim == 2:
            self.xm = (float(xm[0]), float(xm[1]), 0.)
        else:
            self.xm = (float(xm[0]), float(xm[1]), float(xm[2]))

    def get_lx(self):
        """
        Returns block lengths as (lx, ly, lz) tuple

        :returns: Block dimensions (tuple of 3 floats) in x, y, and z dimensions
        :rtype: tuple
        """
        return self.lx

    def set_lx(self,lx):
        """
        Sets block lengths

        Changes block length to ``lx`` (tuple of 2 (2D only) or 3 floats) where the block length
        in each dimension is given by ``(lx, ly, lz)``

        :param lx: New value of block lengths (tuple of 2 (2D) or 3 floats)
        :type lx: tuple or list
        :returns: None
        """
        assert (len(lx) == 3 or (len(lx) == 2 and self.ndim == 2)), "lx must be a list or tuple of length 3 of positive floats"
        for l in lx:
            assert l >= 0., "lx must be a list or tuple of length 3 of positive floats"
        if self.ndim == 3:
            self.lx = (float(lx[0]), float(lx[1]), float(lx[2]))
        else:
            self.lx = (float(lx[0]), float(lx[1]), 0.)

    def get_coords(self):
        """
       Returns block coordinates (tuple of integer indices in each coordinate direction)

       :returns: Block coordinates (tuple of 3 integers)
       :rtype: tuple
       """
        return self.coords

    def set_coords(self,coords):
        """
        Sets block coordinates to a new value

        Set block coordinates to ``coords``, and tuple of nonnegative integers denoting the location
        of the block in the domain.

        :param coords: New coordaintes (tuple or list of nonnegative integers)
        :type coords: tuple or list
        :returns: None
        """
        assert len(coords) == 3, "coords must be a list or tuple of length 3 of nonnegative integers"
        for i in range(3):
            assert coords[i] >= 0, "coords must be a list or tuple of length 3 of floats"
        self.coords = (int(coords[0]), int(coords[1]), int(coords[2]))
        if self.ndim == 2:
            self.coords = (int(coords[0]), int(coords[1]), 0)

    def get_bounds(self, loc = None):
        """
        Returns boundary types

        If ``loc`` (int) is provided, the method returns a specific location (str). Otherwise it returns a list
        of all boundaries, which will have length 4 for 2D problems and length 6 for 3D problems.
        ``loc`` serves effectively as an index into the list, and the indices correspond to the following:
        0 = left, 1 = right, 2 = front, 3 = back, 4 = bottom, 5 = top. Note that the location must be
        0 <= loc < 2*ndim

        :param loc: Location of boundary that is desired (optional). If ``loc`` is not provided, returns
                           a list
        :type loc: int or None
        :returns: Boundary type (if ``loc`` provided, returns a string of the boundary type for the
                      desired location, of not returns a list of strings indicating all boundary types)
        :rtype: str or list
        """
        if loc is None:
            return self.bounds
        elif loc >= 0 and loc < 2*self.ndim:
            return self.bounds[loc]
        else:
            raise TypeError("loc must be None or an integer location")

    def set_bounds(self, bounds, loc = None):
        """
        Sets boundary types

        Changes the type of boundary conditions on a block. Acceptable values are 'absorbing'
        (incoming wave amplitude set to zero), 'free' (no traction on boundary), 'rigid' (no displacement
        of boundary), or 'none' (boundary conditions set by imposing interface conditions).

        There are two ways to use ``set_bounds``:

        1. Set ``loc`` to be ``None`` (default) and provide a list of strings specifying boundary
           type for ``bounds``. The length of ``bounds`` is 4 for a 2D simulation and 6 for 3D.

        2. Set ``loc`` to be an integer denoting location and give ``bounds`` as a single string.
           The possible locations correspond to the following: 0 = left, 1 = right, 2 = front, 3 = back,
           4 = bottom, 5 = top. 4 and 5 are only applicable to 3D simulations (0 <= loc < 2*ndim).

        :param bounds: New boundary condition type (string or list of strings)
        :type bounds: str or list
        :param loc: If provided, only change one type of boundary condition rather than all (optional,
                           loc serves as an index into the list if used)
        :type loc: int or None
        :returns: None
        """
        if loc is None:
            assert len(bounds) == 2*self.ndim, "Must give 2*ndim boundary types"
            for i in range(2*self.ndim):
                assert (bounds[i] == "none") or (bounds[i] == "absorbing") or (bounds[i] == "free") or (bounds[i] == "rigid"), "Boundary types must be none, absorbing, free, or rigid"
            self.bounds = bounds
        elif loc >=0 and loc < 2*self.ndim:
            assert (bounds == "none") or (bounds == "absorbing") or (bounds == "free") or (bounds == "rigid"), "Boundary types must be none, absorbing, free, or rigid"
            self.bounds[loc] = bounds
        else:
            raise TypeError("loc must either be None or an integer location")

    def get_surf(self, loc):
        """
        Returns block boundary surface for a block edge

        Returns the surface assigned to a specific edge. ``loc`` determines the edge that is
        returned (integer, corresponding to an index). Location indices correspond to the
        following: 0 = left, 1 = right, 2 = front, 3 = back, 4 = bottom, 5 = top
        Note that the location must be 0 <= loc < 2*ndim (for 2D problems, ``loc`` cannot be 5 or 6).

        Returns either a curve (2D problems) or surface (3D problems) or None

        If ``loc`` indices are out of bounds, the code will raise an error.

        :param loc: Location of desired boundary (0 = left, 1 = right, 2 = front, 3 = back,
                          4 = bottom, 5 = top). For 2D problems, ``loc`` must be between 0 and 3.
        :type loc: int
        :returns: curve or surface corresponding to the selected location. If the
                      desired edge does not have a bounding surface, returns None.
        :rtype: curve or surface or None
        """
        assert type(loc) is int and (loc >= 0 and loc < 2*self.ndim), "location out of range"
        return self.surfs[loc]

    def set_surf(self, loc, surf):
        """
        Sets boundary surface for a particular block edge

        Changes the bounding surface of a particular block edge. Location is determined
        by ``loc`` which is an integer that indexes into a list. Locations correspond to the
        following: 0 = left, 1 = right, 2 = front, 3 = back, 4 = bottom, 5 = top. Note that the
        location must be 0 <= loc < 2*ndim

        For 2D problems, ``surf`` must be a curve. For 3D problems, ``surf`` must be a surface.
        Other choices will raise an error. If ``loc`` is out of bounds, the code
        will also signal an error.

        :param loc: Location of desired boundary (0 = left, 1 = right, 2 = front, 3 = back,
                          4 = bottom, 5 = top). For 2D problems, ``loc`` must be between 0 and 3.
        :type loc: int
        :param surf: curve or surface corresponding to the selected block and location
        :type surf: curve or surface
        :returns: None
        """
        assert type(loc) is int and (loc >= 0 and loc < 2*self.ndim), "location out of range"
        if self.ndim == 3:
            assert type(surf) is surface
            if loc == 0 or loc == 1:
                assert surf.get_direction() == 'x', "surface direction does not match location"
                assert surf.get_n1() == self.nx[1] and surf.get_n2() == self.nx[2], "number of grid points does not match"
            elif loc == 2 or loc == 3:
                assert surf.get_direction() == 'y', "surface direction does not match location"
                assert surf.get_n1() == self.nx[0] and surf.get_n2() == self.nx[2], "number of grid points does not match"
            else:
                assert surf.get_direction() == 'z', "surface direction does not match location"
                assert surf.get_n1() == self.nx[0] and surf.get_n2() == self.nx[1], "number of grid points does not match"
        else:
            assert type(surf) is curve
            if loc == 0 or loc == 1:
                assert surf.get_direction() == 'x', "surface direction does not match location"
                assert surf.get_n1() == self.nx[1], "number of grid points does not match"
            else:
                assert surf.get_direction() == 'y', "surface direction does not match location"
                assert surf.get_n1() == self.nx[0], "number of grid points does not match"
        self.surfs[loc] = surf

    def delete_surf(self, loc):
        """
        Removes boundary surface for a particular block edge

        Removes the bounding surface of a particular block edge. Location is determined by
        ``loc`` which is an integer that indexes into a list. Locations correspond to the following:
        0 = left, 1 = right, 2 = front, 3 = back, 4 = bottom, 5 = top. Note that the location must be
        0 <= loc < 2*ndim

        If ``loc`` is out of bounds, the code will also signal an error.

        :param loc: Location of desired boundary to be removed (0 = left, 1 = right, 2 = front,
                          3 = back, 4 = bottom, 5 = top). For 2D problems, ``loc`` must be between 0 and 3.
        :type loc: int
        :returns: None
        """
        assert type(loc) is int and (loc >= 0 and loc < 2*self.ndim), "location out of range"
        self.surfs[loc] = None

    def get_material(self):
        """
        Returns material

        Returns the material class associated with this block

        :returns: Material class with properties for this block
        :rtype: material
        """
        return self.m

    def set_mattype(self, mattype):
        """
        Sets block material type ('elastic' or 'plastic')

        Sets the material type for the block. Options are 'elastic' for an elastic simulation
        and 'plastic' for a plastic simulation. Anything else besides these options will cause the
        code to raise an error.

        :param mattype: New material type ('elastic' or 'plastic')
        :type mattype: str
        :returns: None
        """
        self.m.set_type(mattype)

    def set_material(self,mat):
        """
        Sets block material properties

        Sets new material properties stored in an instance of the ``material`` class.

        :param newmaterial: New material properties
        :type newmaterial: material
        :param coords: Coordinates of block to be changed (optional, omitting changes all blocks).
                                 ``coords`` must be a tuple or list of three integers that match the coordinates
                                 of a block.
        :type coords: tuple or list
        :returns: None
        """
        assert type(mat) is material
        self.m = mat

    def get_x(self, coord):
        """
        Returns grid value for given spatial index

        For a given problem set up, returns the location of a particular set of coordinate indices.
        Note that since blocks are set up by setting values only on the edges, coordinates on
        the interior are not specified *a priori* and instead determined using transfinite interpolation
        to generate a regular grid on the block interiors. Calling ``get_x`` generates the interior grid
        to find the coordinates of the desired point.

        Within each call to ``get_x``, the grid is generated on the fly only for the relevant block
        where the desired point is located. It is not stored. This helps reduce memory requirements
        for large 3D problems (since the Python module does not run in parallel), but is slower.
        Because the computational grid is regular, though, it can be done in a single step in closed
        form.

        Returns a numpy array of length 3 holding the spatial location (x, y, z).

        :param coord: Spatial coordinate where grid values are desired (tuple or list of 3 integers
                               or 2 integers for 2D problems)
        :type coord: tuple or list
        :returns: (x, y, z) coordinates of spatial location
        :rtype: ndarray
        """

        if self.ndim == 2:
            assert (len(coord) == 2 or len(coord) == 3), "Coordinates must have length 2 or 3"
            coord = (coord[0], coord[1])
        else:
            assert len(coord) == 3, "Coordinates must have length 3"
        for i in range(self.ndim):
            assert (coord[i] >= 0 and coord[i] < self.nx[i]), "Coordinate value out of range"

        # make temporary surfaces and check that edges match

        tmpsurfs = self.make_tempsurfs()
        self.checksurfs(tmpsurfs)

        p = float(coord[0])/float(self.nx[0]-1)
        q = float(coord[1])/float(self.nx[1]-1)
        x = np.zeros(3)
        if self.ndim == 2:
            x[0] = ((1.-p)*tmpsurfs[0].get_x(coord[1])+p*tmpsurfs[1].get_x(coord[1])+
                        (1.-q)*tmpsurfs[2].get_x(coord[0])+q*tmpsurfs[3].get_x(coord[0])-
                        (1.-p)*(1.-q)*tmpsurfs[0].get_x(0)-(1.-q)*p*tmpsurfs[1].get_x(0)-
                        q*(1.-p)*tmpsurfs[0].get_x(-1)-q*p*tmpsurfs[1].get_x(-1))
            x[1] = ((1.-p)*tmpsurfs[0].get_y(coord[1])+p*tmpsurfs[1].get_y(coord[1])+
                            (1.-q)*tmpsurfs[2].get_y(coord[0])+q*tmpsurfs[3].get_y(coord[0])-
                            (1.-p)*(1.-q)*tmpsurfs[0].get_y(0)-(1.-q)*p*tmpsurfs[1].get_y(0)-
                            q*(1.-p)*tmpsurfs[0].get_y(-1)-q*p*tmpsurfs[1].get_y(-1))
        else:
            r = float(coord[2])/float(self.nx[2]-1)
            x[0] = ((1.-p)*tmpsurfs[0].get_x((coord[1], coord[2]))+p*tmpsurfs[1].get_x((coord[1], coord[2]))+
                    (1.-q)*tmpsurfs[2].get_x((coord[0], coord[2]))+q*tmpsurfs[3].get_x((coord[0], coord[2]))+
                    (1.-r)*tmpsurfs[4].get_x((coord[0], coord[1]))+r*tmpsurfs[5].get_x((coord[0], coord[1])))
            x[1] = ((1.-p)*tmpsurfs[0].get_y((coord[1], coord[2]))+p*tmpsurfs[1].get_y((coord[1], coord[2]))+
                    (1.-q)*tmpsurfs[2].get_y((coord[0], coord[2]))+q*tmpsurfs[3].get_y((coord[0], coord[2]))+
                    (1.-r)*tmpsurfs[4].get_y((coord[0], coord[1]))+r*tmpsurfs[5].get_y((coord[0], coord[1])))
            x[2] = ((1.-p)*tmpsurfs[0].get_z((coord[1], coord[2]))+p*tmpsurfs[1].get_z((coord[1], coord[2]))+
                    (1.-q)*tmpsurfs[2].get_z((coord[0], coord[2]))+q*tmpsurfs[3].get_z((coord[0], coord[2]))+
                    (1.-r)*tmpsurfs[4].get_z((coord[0], coord[1]))+r*tmpsurfs[5].get_z((coord[0], coord[1])))
            x[0] -= ((1.-q)*(1.-p)*tmpsurfs[0].get_x((0, coord[2]))+(1.-q)*p*tmpsurfs[1].get_x((0, coord[2]))+
                     q*(1.-p)*tmpsurfs[0].get_x((-1,coord[2]))+q*p*tmpsurfs[1].get_x((-1,coord[2]))+
                     (1.-p)*(1.-r)*tmpsurfs[0].get_x((coord[1], 0))+p*(1.-r)*tmpsurfs[1].get_x((coord[1], 0))+
                     (1.-q)*(1.-r)*tmpsurfs[2].get_x((coord[0], 0))+q*(1.-r)*tmpsurfs[3].get_x((coord[0], 0))+
                     (1.-p)*r*tmpsurfs[0].get_x((coord[1], -1))+p*r*tmpsurfs[1].get_x((coord[1],-1))+
                     (1.-q)*r*tmpsurfs[2].get_x((coord[0], -1))+q*r*tmpsurfs[3].get_x((coord[0], -1)))
            x[1] -= ((1.-q)*(1.-p)*tmpsurfs[0].get_y((0, coord[2]))+(1.-q)*p*tmpsurfs[1].get_y((0, coord[2]))+
                     q*(1.-p)*tmpsurfs[0].get_y((-1,coord[2]))+q*p*tmpsurfs[1].get_y((-1,coord[2]))+
                     (1.-p)*(1.-r)*tmpsurfs[0].get_y((coord[1], 0))+p*(1.-r)*tmpsurfs[1].get_y((coord[1], 0))+
                     (1.-q)*(1.-r)*tmpsurfs[2].get_y((coord[0], 0))+q*(1.-r)*tmpsurfs[3].get_y((coord[0], 0))+
                     (1.-p)*r*tmpsurfs[0].get_y((coord[1], -1))+p*r*tmpsurfs[1].get_y((coord[1],-1))+
                     (1.-q)*r*tmpsurfs[2].get_y((coord[0], -1))+q*r*tmpsurfs[3].get_y((coord[0], -1)))
            x[2] -= ((1.-q)*(1.-p)*tmpsurfs[0].get_z((0, coord[2]))+(1.-q)*p*tmpsurfs[1].get_z((0, coord[2]))+
                     q*(1.-p)*tmpsurfs[0].get_z((-1,coord[2]))+q*p*tmpsurfs[1].get_z((-1,coord[2]))+
                     (1.-p)*(1.-r)*tmpsurfs[0].get_z((coord[1], 0))+p*(1.-r)*tmpsurfs[1].get_z((coord[1], 0))+
                     (1.-q)*(1.-r)*tmpsurfs[2].get_z((coord[0], 0))+q*(1.-r)*tmpsurfs[3].get_z((coord[0], 0))+
                     (1.-p)*r*tmpsurfs[0].get_z((coord[1], -1))+p*r*tmpsurfs[1].get_z((coord[1],-1))+
                     (1.-q)*r*tmpsurfs[2].get_z((coord[0], -1))+q*r*tmpsurfs[3].get_z((coord[0], -1)))
            x[0] += ((1.-p)*(1.-q)*(1.-r)*tmpsurfs[0].get_x((0,0))+p*(1.-q)*(1.-r)*tmpsurfs[1].get_x((0,0))+
                     (1.-p)*q*(1.-r)*tmpsurfs[0].get_x((-1,0))+(1.-p)*(1.-q)*r*tmpsurfs[0].get_x((0,-1))+
                     p*q*(1.-r)*tmpsurfs[1].get_x((-1,0))+p*(1.-q)*r*tmpsurfs[1].get_x((0,-1))+
                     (1.-p)*q*r*tmpsurfs[0].get_x((-1,-1))+p*q*r*tmpsurfs[1].get_x((-1,-1)))
            x[1] += ((1.-p)*(1.-q)*(1.-r)*tmpsurfs[0].get_y((0,0))+p*(1.-q)*(1.-r)*tmpsurfs[1].get_y((0,0))+
                     (1.-p)*q*(1.-r)*tmpsurfs[0].get_y((-1,0))+(1.-p)*(1.-q)*r*tmpsurfs[0].get_y((0,-1))+
                     p*q*(1.-r)*tmpsurfs[1].get_y((-1,0))+p*(1.-q)*r*tmpsurfs[1].get_y((0,-1))+
                     (1.-p)*q*r*tmpsurfs[0].get_y((-1,-1))+p*q*r*tmpsurfs[1].get_y((-1,-1)))
            x[2] += ((1.-p)*(1.-q)*(1.-r)*tmpsurfs[0].get_z((0,0))+p*(1.-q)*(1.-r)*tmpsurfs[1].get_z((0,0))+
                     (1.-p)*q*(1.-r)*tmpsurfs[0].get_x((-1,0))+(1.-p)*(1.-q)*r*tmpsurfs[0].get_z((0,-1))+
                     p*q*(1.-r)*tmpsurfs[1].get_z((-1,0))+p*(1.-q)*r*tmpsurfs[1].get_z((0,-1))+
                     (1.-p)*q*r*tmpsurfs[0].get_z((-1,-1))+p*q*r*tmpsurfs[1].get_z((-1,-1)))

        return np.array(x)

    def make_tempsurfs(self):
        """
        Create temporary surface list

        This method generates all six (four in 2D) bounding surfaces (curves in 2D). Note that
        these surfaces are not usually stored for rectangular block edges to save memory,
        as they are trivial to create. The temporary surfaces can be used to check that the
        edges of the surfaces/curves match or to use transfinite interpolation to generate the grid.

        :returns: List of all bounding surfaces (not stored beyond the time they are needed)
        :rtype: list
        """

        tmpsurf = []

        if self.ndim == 2:
            for i in range(4):
                if self.surfs[i] is None:
                    if i == 0:
                        tmpsurf.append(curve(self.nx[1], 'x', np.ones(self.nx[1])*self.xm[0],
                                             np.linspace(self.xm[1], self.xm[1]+self.lx[1], self.nx[1])))
                    elif i == 1:
                        tmpsurf.append(curve(self.nx[1], 'x', np.ones(self.nx[1])*(self.xm[0]+self.lx[0]),
                                             np.linspace(self.xm[1], self.xm[1]+self.lx[1], self.nx[1])))
                    elif i == 2:
                        tmpsurf.append(curve(self.nx[0], 'y', np.linspace(self.xm[0], self.xm[0]+self.lx[0], self.nx[0]),
                                             np.ones(self.nx[0])*self.xm[1]))
                    else:
                        tmpsurf.append(curve(self.nx[0], 'y',np.linspace(self.xm[0], self.xm[0]+self.lx[0], self.nx[0]),
                                             np.ones(self.nx[0])*(self.xm[1]+self.lx[1])))
                else:
                    tmpsurf.append(self.surfs[i])

        else:
            for i in range(6):
                if self.surfs[i] is None:
                    if i == 0:
                        tmpsurf.append(surface(self.nx[1], self.nx[2], 'x', np.ones((self.nx[1], self.nx[2]))*self.xm[0],
                                               np.meshgrid(np.linspace(self.xm[1], self.xm[1]+self.lx[1], self.nx[1]),np.linspace(self.xm[2], self.xm[2]+self.lx[2], self.nx[2]), indexing='ij')[0],
                                               np.meshgrid(np.linspace(self.xm[1], self.xm[1]+self.lx[1], self.nx[1]),np.linspace(self.xm[2], self.xm[2]+self.lx[2], self.nx[2]), indexing='ij')[1]))
                    elif i == 1:
                        tmpsurf.append(surface(self.nx[1], self.nx[2], 'x', np.ones((self.nx[1], self.nx[2]))*(self.xm[0]+self.lx[0]),
                                               np.meshgrid(np.linspace(self.xm[1], self.xm[1]+self.lx[1], self.nx[1]),np.linspace(self.xm[2], self.xm[2]+self.lx[2], self.nx[2]), indexing='ij')[0],
                                               np.meshgrid(np.linspace(self.xm[1], self.xm[1]+self.lx[1], self.nx[1]),np.linspace(self.xm[2], self.xm[2]+self.lx[2], self.nx[2]), indexing='ij')[1]))
                    elif i == 2:
                        tmpsurf.append(surface(self.nx[0], self.nx[2], 'y',
                                               np.meshgrid(np.linspace(self.xm[0], self.xm[0]+self.lx[0], self.nx[0]),np.linspace(self.xm[2], self.xm[2]+self.lx[2], self.nx[2]), indexing='ij')[0],
                                               np.ones((self.nx[0], self.nx[2]))*self.xm[1],
                                               np.meshgrid(np.linspace(self.xm[0], self.xm[0]+self.lx[0], self.nx[0]),np.linspace(self.xm[2], self.xm[2]+self.lx[2], self.nx[2]), indexing='ij')[1]))
                    elif i == 3:
                        tmpsurf.append(surface(self.nx[0], self.nx[2], 'y',
                                               np.meshgrid(np.linspace(self.xm[0], self.xm[0]+self.lx[0], self.nx[0]),np.linspace(self.xm[2], self.xm[2]+self.lx[2], self.nx[2]), indexing='ij')[0],
                                               np.ones((self.nx[0], self.nx[2]))*(self.xm[1]+self.lx[1]),
                                               np.meshgrid(np.linspace(self.xm[0], self.xm[0]+self.lx[0], self.nx[0]),np.linspace(self.xm[2], self.xm[2]+self.lx[2], self.nx[2]), indexing='ij')[1]))
                    elif i == 4:
                        tmpsurf.append(surface(self.nx[0], self.nx[1], 'z',
                                               np.meshgrid(np.linspace(self.xm[0], self.xm[0]+self.lx[0], self.nx[0]),np.linspace(self.xm[1], self.xm[1]+self.lx[1], self.nx[1]), indexing='ij')[0],
                                               np.meshgrid(np.linspace(self.xm[0], self.xm[0]+self.lx[0], self.nx[0]),np.linspace(self.xm[1], self.xm[1]+self.lx[1], self.nx[1]), indexing='ij')[1],
                                               np.ones((self.nx[0], self.nx[1]))*self.xm[2]))
                    else:
                        tmpsurf.append(surface(self.nx[0], self.nx[1], 'z',
                                               np.meshgrid(np.linspace(self.xm[0], self.xm[0]+self.lx[0], self.nx[0]),np.linspace(self.xm[1], self.xm[1]+self.lx[1], self.nx[1]), indexing='ij')[0],
                                               np.meshgrid(np.linspace(self.xm[0], self.xm[0]+self.lx[0], self.nx[0]),np.linspace(self.xm[1], self.xm[1]+self.lx[1], self.nx[1]), indexing='ij')[1],
                                               np.ones((self.nx[0], self.nx[1]))*(self.xm[2]+self.lx[2])))
                else:
                    tmpsurf.append(self.surfs[i])

        return tmpsurf

    def check(self):
        """
        Checks for errors before writing input file

        Checks that edges of bounding surfaces match. If the edges are not defined as
        surfaces, the code temporarily creates them to check that they match.

        :returns: None
        """

        tmpsurfs = self.make_tempsurfs()
        self.checksurfs(tmpsurfs)

    def checksurfs(self, tmpsurfs):
        """
        Checks that surface boundaries match

        Input is a list of surfaces, with order corresponding to (left, right, front, back, top, bottom).
        In 2D problems, there is no top or bottom surface.

        :param tmpsurfs: List of surfaces to compare (lenth 4 or 6)
        :type tmpsurfs: list
        :returns: None
        """

        surf1 = [0, 0, 1, 1, 0, 0, 1, 1, 2, 2, 3, 3]
        surf2 = [2, 3, 2, 3, 4, 5, 4, 5, 4, 5, 4, 5]
        edge1 = [1, 3,1, 3, 0, 2, 0, 2, 0, 2, 0, 2]
        edge2 = [1, 1, 3, 3, 1, 1, 3, 3, 0, 0, 2, 2]

        for i in range(2**(self.ndim-1)*self.ndim):
            assert tmpsurfs[surf1[i]].has_same_edge(edge1[i], edge2[i], tmpsurfs[surf2[i]]), "surface edges do not match"

    def write_input(self, f, probname, directory, endian = '='):
        """
        Writes block information to input file

        Method writes information for block to file. It also writes all relevant surface data to file
        to describe non-rectangular geometries. Inputs inlcude the file handle for the input
        file, the problem name (used for naming surface files), the destination directory for
        all files, and endianness (optional, default is native) for binary surface files.

        :param f: file handle for input file
        :type f: file
        :param probname: Problem name
        :type probname: str
        :param directory: Directory where output should be written
        :type directory: str
        :param endian: Byte-ordering for binary files for surface data. Possible values are
                                ``'<'`` (little endian), ``'>'`` (big endian), or ``'='`` (native, default)
        :type endian: str
        :returns: None
        """

        if directory == "":
            inputfiledir = 'problems/'
        else:
            inputfiledir = directory

        f.write("[fdfault.block"+str(self.coords[0])+str(self.coords[1])+str(self.coords[2])+"]\n")
        self.m.write_input(f)
        outstring = ""
        for i in range(self.ndim):
            outstring += repr(self.xm[i])+" "
        outstring = outstring[0:-1]
        f.write(outstring+"\n")
        outstring = ""
        for i in range(self.ndim):
            outstring += repr(self.lx[i])+" "
        outstring = outstring[0:-1]
        f.write(outstring+"\n")
        for btype in self.bounds:
            f.write(btype+"\n")
        nsurfs = 0
        for s in self.surfs:
            if s is None:
                f.write("none\n")
            else:
                f.write(join(inputfiledir, probname)+"_block"+str(self.coords[0])+str(self.coords[1])+str(self.coords[2])+str(nsurfs)+".surf\n")
                s.write(join(directory, probname+"_block"+str(self.coords[0])+str(self.coords[1])+str(self.coords[2])+str(nsurfs)+".surf"), endian)
            nsurfs += 1
        f.write("\n")

    def __str__(self):
        '''
        returns a string representation
        '''
        surfstring = ''
        for surf in self.surfs:
            surfstring += str(surf)+"\n"
        return ("Block "+str(self.coords)+":\nnx = "+str(self.nx)+"\nxm = "+str(self.xm)+"\nlx = "+
                str(self.lx)+"\nbounds = "+str(self.bounds)+"\nsurfaces =\n"+surfstring+str(self.m))


#!/usr/bin/env python

"""Calculate total number of fully bound staples for a simulation set"""

import argparse
import os.path

import numpy as np

from origamipy import biases
from origamipy import conditions
from origamipy import config_process
from origamipy import datatypes
from origamipy import files
from origamipy import outputs
from origamipy import decorrelate
from origamipy import mbar_wrapper

def main():
    args = parse_args()
    system_file = files.JSONStructInpFile(args.system_filename)
    domain_pairs = parse_domain_pairs(args.domain_pairs)

    fileformatter = construct_fileformatter()
    all_conditions = construct_conditions(args, fileformatter, system_file)
    inp_filebase = create_input_filepathbase(args)
    sim_collections = outputs.create_sim_collections(inp_filebase,
                                                     all_conditions, args.reps)
    for sim_collection in sim_collections:
        for rep in sim_collection._reps:
            ops = sim_collection.get_reps_data('ops', concatenate=False)
            runs = len(ops[rep])
            for run in range(runs):
                run_filebase = sim_collection.get_filebase(run, rep)
                trj_filename = '{}.trj'.format(run_filebase)
                trj_file = files.TxtTrajInpFile(trj_filename, system_file)
                ops = datatypes.OrderParams.from_file(run_filebase)
                all_dists = [[] for i in range(len(domain_pairs))]
                for i, step in enumerate(trj_file):
                    config = np.array(step[0]['positions'])
                    for j, domain_pair in enumerate(domain_pairs):
                        pos_i = config[domain_pair[0]]
                        pos_j = config[domain_pair[1]]
                        all_dists[j].append(
                            config_process.calc_dist(pos_i, pos_j))

                for domain_pair, dists in zip(domain_pairs, all_dists):
                    dist_tag = 'dist-d{}-d{}'.format(
                        domain_pair[0], domain_pair[1])
                    if dist_tag in ops.tags:
                        ops[dist_tag] = dists
                    else:
                        ops.add_column(dist_tag, dists)
                    adj_tag = 'adj-d{}-d{}'.format(
                        domain_pair[0], domain_pair[1])
                    adj_sites = np.array(dists) == 1
                    if adj_tag in ops.tags:
                        ops[adj_tag] = adj_sites
                    else:
                        ops.add_column(adj_tag, adj_sites)

                dist_sum = np.sum(all_dists, axis=0)
                tag = 'dist-sum'
                if tag in ops.tags:
                    ops[tag] = dist_sum
                else:
                    ops.add_column(tag, dist_sum)

                ops.to_file(run_filebase)

def construct_conditions(args, fileformatter, system_file):
    stack_biases = []
    for stack_mult in args.stack_mults:
        stack_bias = biases.StackingBias(args.stack_ene, stack_mult)
        stack_biases.append(stack_bias)

    conditions_map = {'temp': args.temps,
                      'staple_m': [args.staple_m],
                      'bias': stack_biases}

    return conditions.AllSimConditions(conditions_map, fileformatter, system_file)

def construct_fileformatter():
    specs = []
    specs.append(conditions.ConditionsFileformatSpec('temp', '{}'))
    specs.append(conditions.ConditionsFileformatSpec('bias', '{}'))

    return conditions.ConditionsFileformatter(specs)

def create_input_filepathbase(args):
    return '{}/{}'.format(args.input_dir, args.filebase)

def create_output_filepathbase(args):
    return '{}/{}'.format(args.output_dir, args.filebase)

def parse_domain_pairs(domain_pair_strings):
    domain_pairs = []
    for domain_pair in domain_pair_strings:
        domain_pairs.append([int(d) for d in domain_pair.split(',')])

    return domain_pairs

def parse_args():
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument(
        'system_filename',
        type=str,
        help='System file')
    parser.add_argument(
        'filebase',
        type=str,
        help='Base name for files')
    parser.add_argument(
        'input_dir',
        type=str,
        help='Directory of inputs')
    parser.add_argument(
        'output_dir',
        type=str,
        help='Directory to output to')
    parser.add_argument(
        'staple_m',
        type=float,
        help='Staple molarity (mol/V)')
    parser.add_argument(
        'stack_ene',
        type=float,
        help='Stacking energy (kb K)')
    parser.add_argument(
        '--reps',
        nargs='+',
        type=int,
        help='Reps (leave empty for all available)')
    parser.add_argument(
        '--temps',
        nargs='+',
        type=str,
        help='Temperatures')
    parser.add_argument(
        '--stack_mults',
        nargs='+',
        type=str,
        help='Stacking energy multipliers')
    parser.add_argument(
        '--domain_pairs',
        nargs='+',
        type=str,
        help='Scaffold domain pairs to calculate distances between')

    return parser.parse_args()

if __name__ == '__main__':
    main()