<commit_before># MRA stands for mechanistic reasoning agent.
# Its task is to use INDRA to construct mechanistic models of
# biochemical systems from natural language, publications
# and databases.

import os
import copy
import json
import logging
import networkx
import subprocess
from datetime import datetime

import kappy

from bioagents import get_img_path
from indra.sources import trips
from indra.statements import Complex, Activation, IncreaseAmount, \
    stmts_from_json
from indra.preassembler.hierarchy_manager import hierarchies
from indra.assemblers.pysb import assembler as pysb_assembler
from indra.assemblers.pysb import PysbAssembler
from pysb.bng import BngInterfaceError
from pysb.tools import render_reactions

from pysb.export import export
from indra.util.kappa_util import im_json_to_graph, cm_json_to_graph
from bioagents.mra.sbgn_colorizer import SbgnColorizer
from bioagents.mra.model_diagnoser import ModelDiagnoser
logger = logging.getLogger('MRA')


class MRA(object):
    def __init__(self):
        self.models = {}
        self.transformations = []
        self.id_counter = 0
        self.default_policy = 'one_step'
        self.default_initial_amount = 100.0
        self.explain = None
        self.context = None

    def get_new_id(self):
        self.id_counter += 1
        return self.id_counter

    def has_id(self, model_id):
        if model_id in self.models:
            return True
        return False

    def assemble_pysb(self, stmts):
        pa = PysbAssembler()
        pa.add_statements(stmts)
        pa.make_model(policies=self.default_policy)
        pa.add_default_initial_conditions(self.default_initial_amount)
        return pa.model

    def build_model_from_ekb(self, model_ekb):
        """Build a model using DRUM extraction knowledge base."""
        tp = trips.process_xml(model_ekb)
        if tp is None:
            return {'error': 'Failed to process EKB.'}
        stmts = tp.statements
        model_id = self.new_model(stmts)
        res = {'model_id': model_id,
               'model': stmts}
        if not stmts:
            return res
        ambiguities = get_ambiguities(tp)
        res['ambiguities'] = ambiguities
        model_exec = self.assemble_pysb(stmts)
        res['model_exec'] = model_exec
        res['diagrams'] = make_diagrams(model_exec, model_id,
                                        self.models[model_id], self.context)
        self.run_diagnoser(res, stmts, model_exec)
        return res

    def build_model_from_json(self, model_json):
        """Build a model using INDRA JSON."""
        stmts = stmts_from_json(json.loads(model_json))
        model_id = self.new_model(stmts)
        res = {'model_id': model_id,
               'model': stmts}
        if not stmts:
            return res
        model_exec = self.assemble_pysb(stmts)
        res['model_exec'] = model_exec
        res['diagrams'] = make_diagrams(model_exec, model_id,
                                        self.models[model_id], self.context)
        return res

    def expand_model_from_ekb(self, model_ekb, model_id):
        """Expand a model using DRUM extraction knowledge base."""
        tp = trips.process_xml(model_ekb)
        if tp is None:
            return {'error': 'Failed to process EKB.'}
        stmts = tp.statements
        new_model_id, new_stmts = self.extend_model(stmts, model_id)
        logger.info('Old model id: %s, New model id: %s' %
                    (model_id, new_model_id))
        model_stmts = self.models[new_model_id]
        res = {'model_id': new_model_id,
               'model': model_stmts}
        if not model_stmts:
            return res
        ambiguities = get_ambiguities(tp)
        res['ambiguities'] = ambiguities
        res['model_new'] = new_stmts
        model_exec = self.assemble_pysb(model_stmts)
        res['model_exec'] = model_exec
        res['diagrams'] = make_diagrams(model_exec, new_model_id,
                                        self.models[new_model_id],
                                        self.context)
        self.run_diagnoser(res, model_stmts, model_exec)
        return res

    def run_diagnoser(self, res, model_stmts, model_exec):
        # Use a model diagnoser to identify explanations given the executable
        # model, the current statements, and the explanation goal
        if self.explain:
            md = ModelDiagnoser(model_stmts, model=model_exec,
                                explain=self.explain)
            md_result = md.check_explanation()
            res.update(md_result)
            # If we got a proposal for a statement, get a specific
            # recommendation
            connect_stmts = res.get('connect_stmts')
            if connect_stmts:
                u_stmt, v_stmt = connect_stmts
                stmt_suggestions = md.suggest_statements(u_stmt, v_stmt)
                if stmt_suggestions:
                    agents = [a.name for a in stmt_suggestions[0].agent_list()
                              if a is not None]
                    if len(set(agents)) > 1:
                        res['stmt_suggestions'] = stmt_suggestions
        md = ModelDiagnoser(model_stmts)
        acts = md.get_missing_activities()
        if acts:
            res['stmt_corrections'] = acts

    def expand_model_from_json(self, model_json, model_id):
        """Expand a model using INDRA JSON."""
        stmts = stmts_from_json(json.loads(model_json))
        new_model_id, new_stmts = self.extend_model(stmts, model_id)
        logger.info('Old model id: %s, New model id: %s' %
                    (model_id, new_model_id))
        model_stmts = self.models[new_model_id]
        res = {'model_id': new_model_id,
               'model': model_stmts}
        if not model_stmts:
            return res
        res['model_new'] = new_stmts
        model_exec = self.assemble_pysb(model_stmts)
        res['model_exec'] = model_exec
        res['diagrams'] = make_diagrams(model_exec, new_model_id,
                                        self.models[model_id], self.context)
        return res

    def has_mechanism(self, mech_ekb, model_id):
        """Return True if the given model contains the given mechanism."""
        tp = trips.process_xml(mech_ekb)
        res = {}
        if not tp.statements:
            res['has_mechanism'] = False
            return res
        query_st = tp.statements[0]
        res['query'] = query_st
        model_stmts = self.models[model_id]
        for model_st in model_stmts:
            if model_st.refinement_of(query_st, hierarchies):
                res['has_mechanism'] = True
                return res
        res['has_mechanism'] = False
        return res

    def remove_mechanism(self, mech_ekb, model_id):
        """Return a new model with the given mechanism having been removed."""
        tp = trips.process_xml(mech_ekb)<commit_msg>Fix json-based update model ID when making diagrams<commit_after># MRA stands for mechanistic reasoning agent.
