import numpy as np

def provide_PSF_2D(x=None,y=None,PSF_version=None):
    """ Provides 2D PSF at any position in the detector plane
        This a version which takes a finite nubmer of pregenerated PSF and \
        creates the interpolated version at required position
        (Future: version which takes interpolated values for Zernike \
        coefficients and generates image on the fly?)

        To be used with the focused data taken on July 25 and 26
        (e.g., 21400 for HgAr, 21604 for Ne, 21808 for Kr)

        Example usage: ``provide_PSF_2D(10,2010)'' 10 is x-coordinate,
        and 2010 is y-coordinate

    @param[in] x            x-coordinate
    @param[in] y            y-coordinate
    @param[in] PSF_version  version of the PSF input files
    @returns                numpy array, 100x100, oversampled 5 times,
                            corresponding to 20x20 physical pixels
                            (300x300 microns)
    """

    # on tiger the directory contaning array of PSFs is at:
    DATA_DIRECTORY='/tigress/ncaplar/PIPE2D-450/'

    if PSF_version is None:
        PSF_version='Sep12_v1'

    positions_of_simulation=np.load(DATA_DIRECTORY+\
                        'positions_of_simulation_00_from_'+PSF_version+'.npy')
    array_of_simulation=np.load(DATA_DIRECTORY+\
                        'array_of_simulation_00_from_'+PSF_version+'.npy')

    # x and y position with simulated PSFs
    x_positions_of_simulation=positions_of_simulation[:,1]
    y_positions_of_simulation=positions_of_simulation[:,2]

    # This is a simple code that finds the closest avaliable PSFs, given the x and y position
    # This will have to be improved in order when we get to work with the full populated dectector plane

    # how far in x-dimension are you willing to search for suitable simulated PSFs
    x_search_distance=20
    # positions of all simulated PSFs in that range
    positions_of_simulation_in_acceptable_x_range=\
    positions_of_simulation[(x_positions_of_simulation<(x+x_search_distance))\
                            &(x_positions_of_simulation>(x-x_search_distance))]

    # if there are no simulated PSF avaliable in the specified x-range we are not able to provide the solution
    if len(positions_of_simulation_in_acceptable_x_range)<2:
        print('No simulated PSFs are avaliable in this x-area of the detector,')
        print('probably because this fiber has not been illuminated;')
        print('returning the closest avaliable PSFs, BUT that is probably not what you want')
        distances=np.sqrt(((x-x_positions_of_simulation)**2+\
                           (y-y_positions_of_simulation)**2).astype(float))
        index_of_closest_distance=np.where(distances[distances==\
                                                     np.min(distances)])[0][0]
        return array_of_simulation[index_of_closest_distance]

    # y-distance from the requested positions for all of the suitable simulated PSFs
    distances_of_y_requested_position_from_avaliable=\
    y-positions_of_simulation_in_acceptable_x_range[:,2]
    # out of the suitable PSFs which 2 are the closest
    index_of_1st_closest_simulated_psf=\
    np.where(np.abs(distances_of_y_requested_position_from_avaliable)==\
             np.sort(np.abs(distances_of_y_requested_position_from_avaliable))[0])[0][0]
    index_of_2nd_closest_simulated_psf=\
    np.where(np.abs(distances_of_y_requested_position_from_avaliable)==\
             np.sort(np.abs(distances_of_y_requested_position_from_avaliable))[1])[0][0]
    # where are these 2 closest PSF in the initial table
    index_of_1st_closest_simulated_psf_in_positions_of_simulation=\
    np.where(np.sum(positions_of_simulation,axis=1)==\
             np.sum(positions_of_simulation_in_acceptable_x_range[index_of_1st_closest_simulated_psf]))[0][0]
    index_of_2nd_closest_simulated_psf_in_positions_of_simulation=\
    np.where(np.sum(positions_of_simulation,axis=1)==\
             np.sum(positions_of_simulation_in_acceptable_x_range[index_of_2nd_closest_simulated_psf]))[0][0]
    # extract the 2 simulated PSFs
    first_array_simulation=\
    array_of_simulation[index_of_1st_closest_simulated_psf_in_positions_of_simulation]
    second_array_simulation=\
    array_of_simulation[index_of_2nd_closest_simulated_psf_in_positions_of_simulation]
    # distance of each PSF from the proposed position
    y1_distance=\
    y-positions_of_simulation[index_of_1st_closest_simulated_psf_in_positions_of_simulation][2]
    y2_distance=\
    y-positions_of_simulation[index_of_2nd_closest_simulated_psf_in_positions_of_simulation][2]

    # if you requested psf at the exact position of existing PSF use that one
    if y1_distance==0:
        return first_array_simulation
    else:
        # create the predicted PSF as a linear interpolation of these two PSFs
        predicted_psf=(second_array_simulation-first_array_simulation*(y2_distance/y1_distance))/(1-y2_distance/y1_distance)
        return predicted_psf


"""
Class for calculating IRC using Morokuma Algorithm
"""

import sys
import os
import psi4
import json
import numpy as np

def JSON2XYZ(input_params):
    charge = input_params['molecule']['molecular_charge']
    mult = input_params['molecule']['molecular_multiplicity']
    geom = ''
    geom += '\n%d %d\n' %(charge, mult)
    symbols = input_params['molecule']['symbols']
    geometry = input_params['molecule']['geometry']
    for i in range(len(symbols)):
        geom += "%s  " %symbols[i]
        for j in range(3):
            if(j==2):
                geom += "   %f   \n" %geometry[(i + 2*i) + j]
            else:
                geom += "   %f   " %geometry[(i + 2*i) + j]
    #print(geom)
    return geom

class ToolKit():
    # Placeholder for Common Data
    def __init__(self,name):
        self.natoms=0
        self.basis = '3-21G'
        self.method = 'SCF'
        self.level_of_theory = "%s/%s" %(self.method,self.basis)
        self.alpha=0.1
        self.basename='.'.join(name.split('.')[:-1])
        self.out=open("%s.log"%(self.basename),'a',1)
        self.delta=0.05
        self.direction=1
        self.energies=[]
        self.energy=0.00
        self.geos=[]
        self.restart=False
        self.geometry=[]
        self.guessfn=''
        self.damp=0.05
        self.algorithm=1
        self.autodamp=False
        self.prevgrad=0.0
        self.hessfn=''
        self.maxdispl=0.01
        self.mode=4
        self.dgrad=0.0
        self.npoints=25
        self.template=[]
        self.keywords={}
        self.tolerance=1.0e-04
        self.orcacmd='UNDEFINED'
        self.ReadInput(name)
        self.ComputeHessian()
    def printPars(self):
        stmp="  %13s: %s\n"
        ftmp="  %13s: %7.4f\n"
        itmp="  %13s: %7d\n"
        self.out.write('')
        self.out.write(itmp%('Algorithm',self.algorithm))
        self.out.write(itmp%('N. Points',self.npoints))
        self.out.write(ftmp%('Grad. Tol.',self.tolerance))
        self.out.write('')
        self.out.write(stmp%('Hessian',self.hessfn))
        self.out.write(itmp%('Mode',self.mode))
        self.out.write(itmp%('Direction',self.direction))
        self.out.write('')
        self.out.write(ftmp%('Alpha',self.alpha))
        self.out.write(ftmp%('Delta',self.delta))
        self.out.write('')
        self.out.write(ftmp%('Damp Factor',self.damp))
        self.out.write(itmp%('Damp Update',self.autodamp))
        self.out.write('')
        self.out.write(ftmp%('Max. Displ.',self.maxdispl))
        self.out.write('')
        self.out.write(stmp%('Guess',self.guessfn))
        self.out.write("\n------------------------------------------------\n")
    def ReadInput(self,json_input):
        # Read in input from JSON input file
        json_data=open(json_input).read()
        input_params = json.loads(json_data)
        if input_params['molecule']['geometry']:
            self.geometry = JSON2XYZ(input_params)
        if input_params['model']['basis']:
            self.basis = input_params['model']['basis']
            self.level_of_theory = "%s/%s" %(self.method,self.basis)
            #print("IN HERE")
        if input_params['model']['method']:
            self.method = input_params['model']['method']
            self.level_of_theory = "%s/%s" %(self.method,self.basis)
        if input_params['molecule']['symbols']:
            self.natoms = len(input_params['molecule']['symbols'])
        if input_params['keywords']:
            self.keywords = input_params['keywords']
        if "irc" in input_params:
            if input_params['irc']['hessfile']:
                self.hessfn = input_params['irc']['hessfile']
            if input_params['irc']['guess']:
                self.guessfn = input_params['irc']['guess']
            if input_params['irc']['alpha']:
                self.alpha = input_params['irc']['alpha']
            if input_params['irc']['delta']:
                self.delta = input_params['irc']['delta']
            if input_params['irc']['damp']:
                self.damp = input_params['irc']['damp']
            if input_params['irc']['restart']:
                if(input_params['irc']['restart']==1):
                    self.restart = True
                else:
                    self.restart = False
            if input_params['irc']['autodamp']:
                if(input_params['irc']['autodamp']==1):
                    self.autodamp=True
                else:
                    self.autodamp=False
            if input_params['irc']['tol']:
                self.tolerance= input_params['irc']['tol']
            if input_params['irc']['alg']:
                self.algorithm = input_params['irc']['alg']
            if input_params['irc']['dir']:
                self.direction = input_params['irc']['dir']
            if input_params['irc']['mode']:
                self.mode = input_params['irc']['mode']
            if input_params['irc']['maxd']:
                self.maxdispl = input_params['irc']['maxd']
            if input_params['irc']['pts']:
                self.npoints = input_params['irc']['pts']
        else:
            pass
   # def JSON2XYZ(self, input_params):
   #     geom = ''
   #     geom += '%d ' %input_params['molecule']['molecular_charge']
   #     geom += '%d \n' %input_params['molecule']['molecular_multiplicity']
   #     symbols = input_params['molecule']['symbols']
   #     geometry = input_params['molecule']['geometry']
   #     for i in range(len(symbols)):
   #         geom += "%s  " %symbols[i]
   #         for j in range(3):
   #             if(j==2):
   #                 geom += "   %f   \n" %geometry[(i + 2*i) + j]
   #             else:
   #                 geom += "   %f   " %geometry[(i + 2*i) + j]
   #     self.geom = geom

    def ComputeHessian(self):
        # Use PSI4 to calculate Hessian matrix
        psi4.core.set_output_file("hessian.out", False)
        psi4.geometry(self.geometry)
        #print(self.keywords)
        psi4.set_options(self.keywords)
        H = psi4.hessian(self.method)
        Hess = np.array(H)
        print(Hess.shape)
        self.displacement = Hess[:,self.mode]
        print(self.displacement.shape)
        self.grad = np.zeros(3*self.natoms)

