"""The data set's balloon files processed the paper's way, against the paper's Table 2 ("Our analysis").

Usage: python3 -I sonde_table2.py DATA_DIR [OUT.json]
Baird et al. (2026), Space Weather 24, e2026SW005279, §2.3.2 (p. 6): a Savitzky-Golay filter (2nd order polynomial,
25 data points), the ascent only, flight levels rounded to the nearest 20, and the mean of the filtered data at each
level. Table 2 (p. 10) gives the ambient dose equivalent rate and absorbed dose rate at FL360, 420, 480 and 600.
Data: Zenodo version 2, doi:10.5281/zenodo.23082794, CC BY 4.0, read as data. Standard library only.
The filter is applied where its full 25-point window fits; nearer the ends the raw samples are used.
"""
import csv
import json
import os
import statistics as st
import sys

# Table 2, p. 10: ambient dose equivalent rate (uSv/h) and absorbed dose rate (uGy/h) at FL360, 420, 480, 600
TABLE2 = {
    'Lerwick_Radiosonde_Flight_1_GLE77.csv': {'h': [12.3, 21.8, 35.4, 160.0], 'd': [2.0, 3.2, 4.3, 14.9]},
    'Lerwick_Radiosonde_Flight_2_GLE77.csv': {'h': [8.0, 20.6, 36.0, 122.1], 'd': [2.0, 3.2, 4.8, 14.0]},
    'Camborne_Radiosonde_GLE77.csv': {'h': [2.8, 8.2, 9.9, 13.3], 'd': [1.1, 1.8, 2.2, 2.8]},
    'DeBilt_Radiosonde_GLE77.csv': {'h': [4.4, 7.4, 8.2, 14.2], 'd': [1.4, 1.8, 2.2, 2.8]},
}
LEVELS = [360, 420, 480, 600]
M = 12   # half-width: 25 points
COEF = [3 * (3 * M * M + 3 * M - 1 - 5 * i * i) / ((2 * M + 3) * (2 * M + 1) * (2 * M - 1)) for i in range(-M, M + 1)]


def sg(v):
    out = list(v)
    for k in range(M, len(v) - M):
        out[k] = sum(c * v[k + j - M] for j, c in enumerate(COEF))
    return out


def main():
    sys.stdout.reconfigure(encoding='utf-8')
    d = sys.argv[1]
    res = {}
    print('Flight level: paper Table 2 / data set, H*(10) rate in uSv/h (ratio); absorbed dose rate in uGy/h (ratio)')
    for f, paper in TABLE2.items():
        with open(os.path.join(d, f), newline='', encoding='utf-8-sig') as fh:
            rows = list(csv.reader(fh))[1:]
        fl = [float(r[5]) for r in rows]
        top = max(range(len(fl)), key=lambda i: fl[i])   # ascent: up to the highest flight level
        asc = rows[:top + 1]
        h = sg([float(r[7]) for r in asc])
        a = sg([float(r[8]) for r in asc])
        lv = [20 * round(float(r[5]) / 20) for r in asc]
        res[f] = {}
        print(f'\n{f}')
        for i, L in enumerate(LEVELS):
            hs = [x for x, l in zip(h, lv) if l == L]
            ds = [x for x, l in zip(a, lv) if l == L]
            if not hs:
                print(f'  FL{L}: no samples')
                continue
            hm, dm = st.mean(hs), st.mean(ds)
            res[f][L] = {'h_data': round(hm, 1), 'h_paper': paper['h'][i], 'h_ratio': round(hm / paper['h'][i], 2),
                         'd_data': round(dm, 2), 'd_paper': paper['d'][i], 'd_ratio': round(dm / paper['d'][i], 2),
                         'n': len(hs)}
            r = res[f][L]
            print(f'  FL{L}: {r["h_paper"]} / {r["h_data"]} (x{r["h_ratio"]});  {r["d_paper"]} / {r["d_data"]} '
                  f'(x{r["d_ratio"]});  n={r["n"]}')
    ratios = [v['h_ratio'] for f in res for v in res[f].values()]
    print(f'\nH*(10) ratios, data set / Table 2: {min(ratios)}-{max(ratios)}')
    for f in res:
        rr = [v['h_ratio'] for v in res[f].values()]
        dd = [v['d_ratio'] for v in res[f].values()]
        print(f'  {f}: H*(10) x{min(rr)}-{max(rr)}; absorbed x{min(dd)}-{max(dd)}')
    if len(sys.argv) > 2:
        with open(sys.argv[2], 'w', encoding='utf-8', newline='\n') as fh:
            json.dump(res, fh, indent=1)


if __name__ == '__main__':
    main()