"""Freeze the full S7 source and sparse witnesses; no saved PASS verdict is imported.

Usage: python freeze_source.py --workspace PATH --output DIRECTORY
The author workspace supplies sources/rbcgem_s7.xml and the named audit records.
The runnable package includes the resulting JSON, so this step is optional.
"""
import argparse
from fractions import Fraction as Q
import hashlib
import json
from pathlib import Path
import re
import shutil
import xml.etree.ElementTree as ET

SOURCE_SHA256='b2703b5a730cf12890034f01df92506de5facd9ecea5751a795e982b88557f64'
BUDGET_ROWS={'M_proteome_budget','M_hemoglobin_budget','M_total_budget','M_relaxation_budget'}
IONIC_SINKS={'R_SK_na1_c','R_SK_k_c','R_SK_ca2_c'}
RETAINED_CARBON={'R_EX_glc__D_e','R_EX_co2_e','R_EX_hco3_e'}


def digest(p):return hashlib.sha256(Path(p).read_bytes()).hexdigest()
def attrs(e):return {k.rsplit('}',1)[-1]:v for k,v in e.attrib.items()}
def tag(e):return e.tag.rsplit('}',1)[-1]


def parse_xml(path):
    rows=[];columns=[];parameters={}
    for _,e in ET.iterparse(path,events=['end']):
        if tag(e)=='species':
            a=attrs(e);rid=a['id'];rows.append(dict(id=rid,compartment=a.get('compartment',''),formula=a.get('chemicalFormula',''),
                kind='auxiliary' if a.get('compartment')=='pc' or rid in BUDGET_ROWS else 'chemical'));e.clear()
        elif tag(e)=='parameter':
            a=attrs(e)
            if a['id'] in parameters:raise ValueError('Duplicate parameter identifier.')
            parameters[a['id']]=a['value'];e.clear()
        elif tag(e)=='reaction':
            a=attrs(e);s={}
            for side in e:
                if tag(side) not in ('listOfReactants','listOfProducts'):continue
                sign=-1 if tag(side)=='listOfReactants' else 1
                for ref in side:
                    b=attrs(ref);rid=b['species'];s[rid]=s.get(rid,Q(0))+sign*Q(b.get('stoichiometry','1'))
            columns.append(dict(id=a['id'],lower=parameters[a['lowerFluxBound']],upper=parameters[a['upperFluxBound']],
                stoichiometry={k:str(v) for k,v in s.items() if v}));e.clear()
    if len({r['id'] for r in rows})!=len(rows) or len({r['id'] for r in columns})!=len(columns):raise ValueError('Duplicate source identifiers.')
    return rows,columns


def medium_masks(rows,columns):
    formulas={r['id']:r['formula'] for r in rows};sinks=[];sulfur=[];carbon=[]
    for c in columns:
        rid=c['id'];stoich=c['stoichiometry']
        if rid.startswith('R_SK_') and rid not in IONIC_SINKS:sinks.append(rid)
        if not rid.startswith('R_EX_'):continue
        hasS=any(re.search(r'S(?![a-z])',formulas[s]) for s in stoich)
        hasC=rid not in RETAINED_CARBON and any(re.search(r'C(?![a-z])',formulas[s]) for s in stoich)
        if hasS or hasC:
            if len(stoich)!=1 or Q(next(iter(stoich.values())))!=-1:raise ValueError('Import closure needs the audited singleton negative exchange convention.')
        if hasS:sulfur.append(rid)
        if hasC:carbon.append(rid)
    return dict(nonionic_sink_imports=sinks,sulfur_exchange_imports=sulfur,other_carbon_exchange_imports=carbon)


def main():
    ap=argparse.ArgumentParser(description=__doc__);ap.add_argument('--workspace',type=Path,required=True);ap.add_argument('--output',type=Path,required=True);a=ap.parse_args();p=a.workspace;out=a.output;out.mkdir(parents=True,exist_ok=True)
    xml=p/'sources/rbcgem_s7.xml'
    if digest(xml)!=SOURCE_SHA256:raise ValueError('Unexpected XML: use a separately reviewed dataset for a new source version.')
    rows,columns=parse_xml(xml);masks=medium_masks(rows,columns)
    assert (len(rows),len(columns),sum(r['kind']=='chemical' for r in rows))==(10411,19620,2157)
    assert (len(masks['sulfur_exchange_imports']),len(masks['other_carbon_exchange_imports']))==(56,435)
    def dump(name,data):(out/name).write_text(json.dumps(data,separators=(',',':'))+'\n')
    dump('s7_model.json',dict(source_xml_sha256=SOURCE_SHA256,attribution='Haiman, Key, D’Alessandro and Palsson (2025), RBC-GEM, PLOS Computational Biology 21:e1012109; S7.',
        source_url='https://doi.org/10.1371/journal.pcbi.1012109',rows=rows,columns=columns,medium_masks=masks))
    full=json.loads((p/'plateau_full_audit.json').read_text());joint=json.loads((p/'finite_joint_certificate.json').read_text())
    cert=dict(primal=full['primal'],turnover=dict(chemical=full['chemical_weights'],auxiliary=full['auxiliary_weights']),
        service=dict(chemical=full['service_chemical_weights'],auxiliary=full['service_auxiliary_weights']),
        joint=dict(chemical=joint['chemical_weights'],auxiliary=joint['auxiliary_weights'],reverse={k:v for k,v in joint['summary']['reverse_supply_coefficients'].items() if Q(v)}))
    dump('certificates.json',cert)
    shutil.copyfile(p/'sources/RBC_GEM_LICENSE.md',out/'RBC_GEM_LICENSE.md')
    dump('source_manifest.json',dict(source_xml_sha256=SOURCE_SHA256,model_sha256=digest(out/'s7_model.json'),certificates_sha256=digest(out/'certificates.json'),
        original_records={name:digest(p/name) for name in ['plateau_full_audit.json','finite_joint_certificate.json']},
        transformation='All 19620 columns and 10411 rows retained. Source decimal bounds and exact rational stoichiometry; explicit row classification and import masks. Sparse primal/dual entries extracted by source ID, saved audit verdicts discarded.',
        source_license='RBC-GEM Creative Commons Attribution 4.0 International; unchanged license text included. Scientific data adapted to a rational sparse representation.'))
    print(f'Frozen {len(columns)} columns, {len(rows)} rows; model bytes {(out/"s7_model.json").stat().st_size}.')


if __name__=='__main__':main()