##########################
# Energies and Gradients #
##########################

def doEnergy(geom,pars):
    psi4.geometry(geom)
    E = psi4.energy(pars.level_of_theory)
    return E

def doGrad(geom,pars):
    psi4.geometry(geom)
    Grad, wfn = psi4.gradient(pars.level_of_theory, return_wfn=True)
    e = wfn.energy()
    G = np.array(Grad)
    print(G.shape)
    pars.grad = G
    return G, e

######################################
# Geometry manipulation and printing #
######################################

def geodisplace(geom, displace):
    new_geom = []
    displaced_geom = ''
    symbols = []
    #print(geom)
    lines = geom.split('\n')
    lines.pop(0)
    del lines[-1]
    charge_mult_line = lines[0].split() #Extract the Charge and Multiplicity from top of string
    charge = int(charge_mult_line[0])
    mult = int(charge_mult_line[1])
    for line in lines[1:]: # Loop Over lines that actually contain atoms
        atom_line = line.split()
        #print(atom_line)
        symbols.append(atom_line[0])
        atom_line.pop(0) # Get rid of atom symbol temporarily
        new_geom.append(list(map(float,atom_line)))
    i=0
    for atom in new_geom:
        atom[0] += displace[i]
        atom[1] += displace[i+1]
        atom[2] += displace[i+2]
        i += 3
#TODO (12/15/18) Finish this, all we have to do is transform the geometry from a list to a
#string like the geometry given as an input to this function.
    displaced_geom += "\n%d %d\n" %(charge, mult)
    for i in range(len(symbols)):
        displaced_geom += "%s  " %symbols[i]
        for j in range(3):
            if(j==2):
                #displaced_geom += "   %f   \n" %new_geom[(i + 2*i) + j]
                displaced_geom += "   %f   \n" %new_geom[i][j]
            else:
                #displaced_geom += "   %f   " %new_geom[(i + 2*i) + j]
                displaced_geom += "   %f   " %new_geom[i][j]
    print(displaced_geom)
    return displaced_geom

def printTrj(params, n):
    trajectory_file = open(params.basename+'.trj', 'a')
    trajectory_file.write("%d\npyREX IRC point %d E=%14.7f\n"%(params.natoms,n,params.energy))
    trajectory_file.write("%s\n" %params.geometry)
    trajectory_file.close()

###############
#IRC Functions#
###############

def Morokuma(pars,start=False):
    #Does a cycle in the Morokuma algorithm, returns the energy and MaxGrad of the
    # new point
    if (start):
        #apply the appropriate sign to the displacement and convert to angs.
        pars.displacement=float(pars.direction)*pars.displacement
        #for i in range(len(pars.mass)):
        #   pars.displacement[(3*i)] /= np.sqrt(pars.mass[i])
        #   pars.displacement[(3*i)+1] /= np.sqrt(pars.mass[i])
        #   pars.displacement[(3*i)+2] /= np.sqrt(pars.mass[i])
        pars.displacement=pars.alpha*(pars.displacement/np.linalg.norm(pars.displacement))
        #Eo,MGo = doGrad(pars.geometry, pars)
        #tmpvec0=pars.grad
        #displace geometry following the normal mode
        geo1=geodisplace(pars.geometry,pars.displacement)
        E1,MG1=doGrad(geo1,pars)
        #generate vector D, adapting from eq 6 in J. Chem. Phys. 66, 2153
        tmpvec1=pars.grad
        D=-(tmpvec0/np.linalg.norm(tmpvec0))+(tmpvec1/np.linalg.norm(tmpvec1))
        #D=(pars.displacement/np.linalg.norm(pars.displacement))-(tmpvec/np.linalg.norm(tmpvec))
    else:
        #scale down gradients and calculate the displacement
        pars.displacement=(pars.damp*pars.displacement)-((1.0-pars.damp)*pars.grad)
        for i in range(len(pars.mass)):
            pars.displacement[(3*i)] /= np.sqrt(pars.mass[i])
            pars.displacement[(3*i)+1] /= np.sqrt(pars.mass[i])
            pars.displacement[(3*i)+2] /= np.sqrt(pars.mass[i])
        pars.displacement=pars.maxdispl*(pars.displacement/np.linalg.norm(pars.displacement))
        #displace geometry following the gradient
        geo1=geodisplace(pars.geometry,pars.displacement)
        E1,MG1=doGrad(geo1,pars)
        #generate vector D, adapting from eq 6 in J. Chem. Phys. 66, 2153
        tmpvec=pars.grad
        D=(pars.displacement/np.linalg.norm(pars.displacement))-(tmpvec/np.linalg.norm(tmpvec))
    # find the optimum delta that will assure the new point is at the
    # local minimum
    if (pars.algorithm==1):
        geo2=geodisplace(geo1,pars.delta*D)
        E2=doEnergy(geo2,pars)
        if (E2>E1):
            newdelta=0.5*pars.delta
        else:
            newdelta=2.0*pars.delta
        geo3=geodisplace(geo1,newdelta*D)
        E3=doEnergy(geo3,pars)
        Evals=np.array([E1, E2, E3])
        Deltavals=np.array([0.0,pars.delta,newdelta])
    else:
        geo2=geodisplace(geo1,pars.delta*D)
        E2=doEnergy(geo2,pars)
        delta3=0.5*pars.delta
        geo3=geodisplace(geo1,delta3*D)
        E3=doEnergy(geo3,pars)
        delta4=2.0*pars.delta
        geo4=geodisplace(geo1,delta4*D)
        E4=doEnergy(geo4,pars)
        delta5=-0.5*pars.delta
        geo5=geodisplace(geo1,delta5*D)
        E5=doEnergy(geo5,pars)
        Evals=np.array([E1, E2, E3, E4, E5])
        Deltavals=np.array([0.0,pars.delta,delta3, delta4, delta5])
    deltaFit=np.polyfit(Deltavals,Evals,deg=2)
    opdelta=-(deltaFit[1]/(2.0*deltaFit[0]))
    # update the geometry to the new point, update and return energy and gradients
    pars.geos.append(pars.geometry)
    pars.energies.append(pars.energy)
    pars.geometry=geodisplace(geo1,opdelta*D)
    newE, newMaxGrad = doGrad(pars.geometry,pars)
    pars.energy=newE
    #the following line is just for debugging purposes
    #print(Evals, Deltavals,np.linalg.norm(pars.displacement), opdelta)
    return (newE, newMaxGrad)

def ircdrv(inpname):
    params=ToolKit(inpname)
    params.out.write("""   IRC wrapper for Orca - version 2.0
   by Filipe Teixeira,
   REQUIMTE
   Faculdade de Ciencias da Universidade do Porto

  Using Orca from: %s

  Parameters for this run: \n"""%(params.orcacmd))
    params.printPars()
    #print(params.geometry)
    keep=True
    n=0
    oldE=0.0
    params.out.write("   Pt. %20s %9s %10s\n"%('Energy','RMS Grad.','Damp'))
    params.out.write("------------------------------------------------\n")
    while (keep):
        if (params.restart or (n>0)):
            E,MG=Morokuma(params,False)
        else:
            E,MG=Morokuma(params,True)
        n=n+1
        if(params.autodamp and (n>1)):
            params.damp = params.damp * (0.1/np.log(n))
            if (params.damp<1.0e-5):
                params.damp=0.0
                params.autodamp=False
        params.out.write('@  %3d %20.9f %9.5f %10.2e\n'%(n, E, MG, params.damp))
        if (n>params.npoints):
            keep=False
            printTrj(params,n) #appends newGeo
            params.out.write("----------------------------------------------\n")
            params.out.write("--          IRC CALCULATION ENDED           --\n")
            params.out.write("--             MAXPTS ACHIEVED!             --\n")
            params.out.write("----------------------------------------------\n")
        elif (MG<params.tolerance):
            keep=False
            printTrj(params,n) #appends newGeo
            params.out.write("----------------------------------------------\n")
            params.out.write("--          IRC CALCULATION ENDED           --\n")
            params.out.write("--              TOL ACHIEVED!               --\n")
            params.out.write("----------------------------------------------\n")
        elif (E>oldE):
            keep=False
            #printTrj(params,n) #appends newGeo
            params.out.write("----------------------------------------------\n")
            params.out.write("--             ENERGY INCREASED             --\n")
            params.out.write("--            GEOMETRY  MIGHT BE            --\n")
            params.out.write("--          VERY CLOSE TO A MINIMUM         --\n")
            params.out.write("--                                          --\n")
            params.out.write("--        IRC CALCULATION TERMINATED!       --\n")
            params.out.write("----------------------------------------------\n")
        else:
            printTrj(params,n) #appends newGeo
            oldE=E
    params.out.close()

if __name__=='__main__':
    if(len(sys.argv)!=2):
        print("""IRC4Orca - Version 2.0
An Implementation of Morokuma's IRC method for the Orca ESS Software.
by Filipe Teixeira

Usage: %s file.inp

Please consult https://github.com/teixeirafilipe/irc4orca for more information.

"""%(sys.argv[0]))
        sys.exit(1)
    ircdrv(sys.argv[1])


import numbers
import logging
from copy import deepcopy
from collections import Counter

import scipy.stats
import kappy
import itertools
import numpy as np
import networkx as nx
from pysb import WILD, export, Observable, ComponentSet, Annotation
from pysb.core import as_complex_pattern, ComponentDuplicateNameError
from pysb.pattern import RulePatternMatcher
from indra.explanation.reporting import stmt_from_rule, agent_from_obs
from indra.statements import *
from indra.assemblers.pysb import assembler as pa
from indra.assemblers.pysb.kappa_util import im_json_to_graph
from indra.statements.agent import default_ns_order
from indra.ontology.bio import bio_ontology

from . import ModelChecker, PathResult, NodesContainer
from .model_checker import signed_edges_to_signed_nodes

logger = logging.getLogger(__name__)

try:
    import paths_graph as pg
    has_pg = True
