"""Export the researcher-supplied nine-design PETase update, verifying every asset.

Inputs are archived in Ouroboros_Showcase_Research/wetlab_20261110.
The earlier workbook and export snapshot remain in the historical research folder.
"""
import csv
import hashlib
import json
import shutil
from decimal import Decimal
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
AUDIT = ROOT.parent / 'Ouroboros_Showcase_Research/wetlab_20261110'
HISTORICAL = ROOT.parent / 'Ouroboros_Showcase_Research/wetlab_20260927'
PUBLIC = ROOT / 'public'
AA = dict(zip('ALA ARG ASN ASP CYS GLN GLU GLY HIS ILE LEU LYS MET PHE PRO SER THR TRP TYR VAL UNK'.split(), 'ARNDCQEGHILKMFPSTWYVX'))


def digest(path):
    return hashlib.sha256(path.read_bytes()).hexdigest()


def save(path, value):
    temporary = path.with_suffix(path.suffix + '.tmp')
    temporary.write_text(json.dumps(value, ensure_ascii=False, indent=2) + '\n')
    temporary.replace(path)


def main():
    sequences = {r['id']: r for r in json.loads((AUDIT / 'sequence_audit.json').read_text())['rows']}
    generation_audit = json.loads((AUDIT / 'generation_audit.json').read_text())
    generations = {r['id']: r for r in generation_audit['rows']}
    assert digest(Path(generation_audit['source'])) == generation_audit['sourceSha256']
    manifest = json.loads((AUDIT / 'fold_results/manifest.json').read_text())
    assert len(manifest) == 9 and {r['sequence_id'] for r in manifest} == {f'{i:04d}' for i in range(9)}
    job = (AUDIT / 'logs/ouroboros-petase-nine-20261110-r1/ouroboros-petase-nine-20261110-r1.job-id').read_text().strip()
    context = json.loads((AUDIT / 'report_context.json').read_text())
    supplied = list(csv.DictReader((AUDIT / 'assay_supplied.csv').open()))
    old_sequences = {r['id']: r for r in json.loads((HISTORICAL / 'sequence_mapping.json').read_text())}
    reference = old_sequences['rcsb_5XG0_IsPETase']
    rows = []
    for row in supplied:
        design = row['id'] != 'Y1'
        seq = sequences[row['id']] if design else reference
        assert abs(Decimal(row['tpa']) + Decimal(row['mhet']) - Decimal(row['total'])) <= Decimal('0.0001')
        fid = row['id'] if design else reference['id']
        rows.append(dict(id=row['id'], name=row['name'], kind='design' if design else 'wildtype',
                         **{k: float(row[k]) for k in ['tpa', 'mhet', 'total']}, length=seq['length'],
                         sequenceSha256=seq['sha256'], fastaUrl=f'/structures/wetlab-{fid}.fasta'))
    ranked = sorted([r for r in rows if r['kind'] == 'design'], key=lambda r: r['total'], reverse=True)
    cases = []
    for fid in ['de2', 'de1'] + [f'de{i}' for i in range(3, 10)] + ['rcsb_5XG0_IsPETase']:
        design = fid.startswith('de')
        seq = sequences[fid] if design else reference
        generation = generations.get(fid)
        if design:
            fold = next(r for r in manifest if r['sequence_id'] == f'{int(fid[2:])-1:04d}')
            pdb = AUDIT / 'fold_results' / fold['pdb_file']
            assert fold['pdb_sha256'] == digest(pdb)
            assert generation['sequenceSha256'] == seq['sha256']
            fasta = AUDIT / 'fastas' / f'{fid}.fasta'
        else:
            fold = next(r for r in json.loads((HISTORICAL / 'fold_results/manifest.json').read_text()) if r['sequence_id'] == '0001')
            pdb = HISTORICAL / 'fold_results' / fold['pdb_file']
            fasta = HISTORICAL / f'{fid}.fasta'
        sequence = ''.join(fasta.read_text().splitlines()[1:])
        assert sequence == seq['sequence']
        assert hashlib.sha256(sequence.encode()).hexdigest() == seq['sha256'] == fold['source_pred_seq_sha256']
        atoms = [line for line in pdb.read_text().splitlines() if line.startswith('ATOM') and line[12:16].strip() == 'CA']
        assert ''.join(AA[s[17:20]] for s in atoms) == sequence
        assert [int(s[22:26]) for s in atoms] == list(range(1, len(sequence)+1))
        residues = [dict(index=int(s[22:26]), chain=s[21].strip(), code=AA[s[17:20]], confidence=float(s[60:66])) for s in atoms]
        assert all(0 <= r['confidence'] <= 100 for r in residues)
        name = fid if design else 'IsPETase'
        provenance = (f"Submitted {fid.upper()} FASTA exactly matches archived LoRA generation {generation['generationId']}, prediction index {generation['predictionIndex0']} (zero-based), in lora_petase_3epoch/test/test_result.json. The original prompt is recovered from that same record. Structure predicted from the supplied full-length sequence."
                      if design else 'Supplied IsPETase reference FASTA labeled 5XG0, matched to assay sample Y1. This is an ESMFold-v1 prediction of that sequence, not an experimental crystal structure.')
        case = dict(id=name, assetId='wetlab-'+fid, title=name.upper() if design else name,
                    kind='design' if design else 'wildtype', modelLabel='LoRA-generated' if design else 'WT reference',
                    assaySample=fid if design else 'Y1', cohort='Screened PETase wet-lab candidates',
                    rank=next(i+1 for i, r in enumerate(ranked) if r['id'] == fid) if design else 0,
                    cohortSize=9, sequence=sequence, length=len(sequence), residues=residues,
                    prompt=generation['prompt'] if design else '', plddt=sum(r['confidence'] for r in residues)/len(residues),
                    pae=fold['mean_pae'], pdbUrl=f'/structures/wetlab-{fid}.pdb', fastaUrl=f'/structures/wetlab-{fid}.fasta',
                    selection='Highest combined TPA + MHET concentration among the nine active designs.' if fid == 'de2' else ('One of nine screened designs with reported PETase activity.' if design else 'Natural IsPETase reference, assay sample Y1.'),
                    structureStatus='ESMFold-v1 prediction of submitted FASTA',
                    experimentalStatus='PETase activity reported by the research team. Product concentrations are provided; assay conditions, replicates and variability were not supplied with the update.',
                    provenance=provenance, pdbSha256=digest(pdb), sequenceSha256=seq['sha256'],
                    foldJob=job if design else 'ouroboros-wetlab-de2-1009b', foldMethod='ESMFold-v1',
                    residueConfidenceDefinition='Mean of CA B factors in the exported ESMFold PDB')
        cases.append(case)
        shutil.copyfile(pdb, PUBLIC / case['pdbUrl'].lstrip('/'))
        shutil.copyfile(fasta, PUBLIC / case['fastaUrl'].lstrip('/'))
    context.update(rows=rows, csvUrl='/data/petase-assay.csv', sourceSha256=digest(AUDIT/'assay_supplied.csv'),
                   rounding='Values preserve four decimals as supplied. DE7 components differ from its supplied total by 0.0001 μM due to rounding.',
                   generation=generations['de2'],
                   structurePrediction=dict(method='ESMFold-v1', job=job, designCount=9,
                                            manifestSha256=digest(AUDIT/'fold_results/manifest.json'), experimentalStructure=False))
    save(PUBLIC/'data/petase.json', cases)
    save(PUBLIC/'data/wetlab.json', context)
    with (PUBLIC/'data/petase-assay.csv').open('w') as f:
        w=csv.writer(f);w.writerow(['id','name','TPA_uM','MHET_uM','total_uM'])
        for row in supplied:w.writerow([row[k] for k in ['id','name','tpa','mhet','total']])
    provenance=json.loads((PUBLIC/'data/provenance.json').read_text())
    provenance['petase']='Nine active, computationally screened PETase designs (DE1–DE9) and the natural IsPETase reference. DE2 is featured by reported total product concentration. Of 12 selected designs submitted for wet-lab testing, nine were active and three were not.'
    provenance['wetlab']={'report':context['source'],'sourceSha256':context['sourceSha256'], 'sequenceMatches':[{k:c[k] for k in ['id','sequenceSha256','foldJob','pdbSha256']} for c in cases], 'structurePrediction':context['structurePrediction']}
    save(PUBLIC/'data/provenance.json',provenance)
    for i in range(10,13):
        (PUBLIC/f'structures/wetlab-de{i}.fasta').unlink(missing_ok=True)
    print(json.dumps([dict(id=c['id'],length=c['length'],plddt=round(c['plddt'],2)) for c in cases],indent=2))


if __name__=='__main__':
    main()