"""Parse IRI 'Model-Based Probabilistic ENSO Forecast' SVGs (ensoforecast.iri.columbia.edu/figure3_plot/Y/M0,
saved as raw/f3/YYYY-MM.svg) -> data/objective_probs.csv.
Bars are matplotlib rect patches; probability = (y0 - y_top) / (y0 - y100) * 100 using the 0 and 100 y-tick positions.
Fill colour identifies the category; x order gives the season (labels from the x-tick text comments)."""
import re, glob, os, csv
R = os.environ.get('ROOT', '/work/enso-objective')
SEAS = ['DJF', 'JFM', 'FMA', 'MAM', 'AMJ', 'MJJ', 'JJA', 'JAS', 'ASO', 'SON', 'OND', 'NDJ']
MON = 'Jan Feb Mar Apr May Jun Jul Aug Sep Oct Nov Dec'.split()
rows, bad = [], []
for f in sorted(glob.glob(f'{R}/raw/f3/*.svg')):
    s = open(f, errors='replace').read()
    if not s.startswith('<?xml'): continue
    tag = os.path.basename(f)[:-4]
    mt = re.search(r'<!-- (\S+ \d{4}) CCSR/IRI Model-Based', s)  # e.g. 'Mid-Jun 2018'
    if not mt: bad.append((tag, 'no title')); continue
    title = mt.group(1)
    # y ticks: <!-- 0 --> ... <!-- 100 --> inside ytick groups; positions from the <use ... y=> just before
    yt = {}
    for m in re.finditer(r'<g id="ytick_\d+">.*?<use [^>]*y="([\d.]+)".*?<!-- (\d+) -->', s, re.S):
        yt[int(m.group(2))] = float(m.group(1))
    y0, y100 = yt[0], yt[100]
    xt = []
    for m in re.finditer(r'<g id="xtick_\d+">.*?<use [^>]*x="([\d.]+)".*?<!-- (\w+) -->', s, re.S):
        xt.append((float(m.group(1)), m.group(2)))
    bars = []
    for m in re.finditer(r'<g id="patch_\d+">\s*<path d="M ([\d.]+) ([\d.]+) \s*L ([\d.]+) ([\d.]+) \s*L ([\d.]+) ([\d.]+) \s*L ([\d.]+) ([\d.]+) \s*z\s*"[^>]*style="fill: (#[0-9a-f]{6})', s):
        v = [float(m.group(i)) for i in range(1, 9)]
        xs, ys = v[0::2], v[1::2]
        if abs(max(ys) - y0) > 0.01: continue  # legend boxes etc.
        bars.append(((min(xs) + max(xs)) / 2, (y0 - min(ys)) / (y0 - y100) * 100, m.group(9)))
    cols = sorted({b[2] for b in bars})
    by = {}
    for xc, p, c in bars:
        sea = min(xt, key=lambda t: abs(t[0] - xc))[1]
        by.setdefault(sea, {})[c] = p
    seasons = [t[1] for t in sorted(xt)]
    for k, sea in enumerate(seasons):
        d = by.get(sea, {})
        rows.append(dict(issue=tag, title=title, k=k, season=sea, colours='|'.join(sorted(d)),
                         **{c: round(d[c], 2) for c in d}))
    if len(bars) != 3 * len(seasons): bad.append((tag, len(bars), len(seasons)))
cols = sorted({k for r in rows for k in r if k.startswith('#')})
print('colours seen:', cols)
with open(f'{R}/data/objective_probs_raw.csv', 'w', newline='') as fh:
    w = csv.DictWriter(fh, fieldnames=['issue', 'title', 'k', 'season', 'colours'] + cols); w.writeheader(); w.writerows(rows)
print(len({r['issue'] for r in rows}), 'issues', len(rows), 'rows; bad', bad[:10], len(bad))