"""Hours to site-days, exactly as `rule.py` defines them. Step 3 of five.

Reads ~/data/red-sky/*.json, writes ~/data/red-sky/site-days.csv, and prints
the three counts the rule says to publish either way: site-days kept, dropped
for a missing hour, and carrying each label.

    python research/red-sky/load.py [--distance 300]

It computes the labels and the outcome and NOTHING ELSE. No comparison between
them lives in this file; that is `controls.py` and then `analyse.py`, in that
order, because the controls have to be printed before the real statistic
exists.

ONE DEPARTURE FROM THE LETTER OF THE RULE, RECORDED RATHER THAN HIDDEN.
`rule.py` says a site-day is dropped if any hour it needs is missing. Applied
literally that would drop a day from BOTH arms when only one arm's outcome
window is incomplete — which is every site's last day, whose `fine_next` lies
outside the fetched period. Here a missing outcome is written as -1 and the
arm that needs it excludes the row; the other arm keeps it. The printed counts
give both, so the effect of the choice is visible and small. No label
definition is touched.

One row per site-day:

    site, date, month, h_rise, h_set,
    clear_w, lit, red_night,          (the night construct and its two halves)
    clear_e, lit_rise, red_morning,   (the morning arm and its two halves)
    clear_e_at_set,                   (D1, the downwind dummy at sunset)
    fine_next, fine_today,            (the outcomes; -1 where undefined)
    dropped_reason                    ('' when the row is usable)
"""
import csv
import json
import os
import sys
from datetime import date, timedelta

OUT = os.path.expanduser("~/data/red-sky")
SITES = ["london", "manchester", "edinburgh", "cardiff", "belfast", "plymouth",
         "aberdeen", "inverness", "carlisle", "hull", "norwich", "swansea",
         "dublin", "galway"]
CHUNK_YEARS = ["2015", "2020"]

# rule.py, fixed before the data:
THIN = 25          # per cent, low and mid cloud at the offset point
LIT = 30           # per cent, high cloud over the site
FINE_MM = 1.0      # mm over the twelve-hour window
WINDOW = 12        # hours


def load_point(name):
    """{'YYYY-MM-DDTHH:00': {var: value}} merged over the chunks, or None."""
    hours, daily = {}, {}
    for year in CHUNK_YEARS:
        path = os.path.join(OUT, f"{name}_{year}.json")
        if not os.path.exists(path):
            return None, None
        with open(path, encoding="utf-8") as fh:
            j = json.load(fh)
        h = j["hourly"]
        variables = [k for k in h if k != "time"]
        for i, t in enumerate(h["time"]):
            hours[t] = {v: h[v][i] for v in variables}
        if "daily" in j:
            d = j["daily"]
            for i, t in enumerate(d["time"]):
                daily[t] = {"sunrise": d["sunrise"][i], "sunset": d["sunset"][i]}
    return hours, daily


def hour_key(day, hour):
    return f"{day}T{hour:02d}:00"


def window_sum(hours, start_day, start_hour, n, var):
    """Sum `var` over n hours from (start_day, start_hour), or None if any is."""
    total = 0.0
    d = date.fromisoformat(start_day)
    h = start_hour
    for _ in range(n):
        if h > 23:
            h -= 24
            d += timedelta(days=1)
        row = hours.get(hour_key(d.isoformat(), h))
        if row is None or row.get(var) is None:
            return None
        total += row[var]
        h += 1
    return total


def thin(hours, day, hour):
    """The offset-point condition: low AND mid cloud thin at that hour."""
    row = hours.get(hour_key(day, hour))
    if row is None:
        return None
    lo, mid = row.get("cloud_cover_low"), row.get("cloud_cover_mid")
    if lo is None or mid is None:
        return None
    return lo <= THIN and mid <= THIN


