"""Rescore existing graph/native candidate placements with one fragment model.

Development evidence preparation; no truth is loaded. Preserves every native
candidate/locus and every graph path with a proposed concordant placement.
"""
import argparse
import collections
import gzip
import hashlib
import json
import math
import multiprocessing
import os
from pathlib import Path
import time

from build_genome_control import records
from build_graph import sequences, sha
from context_fragment import realign, pair_score, overlaps_gene, internal_pair

_WORKER_STATE = None


def initialize_worker(pairs, seqs, meta, mean, stdev):
    global _WORKER_STATE
    _WORKER_STATE = pairs, seqs, meta, mean, stdev


def score_record(record):
    pairs, seqs, meta, mean, stdev = _WORKER_STATE
    name = record['fragment']
    if name not in pairs:
        raise ValueError('Evidence fragment absent from FASTQ')
    return name, score_fragment(pairs[name], proposals(record), seqs, meta, mean, stdev)


def scored_records(records, pairs, seqs, meta, mean, stdev, workers):
    """Preserve fragment order and exact scoring under sample-level parallelism."""
    global _WORKER_STATE
    try:
        if workers == 1:
            initialize_worker(pairs, seqs, meta, mean, stdev)
            yield from map(score_record, records)
        else:
            with multiprocessing.get_context('fork').Pool(workers, initializer=initialize_worker,
                    initargs=(pairs, seqs, meta, mean, stdev)) as pool:
                yield from pool.imap(score_record, records, chunksize=64)
    finally:
        _WORKER_STATE = None


def read_pairs(paths):
    result = {}
    def fastq(path):
        opener = gzip.open if path.suffix == '.gz' else open
        with opener(path, 'rt') as stream:
            while True:
                header = stream.readline()
                if not header:
                    break
                seq, plus, qual = [stream.readline().rstrip('\n\r') for _ in range(3)]
                if not header.startswith('@') or not plus.startswith('+') or len(seq) != len(qual):
                    raise ValueError('Invalid FASTQ')
                name = header[1:].split()[0]
                if name.endswith(('/1', '/2')):
                    name = name[:-2]
                yield name, seq.upper(), qual
    import itertools
    for a, b in itertools.zip_longest(fastq(paths[0]), fastq(paths[1])):
        if a is None or b is None or a[0] != b[0] or a[0] in result:
            raise ValueError('Mismatched or repeated fragment identifier')
        result[a[0]] = (a[1:], b[1:])
    return result


def references(native, panel):
    native_data = records(native.read_text(), validate_coordinates=False)
    seqs, meta = {}, {}
    for name, (header, sequence) in native_data.items():
        identifier = 'IPD:' + name
        gene = name.split('*')[0].removeprefix('HLA-')
        label = name.removeprefix('HLA-')
        seqs[identifier] = sequence
        meta[identifier] = dict(gene=gene, labels=[label], source='ipd',
                                body=[0, len(sequence)], sequence_sha256=hashlib.sha256(sequence.encode()).hexdigest())
    for sidecar in sorted(panel.glob('HLA-*.json')):
        fasta = sidecar.with_suffix('.fa')
        data = sequences(fasta)
        for row in json.loads(sidecar.read_text()):
            path = row['path']
            if path in seqs or hashlib.sha256(data[path].encode()).hexdigest() != row['sequence_sha256']:
                raise ValueError('Graph path/reference mismatch')
            seqs[path] = data[path]
            meta[path] = dict(gene=sidecar.stem.removeprefix('HLA-'),
                             labels=[a.removeprefix('HLA-') for a in row['genomic_labels']],
                             cds_labels=[a.removeprefix('HLA-') for a in row['cds_labels']],
                             source='graph', body=[2000, len(data[path]) - 2000],
                             sources=row['sources'], sequence_sha256=row['sequence_sha256'])
    return seqs, meta


def proposals(record):
    """Raw scores never leave this candidate-location adapter."""
    output = collections.defaultdict(set)
    for name, placements in record['native_candidates'].items():
        for p in placements:
            # Native export provides a fragment envelope, with inclusive end.
            if p['adjusted_weight'] <= 0:
                continue
            for reverse in (False, True):
                output['IPD:' + name].add((p['start'], p['end'] + 1,
                                          p['start'], p['end'] + 1, reverse))
    for source in record['graph_sources'].values():
        for name, placements in source['candidates'].items():
            for p in placements:
                output[name].add((p['first_start'], p['first_end'],
                                  p['second_start'], p['second_end'], p['first_reverse']))
    return output


