from copy import deepcopy
from indra.databases import hgnc_client
from indra.statements import Agent, IncreaseAmount, DecreaseAmount, Evidence


class TrrustProcessor(object):
    """Processor to extract INDRA Statements from Trrust data frame.

    Attributes
    ----------
    df : pandas.DataFrame
        The Trrust table to process.
    statements : list[indra.statements.Statement]
        The list of INDRA Statements extracted from the table.
    """
    def __init__(self, df):
        self.df = df
        self.statements = []

    def extract_statements(self):
        """Process the table to extract Statements."""
        for _, (tf, target, effect, refs) in self.df.iterrows():
            tf_agent = get_grounded_agent(tf)
            target_agent = get_grounded_agent(target)
            if effect == 'Activation':
                stmt_cls = IncreaseAmount
            elif effect == 'Repression':
                stmt_cls = DecreaseAmount
            else:
                continue
            pmids = refs.split(';')
            for pmid in pmids:
                stmt = make_stmt(stmt_cls, tf_agent, target_agent, pmid)
                self.statements.append(stmt)


def make_stmt(stmt_cls, tf_agent, target_agent, pmid):
    """Return a Statement based on its type, agents, and PMID."""
    ev = Evidence(source_api='trrust', pmid=pmid)
    return stmt_cls(deepcopy(tf_agent), deepcopy(target_agent),
                    evidence=[ev])


def get_grounded_agent(gene_name):
    """Return a grounded Agent based on an HGNC symbol."""
    db_refs = {'TEXT': gene_name}
    hgnc_id = hgnc_client.get_hgnc_id(gene_name)
    if hgnc_id:
        db_refs['HGNC'] = hgnc_id
        up_id = hgnc_client.get_uniprot_id(hgnc_id)
        if up_id:
            db_refs['UP'] = up_id
    agent = Agent(gene_name, db_refs=db_refs)
    return agent

@@ -1580,16 +1580,86 @@
 e_name%7D%0A
+    if gene_name in hgnc_map:%0A        gene_name = hgnc_map%5Bgene_name%5D%0A
     hgnc
@@ -1912,12 +1912,196 @@
 eturn agent%0A
+%0A%0Ahgnc_map = %7B%0A    'CTGF': 'CCN2',%0A    'CYR61': 'CCN1',%0A    'MKL1': 'MRTFA',%0A    'NOV': 'CCN3',%0A    'RFWD2': 'COP1',%0A    'SALL4A': 'SALL4',%0A    'STAT5': 'STAT5A',%0A    'TRAP': 'ACP5',%0A%7D

Add manual mapping for outdated HGNC names