except ImportError:
    pg = None
    has_pg = False
    logger.warning('PathsGraph is not available')

class PysbModelChecker(ModelChecker):
    """Check a PySB model against a set of INDRA statements.

    Parameters
    ----------
    model : pysb.Model
        A PySB model to check.
    statements : Optional[list[indra.statements.Statement]]
        A list of INDRA Statements to check the model against.
    agent_obs: Optional[list[indra.statements.Agent]]
        A list of INDRA Agents in a given state to be observed.
    do_sampling : bool
        Whether to use breadth-first search or weighted sampling to
        generate paths. Default is False (breadth-first search).
    seed : int
        Random seed for sampling (optional, default is None).
    model_stmts : list[indra.statements.Statement]
        A list of INDRA statements used to assemble PySB model.
    nodes_to_agents : dict
        A dictionary mapping nodes of intermediate signed edges graph to INDRA
        agents.

    Attributes
    ----------
    graph : nx.Digraph
        A DiGraph with signed nodes to find paths in.
    """

    def __init__(self, model, statements=None, agent_obs=None,
                 do_sampling=False, seed=None, model_stmts=None,
                 nodes_to_agents=None):
        super().__init__(model, statements, do_sampling, seed, nodes_to_agents)
        if agent_obs:
            self.agent_obs = agent_obs
        else:
            self.agent_obs = []
        mps_to_agents, rules_to_mps = pa.get_grounded_agents(model)
        self.mps_to_agents = mps_to_agents
        self.rules_to_mps = rules_to_mps
        self.model_agents = self.get_model_agents()
        self.model_stmts = model_stmts if model_stmts else []
        # Influence map
        self._im = None
        # Map from statements to associated observables
        self.stmt_to_obs = {}
        # Map from agents to associated observables
        self.agent_to_obs = {}
        # Map between rules and downstream observables
        self.rule_obs_dict = {}
        # Map from observables to agents
        self.obs_to_agents = {}

    def generate_im(self, model):
        """Return a graph representing the influence map generated by Kappa

        Parameters
        ----------
        model : pysb.Model
            The PySB model whose influence map is to be generated

        Returns
        -------
        graph : networkx.MultiDiGraph
            A MultiDiGraph representing the influence map
        """
        kappa = kappy.KappaStd()
        model_str = export.export(model, 'kappa')
        kappa.add_model_string(model_str)
        kappa.project_parse()
        imap = kappa.analyses_influence_map(accuracy='medium')
        graph = im_json_to_graph(imap)
        return graph

    def draw_im(self, fname):
        """Draw and save the influence map in a file.

        Parameters
        ----------
        fname : str
            The name of the file to save the influence map in.
            The extension of the file will determine the file format,
            typically png or pdf.
        """
        im = self.get_im()
        im_agraph = nx.nx_agraph.to_agraph(im)
        im_agraph.draw(fname, prog='dot')

    def get_im(self, force_update=False):
        """Get the influence map for the model, generating it if necessary.

        Parameters
        ----------
        force_update : bool
            Whether to generate the influence map when the function is called.
            If False, returns the previously generated influence map if
            available. Defaults to True.

        Returns
        -------
        networkx MultiDiGraph object containing the influence map.
            The influence map can be rendered as a pdf using the dot layout
            program as follows::

                im_agraph = nx.nx_agraph.to_agraph(influence_map)
                im_agraph.draw('influence_map.pdf', prog='dot')
        """
        if self._im and not force_update:
            return self._im
        if not self.model:
            raise Exception("Cannot get influence map if there is no model.")

        def add_obs_for_agents(main_agent, ref_agents=None):
            if ref_agents:
                all_agents = [main_agent] + ref_agents
            else:
                all_agents = [main_agent]
            ag_to_obj_mps = self.get_all_mps(all_agents, mapping=True)
            if all([not v for v in ag_to_obj_mps.values()]):
                logger.debug('No monomer patterns found in model for agents %s'
                             ', skipping' % all_agents)
                return
            obs_nodes = NodesContainer(main_agent, ref_agents)
            main_obs_set = set()
            ref_obs_set = set()
            for agent in ag_to_obj_mps:
                for obj_mp in ag_to_obj_mps[agent]:
                    obs_name = _monomer_pattern_label(obj_mp) + '_obs'
                    self.obs_to_agents[obs_name] = agent
                    # Add the observable
                    obj_obs = Observable(obs_name, obj_mp, _export=False)
                    if agent.matches(main_agent):
                        main_obs_set.add(obs_name)
                    else:
                        ref_obs_set.add(obs_name)
                    try:
                        self.model.add_component(obj_obs)
                        self.model.add_annotation(
                            Annotation(obs_name, agent.name,
                                       'from_indra_agent'))
                    except ComponentDuplicateNameError as e:
                        pass
            obs_nodes.main_interm = main_obs_set
            obs_nodes.ref_interm = ref_obs_set
            return obs_nodes

        # Create observables for all statements to check, and add to model
        # Remove any existing observables in the model
        self.model.observables = ComponentSet([])
        for stmt in self.statements:
            # Generate observables for Modification statements
            if isinstance(stmt, Modification) or \
               isinstance(stmt, SelfModification):
                # If the statement is a regular Mod, the target is stmt.sub
                if isinstance(stmt, Modification):
                    sub = stmt.sub
                # If it's a SelfMod, the target is stmt.enz
                elif isinstance(stmt, SelfModification):
                    sub = stmt.enz
                # Add the mod for the agent
                if sub is None:
                    self.stmt_to_obs[stmt] = NodesContainer(None)
                else:
                    mod_condition_name = modclass_to_modtype[stmt.__class__]
                    if isinstance(stmt, RemoveModification):
                        mod_condition_name = modtype_to_inverse[
                            mod_condition_name]
                    # Add modification to substrate agent
                    modified_sub = _add_modification_to_agent(
                            sub, mod_condition_name, stmt.residue,
                            stmt.position)
                    # Get all refinements of substrate agent
                    ref_subs = self.get_refinements(modified_sub)
                    obs_nodes = add_obs_for_agents(modified_sub, ref_subs)
                    # Associate this statement with this observable
                    self.stmt_to_obs[stmt] = obs_nodes
            # Generate observables for Activation/Inhibition statements
            elif isinstance(stmt, RegulateActivity):
                if stmt.obj is None:
                    self.stmt_to_obs[stmt] = NodesContainer(None)
                else:
                    # Add activity to object agent
                    regulated_obj = _add_activity_to_agent(
                        stmt.obj, stmt.obj_activity, stmt.is_activation)
                    # Get all refinements of object agent
                    ref_objs = self.get_refinements(stmt.obj)
                    obs_nodes = add_obs_for_agents(regulated_obj, ref_objs)
                    # Associate this statement with this observable
                    self.stmt_to_obs[stmt] = obs_nodes
            elif isinstance(stmt, RegulateAmount):
                if stmt.obj is None:
                    self.stmt_to_obs[stmt] = NodesContainer(None)
                else:
                    # Get all refinements of object agent
                    ref_objs = self.get_refinements(stmt.obj)
                    obs_nodes = add_obs_for_agents(stmt.obj, ref_objs)
                    self.stmt_to_obs[stmt] = obs_nodes
            elif isinstance(stmt, Influence):
                if stmt.obj is None:
                    self.stmt_to_obs[stmt] = NodesContainer(None)
                else:
                    # Get all refinements of object agent
                    ref_objs = self.get_refinements(stmt.obj)
                    concepts = [obj.concept for obj in ref_objs]
                    obs_nodes = add_obs_for_agents(stmt.obj.concept, concepts)
                    self.stmt_to_obs[stmt] = obs_nodes
        # Add observables for each agent
        for ag in self.agent_obs:
            obs_nodes = add_obs_for_agents(ag)
            self.agent_to_obs[ag] = obs_nodes

        logger.info("Generating influence map")
        self._im = self.generate_im(self.model)
        # self._im.is_multigraph = lambda: False
        # Now, for every rule in the model, check if there are any observables
        # downstream; alternatively, for every observable in the model, get a
        # list of rules.
        # We'll need the dictionary to check if nodes are observables
        node_attributes = nx.get_node_attributes(self._im, 'node_type')
        for rule in self.model.rules:
            obs_list = []
            # Get successors of the rule node
            for neighb in self._im.neighbors(rule.name):
                # Check if the node is an observable
                if node_attributes[neighb] != 'variable':
                    continue
                # Get the edge and check the polarity
                edge_sign = _get_edge_sign(self._im, (rule.name, neighb))
                obs_list.append((neighb, edge_sign))
            self.rule_obs_dict[rule.name] = obs_list
        return self._im

    def get_graph(self, prune_im=True, prune_im_degrade=True,
                  prune_im_subj_obj=False, add_namespaces=False,
                  edge_filter_func=None):
        """Get influence map and convert it to a graph with signed nodes."""
        if self.graph:
            return self.graph
        # NOTE edge_filter_func is not currently used in PySB
        im = self.get_im(force_update=True)
        if prune_im:
            self.prune_influence_map()
        if prune_im_degrade:
            self.prune_influence_map_degrade_bind_positive(self.model_stmts)
        if prune_im_subj_obj:
            self.prune_influence_map_subj_obj()
        self.get_nodes_to_agents(add_namespaces=add_namespaces)
        self.graph = signed_edges_to_signed_nodes(
            im, prune_nodes=False, edge_signs={'pos': 1, 'neg': -1})
        return self.graph

    def get_nodes_to_agents(self, add_namespaces=False):
        """Return a dictionary mapping influence map nodes to INDRA agents.

        Parameters
        ----------
        add_namespaces : bool
            Whether to propagate namespaces to node data. Default: False.

        Returns
        -------
        nodes_to_agents : dict
            A dictionary mapping influence map nodes to INDRA agents.
        """
        if self.nodes_to_agents:
            return self.nodes_to_agents

        logger.info('Mapping nodes to agents')
        im = self.get_im()
        nodes_to_agents = {}

        # First map rules to their subject agents
        for rule, mps in self.rules_to_mps.items():
            for mp in mps:
                for ann in self.model.annotations:
                    if ann.subject == rule and ann.object == mp.monomer.name:
                        # We usually want to map rule to subject agent
                        if ann.predicate == 'rule_has_subject':
                            nodes_to_agents[rule] = self.mps_to_agents[mp]

        # Add observables to agents stored earlier
        nodes_to_agents.update(self.obs_to_agents)

        # Optionally propagate namespaces to node data
        if add_namespaces:
            logger.info('Adding namespaces to influence map nodes')
            for n, data in im.nodes(data=True):
                ag = nodes_to_agents.get(n)
                if ag:
                    ns, gr = ag.get_grounding()
                    data['ns'] = ns
        self.nodes_to_agents = nodes_to_agents

    def process_statement(self, stmt):
        self.get_im()
        # Check if this is one of the statement types that we can check
        if not isinstance(stmt, (Modification, RegulateAmount,
                                 RegulateActivity, Influence)):
            logger.info('Statement type %s not handled' %
                        stmt.__class__.__name__)
            return (None, None, 'STATEMENT_TYPE_NOT_HANDLED')
        # Get the polarity for the statement
        if isinstance(stmt, Modification):
            target_polarity = 1 if isinstance(stmt, RemoveModification) else 0
        elif isinstance(stmt, RegulateActivity):
            target_polarity = 0 if stmt.is_activation else 1
        elif isinstance(stmt, RegulateAmount):
            target_polarity = 1 if isinstance(stmt, DecreaseAmount) else 0
        elif isinstance(stmt, Influence):
            target_polarity = 1 if stmt.overall_polarity() == -1 else 0
        # Get the subject and object (works also for Modifications)
        subj, obj = stmt.agent_list()
        # Get a list of monomer patterns matching the subject FIXME Currently
        # this will match rules with the corresponding monomer pattern on it.
        # In future, this statement should (possibly) also match rules in which
        # 1) the agent is in its active form, or 2) the agent is tagged as the
        # enzyme in a rule of the appropriate activity (e.g., a phosphorylation
        # rule) FIXME
        if subj is not None:
            ref_agents = self.get_refinements(subj)
            subj_mps = self.get_all_mps([subj], ignore_activities=True)
            subj_ref_mps = self.get_all_mps(ref_agents, ignore_activities=True)
            if not subj_mps and not subj_ref_mps:
                return (None, None, 'SUBJECT_MONOMERS_NOT_FOUND')
            subj_nodes = NodesContainer(subj, ref_agents)
            meaningful_res_code = None
            # Each subject might produce a different input set and we need to
            # combine them
            for subj_mp in subj_mps:
                inp, res_code = self.process_subject(subj_mp)
                if res_code:
                    meaningful_res_code = res_code
                    continue
                subj_nodes.main_nodes += inp
            for subj_mp in subj_ref_mps:
                inp, res_code = self.process_subject(subj_mp)
                if res_code:
                    meaningful_res_code = res_code
                    continue
                subj_nodes.ref_nodes += inp
            subj_nodes.get_all_nodes()
            if not subj_nodes.all_nodes and meaningful_res_code:
                return (None, None, meaningful_res_code)
        else:
            subj_nodes = NodesContainer(None)
            subj_nodes.all_nodes = None
        # Observables may not be found for an activation since there may be no
        # rule in the model activating the object, and the object may not have
        # an "active" site of the appropriate type
        obs_nodes = self.stmt_to_obs[stmt]
        if obs_nodes is None:
            logger.info("No observables for stmt %s, returning False" % stmt)
            return (None, None, 'OBSERVABLES_NOT_FOUND')
        # Statement object is None
        if obs_nodes.main_agent is None:
            # Cannot check modifications in this case
            if isinstance(stmt, Modification):
                return (None, None, 'STATEMENT_TYPE_NOT_HANDLED')
            obs_nodes.all_nodes = None
        else:
            obs_nodes.main_nodes = [
                (obs, target_polarity) for obs in obs_nodes.main_interm]
            obs_nodes.ref_nodes = [
                (obs, target_polarity) for obs in obs_nodes.ref_interm]
            obs_nodes.get_all_nodes()
        result_code = None
        return subj_nodes, obs_nodes, result_code

    def process_subject(self, subj_mp):
        if subj_mp is None:
            input_set_signed = None
        else:
            input_rule_set = self._get_input_rules(subj_mp)
            if not input_rule_set:
                logger.info('Input rules not found for %s' % subj_mp)
                return (None, 'INPUT_RULES_NOT_FOUND')
            input_set_signed = {(rule, 0) for rule in input_rule_set}
        return input_set_signed, None

    def get_model_agents(self):
        return set(self.mps_to_agents.values())

    def get_refinements(self, agent):
        """Return a list of refinement agents that are part of the model."""
        agents = set()
        for ag in self.model_agents:
            if not ag.matches(agent) and ag.refinement_of(agent, bio_ontology):
                agents.add(ag)
        return list(agents)

    def get_all_mps(self, agents, ignore_activities=False, mapping=False):
        """Get a list of all monomer patterns for a list of agents."""
        ag_to_mps = {}
        mps = []
        for ag in agents:
            ag_mps = list(pa.grounded_monomer_patterns(
                self.model, ag, ignore_activities=ignore_activities))
            if ag_mps:
                ag_to_mps[ag] = ag_mps
                mps += ag_mps
        if mapping:
            return ag_to_mps
        return set(mps)

    def _get_input_rules(self, subj_mp):
        if subj_mp is None:
            raise ValueError("Cannot take None as an argument for subj_mp.")
        input_rules = _match_lhs(subj_mp, self.model.rules)
        logger.debug('Found %s input rules matching %s' %
                     (len(input_rules), str(subj_mp)))
        # Filter to include only rules where the subj_mp is actually the
        # subject (i.e., don't pick up upstream rules where the subject
        # is itself a substrate/object)
        # FIXME: Note that this will eliminate rules where the subject
        # being checked is included on the left hand side as
        # a bound condition rather than as an enzyme.
        subj_rules = pa.rules_with_annotation(self.model,
                                              subj_mp.monomer.name,
                                              'rule_has_subject')
        logger.debug('%d rules with %s as subject' %
                     (len(subj_rules), subj_mp.monomer.name))
        input_rule_set = set([r.name for r in input_rules]).intersection(
                             set([r.name for r in subj_rules]))
        logger.debug('Final input rule set contains %d rules' %
                     len(input_rule_set))
        return input_rule_set

    def _sample_paths(self, input_rule_set, obs_name, target_polarity,
                      max_paths=1, max_path_length=5):
        if max_paths == 0:
            raise ValueError("max_paths cannot be 0 for path sampling.")
        if not has_pg:
            raise ImportError("Paths Graph is not imported")
        # Convert path polarity representation from 0/1 to 1/-1

        def convert_polarities(path_list):
            return [tuple((n[0], 0 if n[1] > 0 else 1) for n in path)
                    for path in path_list]

        pg_polarity = 0 if target_polarity > 0 else 1
        nx_graph = self._im_to_signed_digraph(self.get_im())
        # Add edges from dummy node to input rules
        source_node = 'SOURCE_NODE'
        for rule in input_rule_set:
            nx_graph.add_edge(source_node, rule, sign=0)
        # -------------------------------------------------
        # Create combined paths_graph
        f_level, b_level = pg.get_reachable_sets(nx_graph, source_node,
                                                 obs_name, max_path_length,
                                                 signed=True)
        pg_list = []
        for path_length in range(1, max_path_length+1):
            cfpg = pg.CFPG.from_graph(
                    nx_graph, source_node, obs_name, path_length, f_level,
                    b_level, signed=True, target_polarity=pg_polarity)
            pg_list.append(cfpg)
        combined_pg = pg.CombinedCFPG(pg_list)
        # Make sure the combined paths graph is not empty
        if not combined_pg.graph:
            pr = PathResult(
                False, 'NO_PATHS_FOUND', max_paths, max_path_length)
            pr.path_metrics = None
            pr.paths = []
            return pr

        # Get a dict of rule objects
        rule_obj_dict = {}
        for ann in self.model.annotations:
            if ann.predicate == 'rule_has_object':
                rule_obj_dict[ann.subject] = ann.object

        # Get monomer initial conditions
        ic_dict = {}
        for mon in self.model.monomers:
            # FIXME: A hack that depends on the _0 convention
            ic_name = '%s_0' % mon.name
            # TODO: Wrap this in try/except?
            ic_param = self.model.parameters[ic_name]
            ic_value = ic_param.value
            ic_dict[mon.name] = ic_value

        # Set weights in PG based on model initial conditions
        for cur_node in combined_pg.graph.nodes():
            edge_weights = {}
            rule_obj_list = []
            edge_weights_by_gene = {}
            for u, v in combined_pg.graph.out_edges(cur_node):
                v_rule = v[1][0]
                # Get the object of the rule (a monomer name)
                rule_obj = rule_obj_dict.get(v_rule)
                if rule_obj:
                    # Add to list so we can count instances by gene
                    rule_obj_list.append(rule_obj)
                    # Get the abundance of rule object from the initial
                    # conditions
                    # TODO: Wrap in try/except?
                    ic_value = ic_dict[rule_obj]
                else:
                    ic_value = 1.0
                edge_weights[(u, v)] = ic_value
                edge_weights_by_gene[rule_obj] = ic_value
            # Get frequency of different rule objects
            rule_obj_ctr = Counter(rule_obj_list)
            # Normalize results by weight sum and gene frequency at this level
            edge_weight_sum = sum(edge_weights_by_gene.values())
            edge_weights_norm = {}
            for e, v in edge_weights.items():
                v_rule = e[1][1][0]
                rule_obj = rule_obj_dict.get(v_rule)
                if rule_obj:
                    rule_obj_count = rule_obj_ctr[rule_obj]
                else:
                    rule_obj_count = 1
                edge_weights_norm[e] = ((v / float(edge_weight_sum)) /
                                        float(rule_obj_count))
            # Add edge weights to paths graph
            nx.set_edge_attributes(combined_pg.graph, name='weight',
                                   values=edge_weights_norm)

        # Sample from the combined CFPG
        paths = combined_pg.sample_paths(max_paths)
        # -------------------------------------------------
        if paths:
            pr = PathResult(True, 'PATHS_FOUND', max_paths, max_path_length)
            pr.path_metrics = None
            # Convert path polarity representation from 0/1 to 1/-1
            pr.paths = convert_polarities(paths)
            # Strip off the SOURCE_NODE prefix
            pr.paths = [p[1:] for p in pr.paths]
        else:
            assert False
            pr = PathResult(
                False, 'NO_PATHS_FOUND', max_paths, max_path_length)
            pr.path_metrics = None
            pr.paths = []
        return pr

    def score_paths(self, paths, agents_values, loss_of_function=False,
                    sigma=0.15, include_final_node=False):
        """Return scores associated with a given set of paths.

        Parameters
        ----------
        paths : list[list[tuple[str, int]]]
            A list of paths obtained from path finding. Each path is a list
            of tuples (which are edges in the path), with the first element
            of the tuple the name of a rule, and the second element its
            polarity in the path.
        agents_values : dict[indra.statements.Agent, float]
            A dictionary of INDRA Agents and their corresponding measured
            value in a given experimental condition.
        loss_of_function : Optional[boolean]
            If True, flip the polarity of the path. For instance, if the effect
            of an inhibitory drug is explained, set this to True.
            Default: False
        sigma : Optional[float]
            The estimated standard deviation for the normally distributed
            measurement error in the observation model used to score paths
            with respect to data. Default: 0.15
        include_final_node : Optional[boolean]
            Determines whether the final node of the path is included in the
            score. Default: False
        """
        obs_model = lambda x: scipy.stats.norm(x, sigma)
        # Build up dict mapping observables to values
        obs_dict = {}
        for ag, val in agents_values.items():
            obs_list = self.agent_to_obs[ag]
            if obs_list is not None:
                for obs in obs_list:
                    obs_dict[obs] = val
        # For every path...
        path_scores = []
        for path in paths:
            logger.info('------')
            logger.info("Scoring path:")
            logger.info(path)
            # Look at every node in the path, excluding the final
            # observable...
            path_score = 0
            last_path_node_index = -1 if include_final_node else -2
            for node, sign in path[:last_path_node_index]:
                # ...and for each node check the sign to see if it matches the
                # data. So the first thing is to look at what's downstream
                # of the rule
                # affected_obs is a list of observable names alogn
                for affected_obs, rule_obs_sign in self.rule_obs_dict[node]:
                    flip_polarity = -1 if loss_of_function else 1
                    pred_sign = sign * rule_obs_sign * flip_polarity
                    # Check to see if this observable is in the data
                    logger.info('%s %s: effect %s %s' %
                                (node, sign, affected_obs, pred_sign))
                    measured_val = obs_dict.get(affected_obs)
                    if measured_val:
                        # For negative predictions use CDF (prob that given
                        # measured value, true value lies below 0)
                        if pred_sign <= 0:
                            prob_correct = obs_model(measured_val).logcdf(0)
                        # For positive predictions, use log survival function
                        # (SF = 1 - CDF, i.e., prob that true value is
                        # above 0)
                        else:
                            prob_correct = obs_model(measured_val).logsf(0)
                        logger.info('Actual: %s, Log Probability: %s' %
                                    (measured_val, prob_correct))
                        path_score += prob_correct
                if not self.rule_obs_dict[node]:
                    logger.info('%s %s' % (node, sign))
                    prob_correct = obs_model(0).logcdf(0)
                    logger.info('Unmeasured node, Log Probability: %s' %
                                (prob_correct))
                    path_score += prob_correct
            # Normalized path
            # path_score = path_score / len(path)
            logger.info("Path score: %s" % path_score)
            path_scores.append(path_score)
        path_tuples = list(zip(paths, path_scores))
        # Sort first by path length
        sorted_by_length = sorted(path_tuples, key=lambda x: len(x[0]))
        # Sort by probability; sort in reverse order to large values
        # (higher probabilities) are ranked higher
        scored_paths = sorted(sorted_by_length, key=lambda x: x[1],
                              reverse=True)
        return scored_paths

    def prune_influence_map(self):
        """Remove edges between rules causing problematic non-transitivity.

        First, all self-loops are removed. After this initial step, edges are
        removed between rules when they share *all* child nodes except for each
        other; that is, they have a mutual relationship with each other and
        share all of the same children.

        Note that edges must be removed in batch at the end to prevent edge
        removal from affecting the lists of rule children during the comparison
        process.
        """
        im = self.get_im()

        # First, remove all self-loops
        logger.info('Removing self loops')
        edges_to_remove = []
        for e in im.edges():
            if e[0] == e[1]:
                logger.info('Removing self loop: %s', e)
                edges_to_remove.append((e[0], e[1]))
        # Now remove all the edges to be removed with a single call
        im.remove_edges_from(edges_to_remove)

        # Remove parameter nodes from influence map
        remove_im_params(self.model, im)

        # Now compare nodes pairwise and look for overlap between child nodes
        logger.info('Get successors of each node')
        succ_dict = {}
        for node in im.nodes():
            succ_dict[node] = set(im.successors(node))
        # Sort and then group nodes by number of successors
        logger.info('Compare combinations of successors')
        group_key_fun = lambda x: len(succ_dict[x])
        nodes_sorted = sorted(im.nodes(), key=group_key_fun)
        groups = itertools.groupby(nodes_sorted, key=group_key_fun)
        # Now iterate over each group and then construct combinations
        # within the group to check for shared sucessors
        edges_to_remove = []
        for gix, group in groups:
            combos = itertools.combinations(group, 2)
            for ix, (p1, p2) in enumerate(combos):
                # Children are identical except for mutual relationship
                if succ_dict[p1].difference(succ_dict[p2]) == set([p2]) and \
                   succ_dict[p2].difference(succ_dict[p1]) == set([p1]):
                    for u, v in ((p1, p2), (p2, p1)):
                        edges_to_remove.append((u, v))
                        logger.debug('Will remove edge (%s, %s)', u, v)
        logger.info('Removing %d edges from influence map' %
                    len(edges_to_remove))
        # Now remove all the edges to be removed with a single call
        im.remove_edges_from(edges_to_remove)

    def prune_influence_map_subj_obj(self):
        """Prune influence map to include only edges where the object of the
        upstream rule matches the subject of the downstream rule."""
        def get_rule_info(r):
            result = {}
            for ann in self.model.annotations:
                if ann.subject == r:
                    if ann.predicate == 'rule_has_subject':
                        result['subject'] = ann.object
                    elif ann.predicate == 'rule_has_object':
                        result['object'] = ann.object
            return result
        im = self.get_im()
        rules = im.nodes()
        edges_to_prune = []
        for r1, r2 in itertools.permutations(rules, 2):
            if (r1, r2) not in im.edges():
                continue
            r1_info = get_rule_info(r1)
            r2_info = get_rule_info(r2)
            if 'object' not in r1_info or 'subject' not in r2_info:
                continue
            if r1_info['object'] != r2_info['subject']:
                logger.info("Removing edge %s --> %s" % (r1, r2))
                edges_to_prune.append((r1, r2))
        logger.info('Removing %d edges from influence map' %
                    len(edges_to_prune))
        im.remove_edges_from(edges_to_prune)

    def prune_influence_map_degrade_bind_positive(self, model_stmts):
        """Prune positive edges between X degrading and X forming a
        complex with Y."""
        im = self.get_im()
        edges_to_prune = []
        for r1, r2, data in im.edges(data=True):
            s1 = stmt_from_rule(r1, self.model, model_stmts)
            s2 = stmt_from_rule(r2, self.model, model_stmts)
            # Make sure this is a degradation/binding combo
            s1_is_degrad = (s1 and isinstance(s1, DecreaseAmount))
            s2_is_bind = (s2 and isinstance(s2, Complex) and 'bind' in r2)
            if not s1_is_degrad or not s2_is_bind:
                continue
            # Make sure what is degraded is part of the complex
            if s1.obj.name not in [m.name for m in s2.members]:
                continue
            # Make sure we're dealing with a positive influence
            if data['sign'] == 1:
                edges_to_prune.append((r1, r2))
        logger.info('Removing %d edges from influence map' %
                    len(edges_to_prune))
        im.remove_edges_from(edges_to_prune)

    def _im_to_signed_digraph(self, im):
        edges = []
        for e in im.edges():
            edge_sign = _get_edge_sign(im, e)
            polarity = 0 if edge_sign > 0 else 1
            edges.append((e[0], e[1], {'sign': polarity}))
        dg = nx.DiGraph()
        dg.add_edges_from(edges)
        return dg

