#!/usr/bin/env python3
"""Rebuild the proton projected-range comparison from the retained NIST PSTAR output."""
from __future__ import annotations

import csv
import hashlib
import json
import math
from pathlib import Path
import xml.etree.ElementTree as ET

ROOT = Path(__file__).resolve().parent
SOURCE = ROOT / "nist-pstar-polyethylene-source.txt"
RHO = 0.94  # NIST Polyethylene material preset, g/cm^3
NASA_TEXT_RANGES = {1.0: 0.020, 10.0: 0.12, 100.0: 7.7, 1000.0: 330.0}  # cm; report footnote 7


def read_rows():
    rows = []
    for line in SOURCE.read_text(encoding="utf-8-sig").splitlines():
        fields = line.split()
        if len(fields) == 7 and fields[0][0].isdigit() and "E" in fields[0]:
            energy, stop_e, stop_n, stop_t, csda, projected, detour = map(float, fields)
            rows.append({
                "energy_MeV": energy,
                "stopping_electronic_MeV_cm2_g": stop_e,
                "stopping_nuclear_MeV_cm2_g": stop_n,
                "stopping_total_MeV_cm2_g": stop_t,
                "range_CSDA_g_cm2": csda,
                "range_projected_g_cm2": projected,
                "detour_factor": detour,
                "equivalent_depth_at_0p94_g_cm3_cm": projected / RHO,
            })
    return rows


def comparison(rows):
    lookup = {r["energy_MeV"]: r for r in rows}
    result = []
    for energy, nasa_cm in NASA_TEXT_RANGES.items():
        nist = lookup[energy]["range_projected_g_cm2"]
        nasa_areal = nasa_cm * RHO
        result.append({
            "energy_MeV": energy,
            "nist_projected_range_g_cm2": nist,
            "nist_equivalent_depth_at_0p94_g_cm3_cm": nist / RHO,
            "nasa_report_footnote_range_cm": nasa_cm,
            "nasa_report_footnote_equivalent_areal_range_g_cm2": nasa_areal,
            "nasa_to_nist_areal_ratio": nasa_areal / nist,
        })
    return result


def svg(rows, narrow=False):
    width, height = (820, 1020) if narrow else (1200, 820)
    left, right = (115, 785) if narrow else (130, 1160)
    top, bottom = (170, 690) if narrow else (130, 610)
    xmin, xmax, ymin, ymax = -3.0, 4.0, -6.0, 4.0

    def x(v):
        return left + (math.log10(v) - xmin) / (xmax - xmin) * (right - left)

    def y(v):
        return bottom - (math.log10(v) - ymin) / (ymax - ymin) * (bottom - top)

    ticks_x = [10.0 ** k for k in range(-3, 5)]
    ticks_y = [10.0 ** k for k in range(-6, 5, 2)]
    fs = 22 if narrow else 21
    small = 17 if narrow else 18
    out = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" role="img" aria-labelledby="title desc">',
        '<title id="title">NIST PSTAR proton projected range in polyethylene</title>',
        '<desc id="desc">Logarithmic graph of the NIST PSTAR calculated projected areal range for protons in its polyethylene material preset, over 133 standard energy values. Open orange squares show four approximate ranges stated in NASA Technical Paper 2020-220514 footnote 7, converted using its stated polyethylene density of 0.94 grams per cubic centimetre. At 1 MeV the NASA footnote point is about 8.8 times the NIST point; the other three agree to rounding. Neither series is a radiation-shield effectiveness result.</desc>',
        '<rect width="100%" height="100%" fill="#fff"/>',
        f'<text x="{left}" y="48" font-family="Arial,sans-serif" font-size="{fs+5}" font-weight="700" fill="#152640">One proton, one material, an average range</text>',
        f'<text x="{left}" y="80" font-family="Arial,sans-serif" font-size="{small}" fill="#35465e">NIST PSTAR projected areal range · polyethylene · logarithmic axes</text>',
    ]
    for v in ticks_x:
        xx=x(v)
        out.append(f'<path d="M{xx:.2f} {top}V{bottom}" stroke="#d7deea" stroke-width="1"/>')
        out.append(f'<text x="{xx:.2f}" y="{bottom+30}" text-anchor="middle" font-family="Arial,sans-serif" font-size="{small}" fill="#253750">{v:g}</text>')
    for v in ticks_y:
        yy=y(v)
        out.append(f'<path d="M{left} {yy:.2f}H{right}" stroke="#d7deea" stroke-width="1"/>')
        out.append(f'<text x="{left-12}" y="{yy+6:.2f}" text-anchor="end" font-family="Arial,sans-serif" font-size="{small-1}" fill="#253750">10^{int(math.log10(v))}</text>')
    out.append(f'<path d="M{left} {top}V{bottom}H{right}" fill="none" stroke="#253750" stroke-width="2"/>')
    coords=' '.join(f'{x(r["energy_MeV"]):.2f},{y(r["range_projected_g_cm2"]):.2f}' for r in rows)
    out.append(f'<polyline points="{coords}" fill="none" stroke="#1768a5" stroke-width="4" stroke-linejoin="round"/>')
    for r in rows:
        xx, yy=x(r["energy_MeV"]), y(r["range_projected_g_cm2"])
        out.append(f'<circle cx="{xx:.2f}" cy="{yy:.2f}" r="2.2" fill="#1768a5"/>')
    comp=comparison(rows)
    for c in comp:
        xx=x(c["energy_MeV"]); yy=y(c["nasa_report_footnote_equivalent_areal_range_g_cm2"])
        out.append(f'<rect x="{xx-7:.2f}" y="{yy-7:.2f}" width="14" height="14" fill="#fff" stroke="#c65b22" stroke-width="3"/>')
    label_y=bottom+80 if narrow else bottom+72
    out.append(f'<text x="{(left+right)/2:.1f}" y="{label_y}" text-anchor="middle" font-family="Arial,sans-serif" font-size="{fs}" fill="#152640">Initial proton kinetic energy (MeV)</text>')
    out.append(f'<text transform="translate(34 {(top+bottom)/2:.1f}) rotate(-90)" text-anchor="middle" font-family="Arial,sans-serif" font-size="{fs}" fill="#152640">Projected range (g/cm²)</text>')
    legend_y = label_y + (52 if narrow else 46)
    out.append(f'<path d="M{left} {legend_y}h40" stroke="#1768a5" stroke-width="4"/><circle cx="{left+20}" cy="{legend_y}" r="4" fill="#1768a5"/>')
    out.append(f'<text x="{left+52}" y="{legend_y+6}" font-family="Arial,sans-serif" font-size="{small}" fill="#253750">NIST PSTAR calculated values (133 energies)</text>')
    out.append(f'<rect x="{left}" y="{legend_y+25}" width="14" height="14" fill="white" stroke="#c65b22" stroke-width="3"/>')
    out.append(f'<text x="{left+28}" y="{legend_y+39}" font-family="Arial,sans-serif" font-size="{small}" fill="#253750">NASA report footnote values at 1, 10, 100 and 1,000 MeV</text>')
    if narrow:
        note_y=legend_y+90
        note_size=16
    else:
        note_y=legend_y+72
        note_size=17
    out.append(f'<text x="{left}" y="{note_y}" font-family="Arial,sans-serif" font-size="{note_size}" fill="#8f3d16">At 1 MeV, NASA states 0.2 mm; NIST converts to about 0.023 mm.</text>')
    out.append(f'<text x="{left}" y="{note_y+26}" font-family="Arial,sans-serif" font-size="{note_size}" fill="#40516a">NASA’s plotted Fig. 4 is attributed to SRIM; squares are its separate printed text values.</text>')
    out.append('<text x="10" y="0" visibility="hidden">Original data: NIST PSTAR; NASA/TP-2020-220514, footnote 7. No shield performance calculated.</text>')
    out.append('</svg>')
    return '\n'.join(out)+'\n'


