import re
import logging
from indra.databases import uniprot_client
from indra.statements import Agent, Complex, Evidence
from indra.preassembler.grounding_mapper import standardize_agent_name


logger = logging.getLogger(__name__)


class VirhostnetProcessor:
    def __init__(self, df):
        self.df = df
        self.statements = []

    def extract_statements(self):
        for _, row in self.df.iterrows():
            stmt = process_row(row)
            if stmt:
                self.statements.append(stmt)


def process_row(row):
    host_agent = get_agent_from_grounding(row['host_grounding'])
    vir_agent = get_agent_from_grounding(row['vir_grounding'])

    # There's a column that is always a - character
    assert row['dash'] == '-', row['dash']

    exp_method_id, exp_method_name = parse_psi_mi(row['exp_method'])
    int_type__id, int_type_name = parse_psi_mi(row['int_type'])

    assert row['host_tax'].startswith('taxid:'), row['host_tax']
    _, host_tax = row['host_tax'].split(':')
    assert row['vir_tax'].startswith('taxid:'), row['vir_tax']
    _, vir_tax = row['vir_tax'].split(':')

    annotations = {
        'exp_method': {'id': exp_method_id, 'name': exp_method_name},
        'int_type': {'id': int_type__id, 'name': int_type_name},
        'host_tax': host_tax,
        'vir_tax': vir_tax
    }

    text_refs = parse_text_refs(row['publication'])

    ev = Evidence(source_api='virhostnet', annotations=annotations,
                  text_refs=text_refs, pmid=text_refs.get('PMID'))

    stmt = Complex([host_agent, vir_agent], evidence=[ev])
    return stmt


def get_agent_from_grounding(grounding):
    db_ns, db_id = grounding.split(':')
    # Assume UniProt or RefSeq IDs
    assert db_ns in {'uniprotkb', 'refseq', 'ddbj/embl/genbank'}, db_ns
    if db_ns == 'uniprotkb':
        if '-' in db_id:
            up_id, feat_id = db_id.split('-')
            # Assume it's a feature ID
            assert feat_id.startswith('PRO'), feat_id
            db_refs = {'UP': up_id, 'UPPRO': feat_id}
        else:
            db_refs = {'UP': db_id}
    elif db_ns == 'refseq':
        db_refs = {'REFSEQ_PROT': db_id}
    else:
        db_refs = {'GENBANK': db_id}
    agent = Agent(db_id, db_refs=db_refs)
    standardize_agent_name(agent, standardize_refs=True)
    return agent


def parse_psi_mi(psi_mi_str):
    # Example: psi-mi:"MI:0018"(two hybrid)
    match = re.match(r'psi-mi:"(.+)"\((.+)\)', psi_mi_str)
    mi_id, name = match.groups()
    return mi_id, name


def parse_text_refs(text_ref_str):
    tr_ns, tr_id = text_ref_str.split(':')
    assert tr_ns == 'pubmed', text_ref_str
    if re.match(r'^\d+$', tr_id):
        return {'PMID': tr_id}
    else:
        match = re.match(r'^https\(//doi.org/(.+)\)$')
        if not match:
            logger.warning('Failed to parse text ref: %s' % text_ref_str)
        doi = match.groups()[0]
        return {'DOI': doi}

@@ -1109,16 +1109,205 @@
 lit(':')
+%0A    assert row%5B'score'%5D.startswith('virhostnet-miscore:'), row%5B'score'%5D%0A    _, score = row%5B'score'%5D.split(':')%0A    score = float(score)%0A%0A    source_ids = parse_source_ids(row%5B'source_id'%5D)
 %0A%0A    an
@@ -1511,16 +1511,63 @@
  vir_tax
+,%0A        'score': score,%0A        **source_ids,
 %0A    %7D%0A%0A
@@ -1687,16 +1687,16 @@
 ations,%0A
-
         
@@ -1752,16 +1752,78 @@
 ('PMID')
+,%0A                  source_id=source_ids.get('virhostnet-rid')
 )%0A%0A    s
@@ -3053,18 +3053,25 @@
 (.+)%5C)$'
+, tr_id
 )%0A
-
         
@@ -3154,24 +3154,46 @@
 xt_ref_str)%0A
+            return %7B%7D%0A
         doi 
@@ -3212,16 +3212,16 @@
 ps()%5B0%5D%0A
-
         
@@ -3240,8 +3240,193 @@
 ': doi%7D%0A
+%0A%0Adef parse_source_ids(source_id_str):%0A    ids = source_id_str.split('%7C')%0A    assert len(ids) == 2%0A    ids_dict = %7Bid.split(':')%5B0%5D: id.split(':')%5B1%5D for id in ids%7D%0A    return ids_dict%0A

Handle scores and source IDs