def _find_sources_sample(im, target, sources, polarity, rule_obs_dict,
                         agent_to_obs, agents_values):
    # Build up dict mapping observables to values
    obs_dict = {}
    for ag, val in agents_values.items():
        obs_list = agent_to_obs[ag]
        for obs in obs_list:
            obs_dict[obs] = val

    sigma = 0.2

    def obs_model(x):
        return scipy.stats.norm(x, sigma)

    def _sample_pred(im, target, rule_obs_dict, obs_model):
        preds = list(_get_signed_predecessors(im, target, 1))
        if not preds:
            return None
        pred_scores = []
        for pred, sign in preds:
            pred_score = 0
            for affected_obs, rule_obs_sign in rule_obs_dict[pred]:
                pred_sign = sign * rule_obs_sign
                # Check to see if this observable is in the data
                logger.info('%s %s: effect %s %s' %
                            (pred, sign, affected_obs, pred_sign))
                measured_val = obs_dict.get(affected_obs)
                if measured_val:
                    logger.info('Actual: %s' % measured_val)
                    # The tail probability of the real value being above 1
                    tail_prob = obs_model(measured_val).cdf(1)
                    pred_score += (tail_prob if pred_sign == 1 else
                                   1-tail_prob)
            pred_scores.append(pred_score)
        # Normalize scores
        pred_scores = np.array(pred_scores) / np.sum(pred_scores)
        pred_idx = np.random.choice(range(len(preds)), p=pred_scores)
        pred = preds[pred_idx]
        return pred

    preds = []
    for i in range(100):
        pred = _sample_pred(im, target, rule_obs_dict, obs_model)
        preds.append(pred[0])

