#!/usr/bin/env python3
"""Synthetic LC-MS assignment exercise, CC BY 4.0, ChemVerify / Peak & Mass.

Python 3.9+ standard library only. Download solve.py, observations.csv and
parameters.csv into one directory, then run: python3 solve.py
The output my-solutions.csv evaluates every allowed hypothesis for each row.
No measured sample, sequence, instrument or validated tolerance is represented.
"""
import argparse
import csv
from decimal import Decimal, localcontext
from pathlib import Path

PARAMETERS = [
    ('neutral_mass_da', '1234.567890', 'Da', 'Fictional teaching input; no sequence or elemental composition assigned'),
    ('proton_mass_da', '1.0072764665789', 'Da', 'NIST CODATA 2022: https://physics.nist.gov/cuu/pdf/all.pdf'),
    ('sodium_ion_mass_da', '22.989221', 'Da', 'Rounded teaching approximation: sodium-23 atom minus electron; ionization energy neglected. https://physics.nist.gov/cgi-bin/Compositions/stand_alone.pl?ele=Na ; https://physics.nist.gov/cuu/pdf/all.pdf'),
    ('tolerance_ppm', '5', 'ppm', 'Fictional inclusive teaching criterion; not an instrument or method specification'),
]
OBSERVATIONS = [
    {'row_id':'A', 'retention_group':'R1', 'observed_mz':'1235.577638', 'charge_number':'1', 'data_origin':'synthetic'},
    {'row_id':'B', 'retention_group':'R1', 'observed_mz':'618.289367', 'charge_number':'2', 'data_origin':'synthetic'},
    {'row_id':'C', 'retention_group':'R1', 'observed_mz':'1257.558369', 'charge_number':'1', 'data_origin':'synthetic'},
    {'row_id':'D', 'retention_group':'R1', 'observed_mz':'1235.589993', 'charge_number':'1', 'data_origin':'synthetic'},
]
OBS_FIELDS = ['row_id', 'retention_group', 'observed_mz', 'charge_number', 'data_origin']
RESULT_FIELDS = ['row_id','ion_hypothesis','charge_number','observed_mz','calculated_mz','delta_mz','signed_error_ppm','denominator_mz','within_fictional_tolerance','data_origin']

def positive_decimal(value, name):
    result = Decimal(value)
    if not result.is_finite() or result <= 0:
        raise ValueError(name + ' must be a positive finite decimal')
    return result

def evaluate(observations, parameters):
    """Return all charge-compatible hypotheses; never silently force a match."""
    with localcontext() as context:
        context.prec = 40
        m, h, na, tolerance = [positive_decimal(parameters[key], key) for key in
            ['neutral_mass_da','proton_mass_da','sodium_ion_mass_da','tolerance_ppm']]
        hypotheses = [('[M+H]+', 1, m+h), ('[M+2H]2+', 2, (m+2*h)/2), ('[M+Na]+', 1, m+na)]
        results = []
        for row in observations:
            observed = positive_decimal(row['observed_mz'], 'observed_mz')
            if row['charge_number'] not in ('1','2'):
                raise ValueError('This exercise supports supplied charge_number 1 or 2 only')
            if row['data_origin'] != 'synthetic':
                raise ValueError('This solver is restricted to the synthetic teaching exercise')
            for ion, charge, expected in hypotheses:
                if charge != int(row['charge_number']):
                    continue
                delta = observed - expected
                ppm = delta / expected * 1_000_000
                results.append(dict(row_id=row['row_id'], ion_hypothesis=ion,
                    charge_number=str(charge), observed_mz=row['observed_mz'],
                    calculated_mz=format(expected,'f'), delta_mz=format(delta,'+f'),
                    signed_error_ppm=format(ppm,'+.12f'), denominator_mz=format(expected,'f'),
                    within_fictional_tolerance=str(abs(ppm) <= tolerance).lower(), data_origin='synthetic'))
        return results

def write_csv(path, fields, rows):
    with Path(path).open('w', newline='', encoding='utf-8') as stream:
        writer = csv.DictWriter(stream, fieldnames=fields, lineterminator='\n')
        writer.writeheader()
        writer.writerows(rows)

def write_data(folder):
    folder.mkdir(parents=True, exist_ok=True)
    write_csv(folder/'observations.csv', OBS_FIELDS, OBSERVATIONS)
    write_csv(folder/'parameters.csv', ['parameter','value','unit','provenance'],
        [dict(zip(['parameter','value','unit','provenance'], row)) for row in PARAMETERS])
    write_csv(folder/'solutions.csv', RESULT_FIELDS,
        evaluate(OBSERVATIONS, {row[0]:row[1] for row in PARAMETERS}))
    (folder/'solve.py').write_bytes(Path(__file__).read_bytes())

def main():
    directory = Path(__file__).resolve().parent
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--observations', type=Path, default=directory/'observations.csv')
    parser.add_argument('--parameters', type=Path, default=directory/'parameters.csv')
    parser.add_argument('--output', type=Path, default=directory/'my-solutions.csv')
    args = parser.parse_args()
    with args.observations.open(newline='', encoding='utf-8-sig') as stream:
        observations = list(csv.DictReader(stream))
    with args.parameters.open(newline='', encoding='utf-8-sig') as stream:
        parameters = {row['parameter']:row['value'] for row in csv.DictReader(stream)}
    results = evaluate(observations, parameters)
    # Avoid replacing downloaded inputs or an earlier calculation accidentally.
    if args.output.exists():
        parser.error('Output already exists; choose a new path with --output')
    write_csv(args.output, RESULT_FIELDS, results)
    print(f'{len(results)} hypothesis comparisons written to {args.output}')

if __name__ == '__main__':
    main()