"""Run E1--E3 and common-score IPD controls from verified development evidence."""
import argparse
import json
import os
from pathlib import Path
import time
from build_graph import sha
from context_inference import emission_rows, infer, gene_evidence_rows, hard_locus_rows

MODES = {
    'E1': dict(context=False, sequence_first=False, include_graph=True),
    'E2': dict(context=False, sequence_first=True, include_graph=True),
    'E3': dict(context=True, sequence_first=True, include_graph=True),
    'IPD-common': dict(context=False, sequence_first=True, include_graph=False),
    'E3-anchored': dict(context=True, sequence_first=False, include_graph=True),
    'E3-hard-locus': dict(context=True, sequence_first=True, include_graph=True, hard_locus=True),
}


def run(evidence, baseline, output, mode):
    if int(os.environ.get('SLURM_CPUS_PER_TASK', '0')) < 1:
        raise RuntimeError('Run under Slurm')
    if output.exists():
        raise FileExistsError(output)
    manifest = json.loads((evidence / 'COMPLETE.json').read_text())
    if manifest['status'] != 'complete':
        raise ValueError('Incomplete evidence')
    for name, digest in manifest['output_sha256'].items():
        if sha(evidence / name) != digest:
            raise ValueError('Evidence changed')
    meta = json.loads((evidence / 'candidates.json').read_text())
    native = json.loads(baseline.read_text())
    output.mkdir(parents=True)
    parameters = MODES[mode]
    record = dict(status='running', started=time.time(), mode=mode, parameters=parameters,
                  evidence_sha256=sha(evidence / 'COMPLETE.json'), baseline_sha256=sha(baseline),
                  code_sha256={name: sha(Path(__file__).with_name(name)) for name in
                               ['run_context_inference.py', 'context_inference.py', 'joint_model.py']},
                  scope='Exposed development data; no independent accuracy claim')
    def save():
        (output / 'manifest.json').write_text(json.dumps(record, indent=2) + '\n')
    save()
    try:
        with (evidence / 'fragments.jsonl').open() as stream:
            identifiers, rows = emission_rows(map(json.loads, stream), meta,
                                              context=parameters['context'], include_graph=parameters['include_graph'])
        print(json.dumps(dict(informative_fragments=len(rows), observed_candidates=len({p for r in rows for p in r}))), flush=True)
        core = None
        if parameters['context']:
            with (evidence / 'fragments.jsonl').open() as stream:
                core = gene_evidence_rows(map(json.loads, stream), identifiers)
        if parameters.get('hard_locus'):
            rows = hard_locus_rows(rows, meta)
        calls, diagnostic = infer(rows, meta, native, sequence_first=parameters['sequence_first'], core_rows=core)
        (output / 'calls.json').write_text(json.dumps(calls, indent=2) + '\n')
        (output / 'decisions.json').write_text(json.dumps(diagnostic, indent=2) + '\n')
        record.update(status='complete', finished=time.time(), fragments=len(identifiers),
                      output_sha256={n: sha(output / n) for n in ['calls.json', 'decisions.json']})
        save()
        (output / 'COMPLETE.json').write_text(json.dumps(record, indent=2) + '\n')
    except BaseException as error:
        record.update(status='failed', finished=time.time(), error=repr(error))
        save()
        raise


if __name__ == '__main__':
    p = argparse.ArgumentParser(description=__doc__)
    for name in ['evidence', 'baseline', 'output']:
        p.add_argument('--' + name, type=Path, required=True)
    p.add_argument('--mode', choices=MODES, required=True)
    a = p.parse_args()
    run(a.evidence, a.baseline, a.output, a.mode)