def remove_im_params(model, im):
    """Remove parameter nodes from the influence map.

    Parameters
    ----------
    model : pysb.core.Model
        PySB model.
    im : networkx.MultiDiGraph
        Influence map.

    Returns
    -------
    networkx.MultiDiGraph
        Influence map with the parameter nodes removed.
    """
    for param in model.parameters:
        # If the node doesn't exist e.g., it may have already been removed),
        # skip over the parameter without error
        try:
            im.remove_node(param.name)
        except:
            pass

def _get_signed_predecessors(im, node, polarity):
    """Get upstream nodes in the influence map.
    Return the upstream nodes along with the overall polarity of the path
    to that node by account for the polarity of the path to the given node
    and the polarity of the edge between the given node and its immediate
    predecessors.
    Parameters
    ----------
    im : networkx.MultiDiGraph
        Graph containing the influence map.
    node : str
        The node (rule name) in the influence map to get predecessors (upstream
        nodes) for.
    polarity : int
        Polarity of the overall path to the given node.
    Returns
    -------
    generator of tuples, (node, polarity)
        Each tuple returned contains two elements, a node (string) and the
        polarity of the overall path (int) to that node.
    """
    signed_pred_list = []
    for pred in im.predecessors(node):
        pred_edge = (pred, node)
        yield (pred, _get_edge_sign(im, pred_edge) * polarity)