def score_fragment(reads, locations, seqs, meta, mean, stdev):
    output = {}
    cache = {}
    maximum = int(math.ceil(mean + 6 * stdev))
    for identifier, positions in locations.items():
        if identifier not in seqs:
            raise ValueError('Unknown candidate: ' + identifier)
        sequence = seqs[identifier]
        candidates = {}
        for a, b, c, d, reverse in sorted(positions):
            if not (0 <= a < b <= len(sequence) and 0 <= c < d <= len(sequence)):
                raise ValueError('Candidate proposal outside reference')
            placements = []
            for mate, start, end, orientation in [(0, a, b, reverse), (1, c, d, not reverse)]:
                # Identical sequence/windows across aliases have identical evidence.
                key = (meta[identifier]['sequence_sha256'], tuple(meta[identifier]['body']), mate, start, end, orientation)
                if key not in cache:
                    cache[key] = realign(*reads[mate], sequence, start, end, orientation,
                                         body=meta[identifier]['body'])
                placements.append(cache[key])
            score = pair_score(*placements, mean, stdev, maximum)
            if score is None:
                continue
            first, second = placements
            placement = dict(first_start=first['start'], first_end=first['end'],
                             second_start=second['start'], second_end=second['end'],
                             first_reverse=first['reverse'])
            key = tuple(placement.values())
            candidates[key] = dict(score=score, **placement,
                gene_penalty=first['gene_penalty'] + second['gene_penalty'],
                gene_bases=first['gene_bases'] + second['gene_bases'],
                gene_overlap=overlaps_gene(placement, meta[identifier]['body']),
                internal=internal_pair(placement, meta[identifier]['body']))
        if candidates:
            # Preserve placements; downstream code must not count each as a read.
            output[identifier] = list(candidates.values())
    return output


def run(joined, native, panel, calibration, reads, output):
    if not os.environ.get('SLURM_JOB_ID'):
        raise RuntimeError('Use a Slurm allocation')
    if output.exists():
        raise FileExistsError(output)
    jm = json.loads((joined / 'COMPLETE.json').read_text())
    cm = json.loads(calibration.read_text())
    if jm['status'] != 'complete' or sha(joined / 'evidence.jsonl') != jm['output_sha256']:
        raise ValueError('Joined evidence changed')
    if cm['status'] != 'complete' or not cm['estimate']['usable']:
        raise ValueError('Invalid calibration')
    read_hashes = {str(p): sha(p) for p in reads}
    if sorted(read_hashes.values()) != sorted(cm['reads_sha256'].values()):
        raise ValueError('Calibration/read mismatch')
    # Validate graph-chain parents and their recorded read identity before reuse.
    for chain in jm['graph_chains']:
        for folder, digest in chain.items():
            manifest = Path(folder) / 'COMPLETE.json'
            if sha(manifest) != digest:
                raise ValueError('Graph evidence parent changed')
            parent = json.loads(manifest.read_text())
            if 'reads_sha256' in parent and sorted(parent['reads_sha256'].values()) != sorted(read_hashes.values()):
                raise ValueError('Graph/read mismatch')
    seqs, meta = references(native, panel)
    pairs = read_pairs(reads)
    workers = min(4, int(os.environ.get('SLURM_CPUS_PER_TASK', '1')))
    output.mkdir(parents=True)
    report = dict(status='running', started=time.time(), reads_sha256=read_hashes,
                  joined_sha256=sha(joined / 'COMPLETE.json'), native_sha256=sha(native),
                  panel_sha256={p.name: sha(p) for p in sorted(panel.glob('HLA-*'))},
                  calibration_sha256=sha(calibration), driver_sha256=sha(Path(__file__)),
                  workers=workers,
                  model_sha256=sha(Path(__file__).with_name('context_fragment.py')),
                  scope='Development common-score evidence; no genotype accuracy or calibrated posterior claim',
                  limitations=['Candidate retrieval remains incomplete; native-only alignments use a fragment-envelope window',
                               'Edit-optimal alignment rescored with base quality is not exhaustive quality-optimal alignment'])
    def save():
        (output / 'manifest.json').write_text(json.dumps(report, indent=2) + '\n')
    save()
    try:
        counts = collections.Counter()
        previous = None
        with (joined / 'evidence.jsonl').open() as stream, (output / 'fragments.jsonl').open('w') as target:
            scored = scored_records(map(json.loads, stream), pairs, seqs, meta,
                                     cm['estimate']['fragment_mean'], cm['estimate']['fragment_stdev'], workers)
            for name, values in scored:
                if previous is not None and name <= previous:
                    raise ValueError('Unsorted/repeated evidence fragment')
                previous = name
                target.write(json.dumps(dict(fragment=name, candidates=values), separators=(',', ':')) + '\n')
                counts['fragments'] += 1
                counts['supported'] += bool(values)
                counts['cross_locus'] += len({meta[p]['gene'] for p in values}) > 1
                if counts['fragments'] % 1000 == 0:
                    print(json.dumps(dict(counts)), flush=True)
        (output / 'candidates.json').write_text(json.dumps(meta, indent=2) + '\n')
        report.update(status='complete', finished=time.time(), counts=dict(counts),
                      output_sha256={n: sha(output / n) for n in ['fragments.jsonl', 'candidates.json']})
        save()
        (output / 'COMPLETE.json').write_text(json.dumps(report, indent=2) + '\n')
    except BaseException as error:
        report.update(status='failed', finished=time.time(), error=repr(error))
        save()
        raise


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    for name in ['joined', 'native', 'panel', 'calibration', 'output']:
        parser.add_argument('--' + name, type=Path, required=True)
    parser.add_argument('--reads', type=Path, nargs=2, required=True)
    args = parser.parse_args()
    run(args.joined, args.native, args.panel, args.calibration, args.reads, args.output)