"""Reproduce the article with chematic==1.1.4; add --plot for Matplotlib SVG.

The Frank model is independent of chematic's molecular calculations.
Dimensionless equations: Deck et al., Supplementary Information, section 9.1,
https://doi.org/10.26434/chemrxiv-2025-dqczr (equations 26 and 27).
Parameters illustrate a mathematical model; they are not a fit to the Soai reaction.
"""
import argparse
import csv
import json
from pathlib import Path

import chematic


def derivative(r, s, delta):
    return r * (1 - r - (1 + delta) * s), s * (1 - s - (1 + delta) * r)


def simulate(r_fraction=0.51, delta=1.0, step=0.01, end=20.0):
    """Fixed-step RK4. Initial scaled total concentration is 1."""
    r, s = r_fraction, 1 - r_fraction
    rows = []
    for i in range(round(end / step) + 1):
        rows.append({"tau": i * step, "r": r, "s": s,
                     "signed_ee_percent": 100 * (r - s) / (r + s)})
        a, b = derivative(r, s, delta)
        c, d = derivative(r + step*a/2, s + step*b/2, delta)
        e, f = derivative(r + step*c/2, s + step*d/2, delta)
        g, h = derivative(r + step*e, s + step*f, delta)
        r += step * (a + 2*c + 2*e + g) / 6
        s += step * (b + 2*d + 2*f + h) / 6
    return rows


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--out', type=Path, default=Path('.'))
    parser.add_argument('--plot', action='store_true')
    args = parser.parse_args()
    if chematic.__version__ != '1.1.4':
        raise RuntimeError('This example targets chematic==1.1.4')
    args.out.mkdir(parents=True, exist_ok=True)
    inputs = ['C[C@H](O)C(=O)O', 'C[C@@H](O)C(=O)O']
    molecules = [chematic.from_smiles(s) for s in inputs]
    structures = []
    for smiles, mol in zip(inputs, molecules):
        descriptor = mol.cip_stereo(mode='accurate')[0]['descriptor']
        structures.append({'smiles': smiles, 'cip': descriptor, 'formula': mol.formula,
                           'molecular_weight': round(mol.rdkit_mw, 3)})
        (args.out / f'lactic-{descriptor.lower()}.svg').write_text(mol.svg(), encoding='utf-8')
        print(f'{descriptor} {mol.formula} {mol.rdkit_mw:.3f}')
    plain_equal = molecules[0].rdkit_ecfp_config(2, 2048) == molecules[1].rdkit_ecfp_config(2, 2048)
    chiral_equal = molecules[0].rdkit_ecfp_config(2, 2048, include_chirality=True) == molecules[1].rdkit_ecfp_config(2, 2048, include_chirality=True)
    assert [row['cip'] for row in structures] == ['S', 'R']
    assert plain_equal and not chiral_equal
    print('Morgan equal, chirality off:', plain_equal)
    print('Morgan equal, chirality on:', chiral_equal)
    cases = []
    traces = []
    for r_fraction, delta in [(0.51, 1.0), (0.51, 0.0), (0.50, 1.0), (0.49, 1.0)]:
        trace = simulate(r_fraction, delta)
        finer = simulate(r_fraction, delta, step=0.005)
        error = max(abs(a['signed_ee_percent'] - b['signed_ee_percent']) for a,b in zip(trace, finer[::2]))
        assert error < 1e-6
        assert all(row['r'] >= 0 and row['s'] >= 0 for row in trace)
        if delta == 0:
            assert max(abs(row['signed_ee_percent'] - 2) for row in trace) < 1e-10
        if r_fraction == 0.5:
            assert all(row['signed_ee_percent'] == 0 for row in trace)
        cases.append({'initial_r_fraction': r_fraction, 'delta': delta,
                      'max_step_halving_ee_difference': error, 'samples': trace[::100]})
        traces.append(trace)
        print(f'R:S = {r_fraction*100:.0f}:{(1-r_fraction)*100:.0f}, delta = {delta:.0f}, ee at tau 20 = {trace[-1]["signed_ee_percent"]:.4f}%')
    results = {'chematic_version': chematic.__version__, 'structures': structures,
               'morgan': {'radius': 2, 'bits': 2048, 'equal_without_chirality': plain_equal, 'equal_with_chirality': chiral_equal},
               'model': {'name': 'Dimensionless Frank model', 'step': 0.01, 'end': 20, 'initial_total': 1,
                         'interpretation': 'Illustrative parameters; not fitted to a chemical experiment', 'cases': cases}}
    (args.out / 'results.json').write_text(json.dumps(results, indent=2) + '\n', encoding='utf-8')
    with (args.out / 'model.csv').open('w', newline='') as f:
        writer = csv.writer(f, lineterminator='\n')
        writer.writerow(['initial_r_fraction', 'delta', 'tau', 'r', 's', 'signed_ee_percent'])
        for case, trace in zip(cases, traces):
            for row in trace[::10]:
                writer.writerow([case['initial_r_fraction'], case['delta'], *row.values()])
    if args.plot:
        import matplotlib
        matplotlib.use('Agg')
        import matplotlib.pyplot as plt
        matplotlib.rcParams['svg.hashsalt'] = 'chematic-nobel-2026'
        fig, ax = plt.subplots(figsize=(8, 4.8), layout='constrained')
        labels = ['51:49, delta = 1', '51:49, delta = 0', '50:50, delta = 1', '49:51, delta = 1']
        for trace, label, color, style in zip(traces, labels, ['#165fad', '#937029', '#546070', '#a6375b'], ['-', '--', ':', '-.']):
            ax.plot([x['tau'] for x in trace], [x['signed_ee_percent'] for x in trace], label=label, color=color, linestyle=style, linewidth=2)
        ax.set(xlabel='Dimensionless time (tau)', ylabel='Signed ee (%) — R positive', ylim=(-105, 105), xlim=(0, 20), title='Frank model · illustrative parameters')
        ax.set_yticks([-100, -50, 0, 50, 100])
        ax.grid(alpha=0.18)
        ax.legend(loc='upper left', fontsize=9)
        fig.savefig(args.out / 'amplification.svg', metadata={'Date': None})
        plt.close(fig)


if __name__ == '__main__':
    main()