def _get_edge_sign(im, edge):
    """Get the polarity of the influence by examining the edge sign."""
    edge_data = im[edge[0]][edge[1]]
    # Handle possible multiple edges between nodes
    signs = list(set([v['sign'] for v in edge_data.values()
                                  if v.get('sign')]))
    if len(signs) > 1:
        logger.warning("Edge %s has conflicting polarities; choosing "
                       "positive polarity by default" % str(edge))
        sign = 1
    else:
        sign = signs[0]
    if sign is None:
        raise Exception('No sign attribute for edge.')
    elif abs(sign) == 1:
        return sign
    else:
        raise Exception('Unexpected edge sign: %s' % edge.attr['sign'])

def _add_modification_to_agent(agent, mod_type, residue, position):
    """Add a modification condition to an Agent."""
    new_mod = ModCondition(mod_type, residue, position)
    # Check if this modification already exists
    for old_mod in agent.mods:
        if old_mod.equals(new_mod):
            return agent
    new_agent = deepcopy(agent)
    new_agent.mods.append(new_mod)
    return new_agent

def _add_activity_to_agent(agent, act_type, is_active):
    # Default to active, and return polarity if it's an inhibition
    new_act = ActivityCondition(act_type, True)
    # Check if this state already exists
    if agent.activity is not None and agent.activity.equals(new_act):
        return agent
    new_agent = deepcopy(agent)
    new_agent.activity = new_act
    polarity = 1 if is_active else -1
    return new_agent

def _match_lhs(cp, rules):
    """Get rules with a left-hand side matching the given ComplexPattern."""
    rule_matches = []
    for rule in rules:
        reactant_pattern = rule.rule_expression.reactant_pattern
        for rule_cp in reactant_pattern.complex_patterns:
            if _cp_embeds_into(rule_cp, cp):
                rule_matches.append(rule)
                break
    return rule_matches

def _cp_embeds_into(cp1, cp2):
    """Check that any state in ComplexPattern2 is matched in ComplexPattern1.
    """
    # Check that any state in cp2 is matched in cp1
    # If the thing we're matching to is just a monomer pattern, that makes
    # things easier--we just need to find the corresponding monomer pattern
    # in cp1
    if cp1 is None or cp2 is None:
        return False
    cp1 = as_complex_pattern(cp1)
    cp2 = as_complex_pattern(cp2)
    if len(cp2.monomer_patterns) == 1:
        mp2 = cp2.monomer_patterns[0]
        # Iterate over the monomer patterns in cp1 and see if there is one
        # that has the same name
        for mp1 in cp1.monomer_patterns:
            if _mp_embeds_into(mp1, mp2):
                return True
    return False

def _mp_embeds_into(mp1, mp2):
    """Check that conditions in MonomerPattern2 are met in MonomerPattern1."""
    sc_matches = []
    if mp1.monomer.name != mp2.monomer.name:
        return False
    # Check that all conditions in mp2 are met in mp1
    for site_name, site_state in mp2.site_conditions.items():
        if site_name not in mp1.site_conditions or \
           site_state != mp1.site_conditions[site_name]:
            return False
    return True

def _monomer_pattern_label(mp):
    """Return a string label for a MonomerPattern."""
    site_strs = []
    for site, cond in mp.site_conditions.items():
        if isinstance(cond, tuple) or isinstance(cond, list):
            assert len(cond) == 2
            if cond[1] == WILD:
                site_str = '%s_%s' % (site, cond[0])
            else:
                site_str = '%s_%s%s' % (site, cond[0], cond[1])
        elif isinstance(cond, numbers.Real):
            continue
        else:
            site_str = '%s_%s' % (site, cond)
        site_strs.append(site_str)
    return '%s_%s' % (mp.monomer.name, '_'.join(site_strs))

import numpy,import scipy,import networkx


# Copyright 2020 Alibaba Group Holding Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# =============================================================================
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
import tensorflow as tf
from graphlearn.python.nn.tf.module import Module
from graphlearn.python.nn.tf.layers.linear_layer import LinearLayer

class EgoSAGELayerGroup(Module):
  def __init__(self, layers):
    super(EgoSAGELayerGroup, self).__init__()
    self.layers = layers

  def append(self, layer):
    self.layers.append(layer)

  def forward(self, x_list, expands):
    """ Compute node embeddings based on GraphSAGE.

    x_list = [nodes, hop1, hop2, ... , hopK-1, hopK]
               |   /  |   /  |   /        |    /
               |  /   |  /   |  /         |   /
               | /    | /    | /          |  /
    output = [ret0,  ret1, ret2, ... , retK-1]

    Args:
      x_list: A list of tensors, representing input nodes and their neighbors.
        If len(x_list) is K+1, that means x_list[0], x_list[1], ... , x_list[K]
        are the hidden embedding values at each hop. Tensors in x_list[i] are
        the neighbors of that in x_list[i-1]. In this layer, we will do
        convolution for each adjencent pair and return a list with length K.

        The shape of x_list[0] is `[n, input_dim_0]`, and the shape of x_list[i]
        is `[n * k_1 * ... * k_i, input_dim_i]`, where `k_i` means the neighbor
        count of each node at (i-1)th hop. Each `input_dim_i` must match with
        `input_dim` parameter when layer construction.

      expands: An integer list of neighbor count at each hop. For the above
        x_list, expands = [k_1, k_2, ... , k_K]

    Return:
      A list with K tensors, and the ith shape is
      `[n * k_1 * ... * k_i, output_dim]`.
    """
    assert len(self.layers) == (len(x_list) - 1)
    assert len(self.layers) == len(expands)

    rets = []
    for i in range(1, len(x_list)):
      x = x_list[i - 1]
      neighbors = x_list[i]
      ret = self.layers[i - 1](x, neighbors, expands[i - 1])
      rets.append(ret)
    return rets

class EgoSAGELayer(Module):
  """ GraphSAGE. https://arxiv.org/abs/1706.02216.

  Args:
    name: A string, layer name.
    input_dim: An integer or a two elements tuple. Dimension of input features.
      If an integer, nodes and neighbors share the same dimension.
      If an tuple, the two elements represent the dimensions of node features
      and neighbor features.
      Usually, different dimensions happen in the heterogeneous graph.
    output_dim: An integer, dimension of the output embeddings. Both the node
      features and neighbor features will be encoded into the same dimension,
      and then do some combination.
    agg_type: A string, how to merge neighbor values. The optional values are
      'mean', 'sum', 'max'.
    com_type: A string, how to combine neighbors to self. The optional values
      are 'add', 'concat'.
    use_bias: A boolean, whether add bias after computation.
  """

  def __init__(self,
               name,
               input_dim,
               output_dim,
               agg_type="mean",
               com_type="add",
               use_bias=False,
               parameter_share=False,
               **kwargs):
    super(EgoSAGELayer, self).__init__()
    assert agg_type in {"mean", "sum", "max", "gcn"}
    assert com_type in {"add", "concat", "gcn"}

    self.agg_type = agg_type
    self.com_type = com_type
    self.out_dim = output_dim

    if isinstance(input_dim, list) or isinstance(input_dim, tuple):
      self.in_dim = input_dim
      assert len(self.in_dim) == 2
    else:
      self.in_dim = [input_dim, input_dim]

    with tf.variable_scope("ego_sage_layer_" + name, reuse=tf.AUTO_REUSE):
      self.linears = self.add_transform_layers(parameter_share, use_bias)

  def add_transform_layers(self, parameter_share, use_bias):
    layers = []
    if self.com_type == "concat":
      dim = self.in_dim[0] + self.in_dim[1]
      layers.append(LinearLayer("trans_nodes", dim, self.out_dim, use_bias))
    elif parameter_share and self.in_dim[0] == self.in_dim[1]:
      layer = LinearLayer("trans_nodes", self.in_dim[0], self.out_dim, use_bias)
      layers.append(layer)
      layers.append(layer)
    else:
      layers.append(
          LinearLayer("trans_nodes", self.in_dim[0], self.out_dim, use_bias))
      layers.append(
          LinearLayer("trans_nbrs", self.in_dim[1], self.out_dim, use_bias))
    return layers

  def forward(self, x, neighbor, expand):
    """ Compute node embeddings based on GraphSAGE.
    Args:
      x: A float tensor with shape = [batch_size, input_dim].
      neighbor: A float tensor with shape = [batch_size * expand, input_dim].
      expand: An integer, the neighbor count.

    Return:
      A float tensor with shape=[batch_size, output_dim].
    """
    # aggregate neighbors at each hop
    agg_func = self.aggregator()
    agg_info = agg_func(neighbor, expand, self.in_dim[1])

    # combine self info with aggregated neighbors
    comb_func = self.combiner()
    return comb_func(x, agg_info)

  def aggregator(self):
    func_name = self.agg_type + "_agg"
    if not hasattr(self, func_name):
      raise TypeError("Unsupported agg_type: " + self.agg_type)
    return getattr(self, func_name)

  def sum_agg(self, x, expand, dim):
    t = tf.reshape(x, [-1, expand, dim])
    return tf.math.reduce_sum(t, axis=1)

  def mean_agg(self, x, expand, dim):
    t = tf.reshape(x, [-1, expand, dim])
    return tf.math.reduce_mean(t, axis=1)

  def gcn_agg(self, x, expand, dim):
    return tf.reshape(x, [-1, expand, dim])

  def max_agg(self, x, expand, dim):
    t = tf.reshape(x, [-1, expand, dim])
    return tf.math.reduce_max(t, axis=1)

  def combiner(self):
    func_name = self.com_type + "_comb"
    if not hasattr(self, func_name):
      raise TypeError("Unsupported com_type: " + self.com_type)
    return getattr(self, func_name)

  def add_comb(self, x, neighbors):
    current = self.linears[0](x)
    nbr_info = self.linears[1](neighbors)
    return tf.add(current, nbr_info)

  def concat_comb(self, x, neighbors):
    tmp = tf.concat([x, neighbors], axis=1)
    ret = self.linears[0](tmp)
    return ret

  def gcn_comb(self, x, neighbors):
    x = tf.reduce_mean(tf.concat(
          [neighbors, tf.expand_dims(x, axis=1)], axis=1), axis=1)
    return self.linears[0](x)


