#!/usr/bin/env python3
"""How often NOAA's geomagnetic storm levels G1-G5 are reached, counted in GFZ's Kp record since 1932.
Fact-Checked Science, "Earth's magnetic field" guide, section 7 ("Our analysis"). Standard library only.

Run it
  python3 kp_gscale_counts.py [folder]
  The folder (default: the script's own) must hold, or will receive (5.5 MB):
    Kp_ap_Ap_SN_F107_since_1932.txt   https://kp.gfz.de/app/files/Kp_ap_Ap_SN_F107_since_1932.txt
  Kp is from Geomagnetic Observatory Niemegk, GFZ Helmholtz Centre for Geosciences, CC BY 4.0: Matzka, Stolle,
  Yamazaki, Bronkalla and Morschhauser (2021), Space Weather, https://doi.org/10.1029/2020SW002641; data
  https://doi.org/10.5880/Kp.0001. The file grows daily and its newest values are nowcasts; the means over the
  complete cycles 17-24 use definitive values only, so they should not change. The file's sunspot numbers
  (CC BY-NC 4.0) are not used. It prints the counts and writes kp-gscale-counts.json.

NOAA's G levels by Kp: G1 Kp = 5, G2 6, G3 7, G4 8 including 9-, G5 9. With Kp in thirds (5- = 4.667, 5o = 5, 5+ =
5.333 ...): G1 4.667-5.333, G2 5.667-6.333, G3 6.667-7.333, G4 7.667-8.667, G5 9. For each solar cycle (starting in
the month of the smoothed sunspot minimum, as listed by WDC-SILSO) it counts the 3-hour intervals at each level and
the days whose highest Kp reached each level, ranks the days by daily Ap and lists the days with Kp 9 since 1985.
"""
import json
import os
import sys
import urllib.request
from collections import defaultdict

NAME = "Kp_ap_Ap_SN_F107_since_1932.txt"
URL = "https://kp.gfz.de/app/files/" + NAME
CYCLE_START = {17: (1933, 9), 18: (1944, 2), 19: (1954, 4), 20: (1964, 10), 21: (1976, 3),
               22: (1986, 9), 23: (1996, 8), 24: (2008, 12), 25: (2019, 12)}
DAYS = ("1989-03-13", "1989-03-14", "2003-10-29", "2003-10-30", "2003-11-20", "2022-02-03", "2022-02-04",
        "2024-05-10", "2024-05-11", "2024-10-10", "2024-10-11")


def level(kp):
    for lv, lo in ((5, 8.99), (4, 7.66), (3, 6.66), (2, 5.66), (1, 4.66)):
        if kp >= lo:
            return lv
    return 0


def cycle_of(y, m):
    c = None
    for k in sorted(CYCLE_START):
        if (y, m) >= CYCLE_START[k]:
            c = k
    return c


def main():
    sys.stdout.reconfigure(encoding="utf-8")
    folder = sys.argv[1] if len(sys.argv) > 1 else os.path.dirname(os.path.abspath(__file__))
    path = os.path.join(folder, NAME)
    if not os.path.exists(path):
        print("downloading", URL)
        req = urllib.request.Request(URL, headers={"User-Agent": "factchecked.science kp_gscale_counts.py"})
        with urllib.request.urlopen(req, timeout=300) as r, open(path, "wb") as f:
            f.write(r.read())
    intervals = defaultdict(lambda: defaultdict(int))
    days = defaultdict(lambda: defaultdict(int))
    ndays = defaultdict(int)
    ap_days, shown, last, kp9 = [], {}, None, []
    with open(path, encoding="utf-8", errors="replace") as f:
        for line in f:
            if line.startswith("#") or not line.strip():
                continue
            p = line.split()
            y, m, d = int(p[0]), int(p[1]), int(p[2])
            kps = [float(x) for x in p[7:15]]
            if any(k < 0 for k in kps):
                continue
            key = f"{y:04d}-{m:02d}-{d:02d}"
            ap = int(p[23])
            ap_days.append((ap, key))
            if max(kps) >= 8.99:
                kp9.append(key)
            if key in DAYS:
                shown[key] = {"Kp": kps, "max_level": level(max(kps)), "Ap": ap}
            c = cycle_of(y, m)
            if c is None:
                continue
            ndays[c] += 1
            for k in kps:
                if level(k):
                    intervals[c][level(k)] += 1
            if level(max(kps)):
                days[c][level(max(kps))] += 1
            last = key
    complete = [c for c in sorted(ndays) if c + 1 in CYCLE_START]
    out = {"record_ends": last, "complete_cycles": complete, "per_cycle": {}, "mean_per_complete_cycle": {}, "days": shown}
    print("cycle  days   3-hour intervals at G1..G5             days whose highest level was G1..G5")
    for c in sorted(ndays):
        out["per_cycle"][c] = {"days": ndays[c], "intervals": [intervals[c][g] for g in range(1, 6)], "days_at_level": [days[c][g] for g in range(1, 6)]}
        print(f"{c:5d} {ndays[c]:5d}   " + " ".join(f"{intervals[c][g]:5d}" for g in range(1, 6)) + "     "
              + " ".join(f"{days[c][g]:5d}" for g in range(1, 6)))
    print("mean per complete cycle, cycles", complete[0], "to", complete[-1])
    for g in range(1, 6):
        mi = sum(intervals[c][g] for c in complete) / len(complete)
        md = sum(days[c][g] for c in complete) / len(complete)
        ma = sum(sum(days[c][h] for h in range(g, 6)) for c in complete) / len(complete)
        out["mean_per_complete_cycle"][f"G{g}"] = {"three_hour_intervals": round(mi, 1), "days_with_this_as_highest": round(md, 1),
                                                   "days_at_or_above": round(ma, 1)}
        print(f"  G{g}: {mi:7.1f} three-hour intervals, {md:6.1f} days with this as the highest level, {ma:6.1f} days at or above")
    ap_days.sort(reverse=True)
    out["top_daily_Ap"] = [{"date": k, "Ap": a} for a, k in ap_days[:10]]
    out["days_with_Kp_9_since_1985"] = [k for k in kp9 if k >= "1985"]
    print("days with Kp 9 since 1985:", ", ".join(out["days_with_Kp_9_since_1985"]))
    print("highest daily Ap since 1932:", ", ".join(f"{k} {a}" for a, k in ap_days[:10]))
    for k in sorted(shown):
        print(k, "Kp", " ".join(f"{v:.3f}" for v in shown[k]["Kp"]), "| highest level G", shown[k]["max_level"], "| Ap", shown[k]["Ap"])
    print("record ends", last)
    with open(os.path.join(folder, "kp-gscale-counts.json"), "w", encoding="utf-8", newline="\n") as f:
        json.dump(out, f, indent=1)
    print("Wrote", os.path.join(folder, "kp-gscale-counts.json"))


if __name__ == "__main__":
    main()