# Its task is to use INDRA to construct mechanistic models of
# biochemical systems from natural language, publications
# and databases.

import os
import copy
import json
import logging
import networkx
import subprocess
from datetime import datetime

import kappy

from bioagents import get_img_path
from indra.sources import trips
from indra.statements import Complex, Activation, IncreaseAmount, \
    stmts_from_json
from indra.preassembler.hierarchy_manager import hierarchies
from indra.assemblers.pysb import assembler as pysb_assembler
from indra.assemblers.pysb import PysbAssembler
from pysb.bng import BngInterfaceError
from pysb.tools import render_reactions

from pysb.export import export
from indra.util.kappa_util import im_json_to_graph, cm_json_to_graph
from bioagents.mra.sbgn_colorizer import SbgnColorizer
from bioagents.mra.model_diagnoser import ModelDiagnoser
logger = logging.getLogger('MRA')


class MRA(object):
    def __init__(self):
        self.models = {}
        self.transformations = []
        self.id_counter = 0
        self.default_policy = 'one_step'
        self.default_initial_amount = 100.0
        self.explain = None
        self.context = None

    def get_new_id(self):
        self.id_counter += 1
        return self.id_counter

    def has_id(self, model_id):
        if model_id in self.models:
            return True
        return False

    def assemble_pysb(self, stmts):
        pa = PysbAssembler()
        pa.add_statements(stmts)
        pa.make_model(policies=self.default_policy)
        pa.add_default_initial_conditions(self.default_initial_amount)
        return pa.model

    def build_model_from_ekb(self, model_ekb):
        """Build a model using DRUM extraction knowledge base."""
        tp = trips.process_xml(model_ekb)
        if tp is None:
            return {'error': 'Failed to process EKB.'}
        stmts = tp.statements
        model_id = self.new_model(stmts)
        res = {'model_id': model_id,
               'model': stmts}
        if not stmts:
            return res
        ambiguities = get_ambiguities(tp)
        res['ambiguities'] = ambiguities
        model_exec = self.assemble_pysb(stmts)
        res['model_exec'] = model_exec
        res['diagrams'] = make_diagrams(model_exec, model_id,
                                        self.models[model_id], self.context)
        self.run_diagnoser(res, stmts, model_exec)
        return res

    def build_model_from_json(self, model_json):
        """Build a model using INDRA JSON."""
        stmts = stmts_from_json(json.loads(model_json))
        model_id = self.new_model(stmts)
        res = {'model_id': model_id,
               'model': stmts}
        if not stmts:
            return res
        model_exec = self.assemble_pysb(stmts)
        res['model_exec'] = model_exec
        res['diagrams'] = make_diagrams(model_exec, model_id,
                                        self.models[model_id], self.context)
        return res

    def expand_model_from_ekb(self, model_ekb, model_id):
        """Expand a model using DRUM extraction knowledge base."""
        tp = trips.process_xml(model_ekb)
        if tp is None:
            return {'error': 'Failed to process EKB.'}
        stmts = tp.statements
        new_model_id, new_stmts = self.extend_model(stmts, model_id)
        logger.info('Old model id: %s, New model id: %s' %
                    (model_id, new_model_id))
        model_stmts = self.models[new_model_id]
        res = {'model_id': new_model_id,
               'model': model_stmts}
        if not model_stmts:
            return res
        ambiguities = get_ambiguities(tp)
        res['ambiguities'] = ambiguities
        res['model_new'] = new_stmts
        model_exec = self.assemble_pysb(model_stmts)
        res['model_exec'] = model_exec
        res['diagrams'] = make_diagrams(model_exec, new_model_id,
                                        self.models[new_model_id],
                                        self.context)
        self.run_diagnoser(res, model_stmts, model_exec)
        return res

    def run_diagnoser(self, res, model_stmts, model_exec):
        # Use a model diagnoser to identify explanations given the executable
        # model, the current statements, and the explanation goal
        if self.explain:
            md = ModelDiagnoser(model_stmts, model=model_exec,
                                explain=self.explain)
            md_result = md.check_explanation()
            res.update(md_result)
            # If we got a proposal for a statement, get a specific
            # recommendation
            connect_stmts = res.get('connect_stmts')
            if connect_stmts:
                u_stmt, v_stmt = connect_stmts
                stmt_suggestions = md.suggest_statements(u_stmt, v_stmt)
                if stmt_suggestions:
                    agents = [a.name for a in stmt_suggestions[0].agent_list()
                              if a is not None]
                    if len(set(agents)) > 1:
                        res['stmt_suggestions'] = stmt_suggestions
        md = ModelDiagnoser(model_stmts)
        acts = md.get_missing_activities()
        if acts:
            res['stmt_corrections'] = acts

    def expand_model_from_json(self, model_json, model_id):
        """Expand a model using INDRA JSON."""
        stmts = stmts_from_json(json.loads(model_json))
        new_model_id, new_stmts = self.extend_model(stmts, model_id)
        logger.info('Old model id: %s, New model id: %s' %
                    (model_id, new_model_id))
        model_stmts = self.models[new_model_id]
        res = {'model_id': new_model_id,
               'model': model_stmts}
        if not model_stmts:
            return res
        res['model_new'] = new_stmts
        model_exec = self.assemble_pysb(model_stmts)
        res['model_exec'] = model_exec
        res['diagrams'] = make_diagrams(model_exec, new_model_id,
                                        self.models[new_model_id],
                                        self.context)
        return res

    def has_mechanism(self, mech_ekb, model_id):
        """Return True if the given model contains the given mechanism."""
        tp = trips.process_xml(mech_ekb)
        res = {}
        if not tp.statements:
            res['has_mechanism'] = False
            return res
        query_st = tp.statements[0]
        res['query'] = query_st
        model_stmts = self.models[model_id]
        for model_st in model_stmts:
            if model_st.refinement_of(query_st, hierarchies):
                res['has_mechanism'] = True
                return res
        res['has_mechanism'] = False
        return res

    def remove_mechanism(self, mech_ekb, model_id):
        """Return a new model with the given mechanism having been removed."""
        tp = trips.process_xml(mech_ekb)
{'fraction_non_alphanumeric': 0.08120766681077965, 'fraction_numerical': 0.0012970168612191958, 'mean_word_length': 2.211987965748669, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 8}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 13, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '877435', 'n_tokens_mistral': 4282, 'n_tokens_neox': 4129, 'n_words': 1000}