# This file is part of QuTiP: Quantum Toolbox in Python.
#
#    Copyright (c) 2011 and later, Paul D. Nation and Robert J. Johansson.
#    All rights reserved.
#
#    Redistribution and use in source and binary forms, with or without
#    modification, are permitted provided that the following conditions are
#    met:
#
#    1. Redistributions of source code must retain the above copyright notice,
#       this list of conditions and the following disclaimer.
#
#    2. Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#
#    3. Neither the name of the QuTiP: Quantum Toolbox in Python nor the names
#       of its contributors may be used to endorse or promote products derived
#       from this software without specific prior written permission.
#
#    THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
#    "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
#    LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
#    PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
#    HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
#    SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
#    LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
#    DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
#    THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
#    (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
#    OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
###############################################################################
"""
This module contains a collection of routines for operating on sparse
matrices on the scipy.sparse formats, for use internally by other modules
throughout QuTiP.
"""

__all__ = ['sp_fro_norm', 'sp_inf_norm', 'sp_L2_norm', 'sp_max_norm',
           'sp_one_norm', 'sp_reshape', 'sp_eigs', 'sp_expm', 'sp_permute',
           'sp_reverse_permute', 'sp_bandwidth', 'sp_profile']

import scipy.sparse as sp
import scipy.sparse.linalg as spla
import numpy as np
import scipy.linalg as la
from scipy.linalg.blas import get_blas_funcs
_dznrm2 = get_blas_funcs("znrm2")
from qutip.cy.sparse_utils import (_sparse_profile, _sparse_permute,
                                   _sparse_reverse_permute, _sparse_bandwidth)
from qutip.settings import debug

import qutip.logging
logger = qutip.logging.get_logger()

if debug:
    import inspect

def sp_fro_norm(data):
    """
    Frobius norm for sparse matrix
    """
    out = np.sum(np.abs(data.data)**2)
    return np.sqrt(out)

def sp_inf_norm(data):
    """
    Infinity norm for sparse matrix
    """
    return np.max([np.sum(np.abs(data.getrow(k).data))
                   for k in range(data.shape[0])])

def sp_L2_norm(data):
    """
    L2 norm sparse vector
    """
    if 1 not in data.shape:
        raise TypeError("Use L2-norm only for vectors.")

    if len(data.data):
        return _dznrm2(data.data)
    else:
        return 0

def sp_max_norm(data):
    """
    Max norm for sparse matrix
    """
    return np.max(np.abs(data.data)) if any(data.data) else 0

def sp_one_norm(data):
    """
    One norm for sparse matrix
    """
    return np.max(np.array([np.sum(np.abs((data.getcol(k).data)))
                            for k in range(data.shape[1])]))

def sp_reshape(A, shape, format='csr'):
    """
    Reshapes a sparse matrix.

    Parameters
    ----------
    A : sparse_matrix
        Input matrix in any format
    shape : list/tuple
        Desired shape of new matrix
    format : string {'csr','coo','csc','lil'}
        Optional string indicating desired output format

    Returns
    -------
    B : csr_matrix
        Reshaped sparse matrix

    References
    ----------

        http://stackoverflow.com/questions/16511879/reshape-sparse-matrix-efficiently-python-scipy-0-12

    """
    if not hasattr(shape, '__len__') or len(shape) != 2:
        raise ValueError('Shape must be a list of two integers')

    C = A.tocoo()
    nrows, ncols = C.shape
    size = nrows * ncols
    new_size = shape[0] * shape[1]

    if new_size != size:
        raise ValueError('Total size of new array must be unchanged.')

    flat_indices = ncols * C.row + C.col
    new_row, new_col = divmod(flat_indices, shape[1])
    B = sp.coo_matrix((C.data, (new_row, new_col)), shape=shape)

    if format == 'csr':
        return B.tocsr()
    elif format == 'coo':
        return B
    elif format == 'csc':
        return B.tocsc()
    elif format == 'lil':
        return B.tolil()
    else:
        raise ValueError('Return format not valid.')

def _dense_eigs(data, isherm, vecs, N, eigvals, num_large, num_small):
    """
    Internal functions for computing eigenvalues and eigenstates for a dense
    matrix.
    """
    if debug:
        logger.debug(inspect.stack()[0][3] + ": vectors = " + str(vecs))

    evecs = None

    if vecs:
        if isherm:
            if eigvals == 0:
                evals, evecs = la.eigh(data)
            else:
                if num_small > 0:
                    evals, evecs = la.eigh(
                        data, eigvals=[0, num_small - 1])
                if num_large > 0:
                    evals, evecs = la.eigh(
                        data, eigvals=[N - num_large, N - 1])
        else:
            evals, evecs = la.eig(data)
    else:
        if isherm:
            if eigvals == 0:
                evals = la.eigvalsh(data)
            else:
                if num_small > 0:
                    evals = la.eigvalsh(data, eigvals=[0, num_small - 1])
                if num_large > 0:
                    evals = la.eigvalsh(data, eigvals=[N - num_large, N - 1])
        else:
            evals = la.eigvals(data)

    _zipped = list(zip(evals, range(len(evals))))
    _zipped.sort()
    evals, perm = list(zip(*_zipped))

    if vecs:
        evecs = np.array([evecs[:, k] for k in perm])

    if not isherm and eigvals > 0:
        if vecs:
            if num_small > 0:
                evals, evecs = evals[:num_small], evecs[:num_small]
            elif num_large > 0:
                evals, evecs = evals[(N - num_large):], evecs[(N - num_large):]
        else:
            if num_small > 0:
                evals = evals[:num_small]
            elif num_large > 0:
                evals = evals[(N - num_large):]

    return np.array(evals), np.array(evecs)

def _sp_eigs(data, isherm, vecs, N, eigvals, num_large, num_small, tol,
             maxiter):
    """
    Internal functions for computing eigenvalues and eigenstates for a sparse
    matrix.
    """
    if debug:
        print(inspect.stack()[0][3] + ": vectors = " + str(vecs))

    big_vals = np.array([])
    small_vals = np.array([])
    evecs = None

    remove_one = False
    if eigvals == (N - 1):
        # calculate all eigenvalues and remove one at output if using sparse
        eigvals = 0
        num_small = int(np.ceil(N / 2.0))
        num_large = N - num_small
        remove_one = True

    if vecs:
        if isherm:
            if num_large > 0:
                big_vals, big_vecs = sp.linalg.eigsh(data, k=num_large,
                                                     which='LA', tol=tol,
                                                     maxiter=maxiter)
                big_vecs = sp.csr_matrix(big_vecs, dtype=complex)
            if num_small > 0:
                small_vals, small_vecs = sp.linalg.eigsh(
                    data, k=num_small, which='SA',
                    tol=tol, maxiter=maxiter)

        else:
            if num_large > 0:
                big_vals, big_vecs = sp.linalg.eigs(data, k=num_large,
                                                    which='LR', tol=tol,
                                                    maxiter=maxiter)
                big_vecs = sp.csr_matrix(big_vecs, dtype=complex)
            if num_small > 0:
                small_vals, small_vecs = sp.linalg.eigs(
                    data, k=num_small, which='SR',
                    tol=tol, maxiter=maxiter)

        if num_large != 0 and num_small != 0:
            evecs = sp.hstack([small_vecs, big_vecs], format='csr')
        elif num_large != 0 and num_small == 0:
            evecs = big_vecs
        elif num_large == 0 and num_small != 0:
            evecs = small_vecs
    else:
        if isherm:
            if num_large > 0:
                big_vals = sp.linalg.eigsh(
                    data, k=num_large, which='LA',
                    return_eigenvectors=False, tol=tol, maxiter=maxiter)
            if num_small > 0:
                small_vals = sp.linalg.eigsh(
                    data, k=num_small, which='SA',
                    return_eigenvectors=False, tol=tol, maxiter=maxiter)
        else:
            if num_large > 0:
                big_vals = sp.linalg.eigs(
                    data, k=num_large, which='LR',
                    return_eigenvectors=False, tol=tol, maxiter=maxiter)
            if num_small > 0:
                small_vals = sp.linalg.eigs(
                    data, k=num_small, which='SR',
                    return_eigenvectors=False, tol=tol, maxiter=maxiter)

    evals = np.hstack((small_vals, big_vals))
    if isherm:
        evals = np.real(evals)

    _zipped = list(zip(evals, range(len(evals))))
    _zipped.sort()
    evals, perm = list(zip(*_zipped))

    if vecs:
        evecs = np.array([evecs[:, k] for k in perm])

    # remove last element if requesting N-1 eigs and using sparse
    if remove_one:
        evals = np.delete(evals, -1)
        if vecs:
            evecs = np.delete(evecs, -1)

    return np.array(evals), np.array(evecs)