def main():
    km = 300
    if "--distance" in sys.argv:
        km = int(sys.argv[sys.argv.index("--distance") + 1])

    rows, kept, dropped, reasons = [], 0, 0, {}
    for site in SITES:
        site_h, site_d = load_point(site)
        west_h, _ = load_point(f"{site}_W{km}")
        east_h, _ = load_point(f"{site}_E{km}")
        if site_h is None or west_h is None or east_h is None:
            print(f"SKIP {site}: a point is not downloaded at {km} km")
            continue

        for day in sorted(site_d):
            rec = {"site": site, "date": day, "month": int(day[5:7])}
            why = ""
            sun = site_d[day]
            h_rise = int(sun["sunrise"][11:13])
            h_set = int(sun["sunset"][11:13])
            rec["h_rise"], rec["h_set"] = h_rise, h_set

            clear_w = thin(west_h, day, h_set)
            clear_e_at_set = thin(east_h, day, h_set)
            clear_e = thin(east_h, day, h_rise)
            set_row = site_h.get(hour_key(day, h_set))
            rise_row = site_h.get(hour_key(day, h_rise))
            lit_set = None if not set_row or set_row.get("cloud_cover_high") is None \
                else set_row["cloud_cover_high"] >= LIT
            lit_rise = None if not rise_row or rise_row.get("cloud_cover_high") is None \
                else rise_row["cloud_cover_high"] >= LIT

            # The outcomes. FINE_NEXT is tomorrow 06:00..17:59; FINE_TODAY is
            # the twelve hours beginning at h_rise + 1 (rule.py).
            tomorrow = (date.fromisoformat(day) + timedelta(days=1)).isoformat()
            rain_next = window_sum(site_h, tomorrow, 6, WINDOW, "precipitation")
            rain_today = window_sum(site_h, day, h_rise + 1, WINDOW, "precipitation")

            if None in (clear_w, clear_e, clear_e_at_set, lit_set, lit_rise):
                why = "missing cloud hour"
            elif rain_next is None and rain_today is None:
                why = "missing outcome window"

            rec["clear_w"] = int(bool(clear_w)) if clear_w is not None else -1
            rec["lit"] = int(bool(lit_set)) if lit_set is not None else -1
            rec["red_night"] = 1 if (clear_w and lit_set) else 0 if why == "" else -1
            rec["clear_e"] = int(bool(clear_e)) if clear_e is not None else -1
            rec["lit_rise"] = int(bool(lit_rise)) if lit_rise is not None else -1
            rec["red_morning"] = 1 if (clear_e and lit_rise) else 0 if why == "" else -1
            rec["clear_e_at_set"] = int(bool(clear_e_at_set)) \
                if clear_e_at_set is not None else -1
            rec["fine_next"] = -1 if rain_next is None else int(rain_next < FINE_MM)
            rec["fine_today"] = -1 if rain_today is None else int(rain_today < FINE_MM)
            rec["dropped_reason"] = why
            rows.append(rec)
            if why:
                dropped += 1
                reasons[why] = reasons.get(why, 0) + 1
            else:
                kept += 1

    path = os.path.join(OUT, f"site-days-{km}.csv")
    with open(path, "w", newline="", encoding="utf-8") as fh:
        w = csv.DictWriter(fh, fieldnames=list(rows[0]))
        w.writeheader()
        w.writerows(rows)

    usable = [r for r in rows if not r["dropped_reason"]]
    print(f"\n{path}")
    print(f"site-days      {len(rows)}   kept {kept}   dropped {dropped} "
          f"({100 * dropped / max(len(rows), 1):.2f} %)")
    for why, n in sorted(reasons.items()):
        print(f"  dropped: {why}  {n}")
    print(f"RED_NIGHT      {sum(r['red_night'] == 1 for r in usable)}")
    print(f"  CLEAR_WEST   {sum(r['clear_w'] == 1 for r in usable)}")
    print(f"  LIT_SKY      {sum(r['lit'] == 1 for r in usable)}")
    print(f"RED_MORNING    {sum(r['red_morning'] == 1 for r in usable)}")
    print(f"  CLEAR_EAST   {sum(r['clear_e'] == 1 for r in usable)}")
    print(f"  LIT_AT_RISE  {sum(r['lit_rise'] == 1 for r in usable)}")
    print(f"fine_next known {sum(r['fine_next'] != -1 for r in usable)}, "
          f"fine_today known {sum(r['fine_today'] != -1 for r in usable)}")
    print("\nNo statistic relating a label to an outcome is computed here. "
          "Next: controls.py, then analyse.py.")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())