def main():
    rows=read_rows()
    assert len(rows)==133, f"Expected 133 NIST default energy records, found {len(rows)}"
    energies=[r["energy_MeV"] for r in rows]
    assert len(energies)==len(set(energies)) and all(a < b for a,b in zip(energies,energies[1:]))
    points=comparison(rows)
    assert [x["energy_MeV"] for x in points]==[1.0,10.0,100.0,1000.0]
    assert math.isclose(points[0]["nist_equivalent_depth_at_0p94_g_cm3_cm"]*10,0.0226808511,rel_tol=1e-6)
    assert 8.7 < points[0]["nasa_to_nist_areal_ratio"] < 8.9
    for p in points[1:]:
        assert abs(p["nasa_to_nist_areal_ratio"]-1)<0.01
    fields=list(rows[0])
    with (ROOT/"nist-pstar-polyethylene.csv").open("w",newline="",encoding="utf-8") as f:
        w=csv.DictWriter(f,fieldnames=fields); w.writeheader(); w.writerows(rows)
    with (ROOT/"nasa-report-cross-check.csv").open("w",newline="",encoding="utf-8") as f:
        w=csv.DictWriter(f,fieldnames=list(points[0])); w.writeheader(); w.writerows(points)
    (ROOT/"proton-range.svg").write_text(svg(rows),encoding="utf-8")
    (ROOT/"proton-range-phone.svg").write_text(svg(rows,True),encoding="utf-8")
    for name in ["proton-range.svg","proton-range-phone.svg"]:
        ET.parse(ROOT/name)
        text=(ROOT/name).read_text(encoding="utf-8")
        assert "aria-labelledby" in text and "At 1 MeV" in text
    source_hash=hashlib.sha256(SOURCE.read_bytes()).hexdigest()
    metadata={
        "source":"NIST PSTAR online calculator, Polyethylene preset (material 221), default energy grid",
        "source_url":"https://physics.nist.gov/PhysRefData/Star/Text/PSTAR.html",
        "query_date":"2026-09-29",
        "source_text_file":"nist-pstar-polyethylene-source.txt",
        "source_text_sha256":source_hash,
        "n_rows":len(rows),
        "plotted_column":"RANGE(p), projected range in g/cm²",
        "density_g_cm3":RHO,
        "conversion":"equivalent depth in cm = projected areal range (g/cm²) / 0.94 g/cm³",
        "NASA_comparison":"NASA/TP-2020-220514, printed p.4, footnote 7; report figure 4 separately attributes its curve to SRIM/Nikezic et al. (2016). The four printed footnote values are preserved separately from NIST values.",
        "limits":["A mean projected range is not a sharp stopping boundary or shield thickness recommendation.","No particle fluence, mixed space spectrum, secondary-particle transport, dose, dose equivalent, biological effect, spacecraft geometry, or shielding performance is calculated.","NIST and NASA report are model/text comparisons, not a new measurement."],
    }
    (ROOT/"study.json").write_text(json.dumps(metadata,indent=2)+"\n",encoding="utf-8")
    print(json.dumps({"nist_points":len(rows),"comparison":points,"source_sha256":source_hash},indent=2))

if __name__=="__main__":
    main()