def sp_eigs(data, isherm, vecs=True, sparse=False, sort='low',
            eigvals=0, tol=0, maxiter=100000):
    """Returns Eigenvalues and Eigenvectors for a sparse matrix.
    Uses dense eigen-solver unless user sets sparse=True.

    Parameters
    ----------
    data : csr_matrix
        Input matrix
    isherm : bool
        Indicate whether the matrix is hermitian or not
    vecs : bool {True , False}
        Flag for requesting eigenvectors
    sparse : bool {False , True}
        Flag to use sparse solver
    sort : str {'low' , 'high}
        Return lowest or highest eigenvals/vecs
    eigvals : int
        Number of eigenvals/vecs to return.  Default = 0 (return all)
    tol : float
        Tolerance for sparse eigensolver.  Default = 0 (Machine precision)
    maxiter : int
        Max. number of iterations used by sparse sigensolver.

    Returns
    -------
    Array of eigenvalues and (by default) array of corresponding Eigenvectors.

    """

    if debug:
        print(inspect.stack()[0][3])

    if data.shape[0] != data.shape[1]:
        raise TypeError("Can only diagonalize square matrices")

    N = data.shape[0]
    if eigvals == N:
        eigvals = 0

    if eigvals > N:
        raise ValueError("Number of requested eigen vals/vecs must be <= N.")

    # set number of large and small eigenvals/vecs
    if eigvals == 0:  # user wants all eigs (default)
        D = int(np.ceil(N / 2.0))
        num_large = N - D
        if not np.mod(N, 2):
            M = D
        else:
            M = D - 1
        num_small = N - M
    else:  # if user wants only a few eigen vals/vecs
        if sort == 'low':
            num_small = eigvals
            num_large = 0
        elif sort == 'high':
            num_large = eigvals
            num_small = 0
        else:
            raise ValueError("Invalid option for 'sort'.")

    # Dispatch to sparse/dense solvers
    if sparse:
        evals, evecs = _sp_eigs(data, isherm, vecs, N, eigvals, num_large,
                                num_small, tol, maxiter)
    else:
        evals, evecs = _dense_eigs(data.todense(), isherm, vecs, N, eigvals,
                                   num_large, num_small)

    if sort == 'high':  # flip arrays to largest values first
        if vecs:
            evecs = np.flipud(evecs)
        evals = np.flipud(evals)

    return (evals, evecs) if vecs else evals

def sp_expm(data, sparse=True):
    """
    Sparse matrix exponential.
    """
    A = data.tocsc()  # extract Qobj data (sparse matrix)
    m_vals = np.array([3, 5, 7, 9, 13])
    theta = np.array([0.01495585217958292, 0.2539398330063230,
                      0.9504178996162932, 2.097847961257068,
                      5.371920351148152], dtype=float)
    normA = sp_one_norm(data)
    if normA <= theta[-1]:
        for ii in range(len(m_vals)):
            if normA <= theta[ii]:
                F = _pade(A, m_vals[ii], sparse)
                break
    else:
        t, s = np.frexp(normA / theta[-1])
        s = s - (t == 0.5)
        A = A / 2.0 ** s
        F = _pade(A, m_vals[-1], sparse)
        for i in range(s):
            F = F * F

    return F

def _pade(A, m, sparse):
    n = np.shape(A)[0]
    c = _padecoeff(m)

    if m != 13:
        apows = [[] for jj in range(int(np.ceil((m + 1) / 2)))]
        apows[0] = sp.eye(n, n, format='csc')
        apows[1] = A * A
        for jj in range(2, int(np.ceil((m + 1) / 2))):
            apows[jj] = apows[jj - 1] * apows[1]
        U = sp.lil_matrix((n, n)).tocsc()
        V = sp.lil_matrix((n, n)).tocsc()
        for jj in range(m, 0, -2):
            U = U + c[jj] * apows[jj // 2]
        U = A * U
        for jj in range(m - 1, -1, -2):
            V = V + c[jj] * apows[(jj + 1) // 2]

        if sparse:
            F = spla.spsolve((-U + V), (U + V))
            return F.tocsr()
        else:
            F = la.solve((-U + V).todense(), (U + V).todense())
            return sp.lil_matrix(F).tocsr()

    elif m == 13:
        A2 = A * A
        A4 = A2 * A2
        A6 = A2 * A4
        U = A * (A6 * (c[13] * A6 + c[11] * A4 + c[9] * A2) +
                 c[7] * A6 + c[5] * A4 + c[3] * A2 +
                 c[1] * sp.eye(n, n).tocsc())
        V = A6 * (c[12] * A6 + c[10] * A4 + c[8] * A2) + c[6] * A6 + c[4] * \
            A4 + c[2] * A2 + c[0] * sp.eye(n, n).tocsc()

        if sparse:
            F = spla.spsolve((-U + V), (U + V))
            return F.tocsr()
        else:
            F = la.solve((-U + V).todense(), (U + V).todense())
            return sp.csr_matrix(F)

def _padecoeff(m):
    """
    Private function returning coefficients for Pade approximation.
    """
    if m == 3:
        return np.array([120, 60, 12, 1])
    elif m == 5:
        return np.array([30240, 15120, 3360, 420, 30, 1])
    elif m == 7:
        return np.array([17297280, 8648640, 1995840, 277200,
                         25200, 1512, 56, 1])
    elif m == 9:
        return np.array([17643225600, 8821612800, 2075673600,
                         302702400, 30270240, 2162160, 110880,
                         3960, 90, 1])
    elif m == 13:
        return np.array([64764752532480000, 32382376266240000,
                         7771770303897600, 1187353796428800,
                         129060195264000, 10559470521600, 670442572800,
                         33522128640, 1323241920, 40840800,
                         960960, 16380, 182, 1])

def sp_permute(A, rperm=(), cperm=(), safe=True):
    """
    Permutes the rows and columns of a sparse CSR/CSC matrix
    according to the permutation arrays rperm and cperm, respectively.
    Here, the permutation arrays specify the new order of the rows and
    columns. i.e. [0,1,2,3,4] -> [3,0,4,1,2].

    Parameters
    ----------
    A : csr_matrix, csc_matrix
        Input matrix.
    rperm : array_like of integers
        Array of row permutations.
    cperm : array_like of integers
        Array of column permutations.
    safe : bool
        Check structure of permutation arrays.

    Returns
    -------
    perm_csr : csr_matrix, csc_matrix
        CSR or CSC matrix with permuted rows/columns.

    """
    rperm = np.asarray(rperm, dtype=np.int32)
    cperm = np.asarray(cperm, dtype=np.int32)
    nrows = A.shape[0]
    ncols = A.shape[1]
    if len(rperm) == 0:
        rperm = np.arange(nrows, dtype=np.int32)
    if len(cperm) == 0:
        cperm = np.arange(ncols, dtype=np.int32)
    if safe:
        if len(np.setdiff1d(rperm, np.arange(nrows))) != 0:
            raise Exception('Invalid row permutation array.')
        if len(np.setdiff1d(cperm, np.arange(ncols))) != 0:
            raise Exception('Invalid column permutation array.')

    shp = A.shape
    kind = A.getformat()
    if kind == 'csr':
        flag = 0
    elif kind == 'csc':
        flag = 1
    else:
        raise Exception('Input must be Qobj, CSR, or CSC matrix.')

    data, ind, ptr = _sparse_permute(A.data, A.indices, A.indptr,
                                     nrows, ncols, rperm, cperm, flag)
    if kind == 'csr':
        return sp.csr_matrix((data, ind, ptr), shape=shp, dtype=data.dtype)
    elif kind == 'csc':
        return sp.csc_matrix((data, ind, ptr), shape=shp, dtype=data.dtype)

def sp_reverse_permute(A, rperm=(), cperm=(), safe=True):
    """
    Performs a reverse permutations of the rows and columns of a sparse CSR/CSC
    matrix according to the permutation arrays rperm and cperm, respectively.
    Here, the permutation arrays specify the order of the rows and columns used
    to permute the original array.

    Parameters
    ----------
    A : csr_matrix, csc_matrix
        Input matrix.
    rperm : array_like of integers
        Array of row permutations.
    cperm : array_like of integers
        Array of column permutations.
    safe : bool
        Check structure of permutation arrays.

    Returns
    -------
    perm_csr : csr_matrix, csc_matrix
        CSR or CSC matrix with permuted rows/columns.

    """
    rperm = np.asarray(rperm, dtype=np.int32)
    cperm = np.asarray(cperm, dtype=np.int32)
    nrows = A.shape[0]
    ncols = A.shape[1]
    if len(rperm) == 0:
        rperm = np.arange(nrows, dtype=np.int32)
    if len(cperm) == 0:
        cperm = np.arange(ncols, dtype=np.int32)
    if safe:
        if len(np.setdiff1d(rperm, np.arange(nrows))) != 0:
            raise Exception('Invalid row permutation array.')
        if len(np.setdiff1d(cperm, np.arange(ncols))) != 0:
            raise Exception('Invalid column permutation array.')

    shp = A.shape
    kind = A.getformat()
    if kind == 'csr':
        flag = 0
    elif kind == 'csc':
        flag = 1
    else:
        raise Exception('Input must be Qobj, CSR, or CSC matrix.')

    data, ind, ptr = _sparse_reverse_permute(A.data, A.indices, A.indptr,
                                             nrows, ncols, rperm, cperm, flag)

    if kind == 'csr':
        return sp.csr_matrix((data, ind, ptr), shape=shp, dtype=data.dtype)
    elif kind == 'csc':
        return sp.csc_matrix((data, ind, ptr), shape=shp, dtype=data.dtype)

def sp_bandwidth(A):
    """
    Returns the max(mb), lower(lb), and upper(ub) bandwidths of a
    sparse CSR/CSC matrix.

    If the matrix is symmetric then the upper and lower bandwidths are
    identical. Diagonal matrices have a bandwidth equal to one.

    Parameters
    ----------
    A : csr_matrix, csc_matrix
        Input matrix

    Returns
    -------
    mb : int
        Maximum bandwidth of matrix.
    lb : int
        Lower bandwidth of matrix.
    ub : int
        Upper bandwidth of matrix.

    """
    nrows = A.shape[0]
    ncols = A.shape[1]

    if A.getformat() == 'csr':
        return _sparse_bandwidth(A.indices, A.indptr, nrows)
    elif A.getformat() == 'csc':
        # Normal output is mb,lb,ub but since CSC
        # is transpose of CSR switch lb and ub
        mb, ub, lb = _sparse_bandwidth(A.indices, A.indptr, ncols)
        return mb, lb, ub
    else:
        raise Exception('Invalid sparse input format.')

def sp_profile(A):
    """Returns the total, lower, and upper profiles of a sparse matrix.

    If the matrix is symmetric then the upper and lower profiles are
    identical. Diagonal matrices have zero profile.

    Parameters
    ----------
    A : csr_matrix, csc_matrix
        Input matrix
    """
    if sp.isspmatrix_csr(A):
        up = _sparse_profile(A.indices, A.indptr, A.shape[0])
        A = A.tocsc()
        lp = _sparse_profile(A.indices, A.indptr, A.shape[0])

    elif sp.isspmatrix_csc(A):
        lp = _sparse_profile(A.indices, A.indptr, A.shape[0])
        A = A.tocsr()
        up = _sparse_profile(A.indices, A.indptr, A.shape[0])

    else:
        raise TypeError('Input sparse matrix must be in CSR or CSC format.')

    return up+lp, lp, up