"""Aggregate NASA FIRMS fire detections and NASA POWER (MERRA-2) weather to Ukraine oblast x month.

Usage: python3 -I process.py <firms_dir> <power_dir> <ne_admin1_geojson> <out_dir>

Outputs (out_dir):
  regions.geojson        25 regions (24 oblasts + Crimea; Kyiv City merged into Kyiv oblast,
                         Sevastopol into Crimea), internationally recognised borders
  fires_monthly.csv      region, sensor, year, month, n (detections), frp_sum
  weather_monthly.csv    region, year, month, T2M_MAX, PRECTOTCORR, RH2M, GWETTOP
  qa.json                filtering counts per sensor-year

Filtering: FIRMS 'type' == 0 (presumed vegetation fire; drops static industrial sources,
volcanoes, offshore). MODIS confidence >= 30; VIIRS confidence in {n, h}.
Only aggregate counts leave this script: no point coordinates are written out.
"""
import glob
import json
import os
import sys

import geopandas as gpd
import pandas as pd

NAME_FIX = {
    "UA-74": "Chernihiv", "UA-07": "Volyn", "UA-56": "Rivne", "UA-18": "Zhytomyr",
    "UA-32": "Kyiv", "UA-30": "Kyiv", "UA-21": "Zakarpattia", "UA-77": "Chernivtsi",
    "UA-26": "Ivano-Frankivsk", "UA-51": "Odesa", "UA-05": "Vinnytsia", "UA-46": "Lviv",
    "UA-59": "Sumy", "UA-63": "Kharkiv", "UA-09": "Luhansk", "UA-14": "Donetsk",
    "UA-43": "Crimea", "UA-40": "Crimea", "UA-65": "Kherson", "UA-23": "Zaporizhzhia",
    "UA-48": "Mykolaiv", "UA-53": "Poltava", "UA-68": "Khmelnytskyi", "UA-61": "Ternopil",
    "UA-12": "Dnipropetrovsk", "UA-71": "Cherkasy", "UA-35": "Kirovohrad",
}


def regions(ne_path):
    g = gpd.read_file(ne_path)
    g = g[g["iso_3166_2"].isin(NAME_FIX.keys())].copy()
    g["region"] = g["iso_3166_2"].map(NAME_FIX)
    r = g.dissolve(by="region", as_index=False)[["region", "geometry"]]
    assert len(r) == 25, len(r)
    return r


def load_firms(firms_dir):
    frames, qa = [], {}
    for f in sorted(glob.glob(os.path.join(firms_dir, "*_Ukraine.csv"))):
        base = os.path.basename(f)
        sensor = "MODIS" if base.startswith("modis") else "VIIRS"
        year = int(base.split("_")[1])
        d = pd.read_csv(f, usecols=["latitude", "longitude", "acq_date", "confidence", "frp", "type"],
                        dtype={"confidence": str})
        n0 = len(d)
        d = d[d["type"] == 0]
        if sensor == "MODIS":
            d = d[pd.to_numeric(d["confidence"], errors="coerce") >= 30]
        else:
            d = d[d["confidence"].str.lower().isin(["n", "h", "nominal", "high"])]
        qa[f"{sensor}_{year}"] = {"raw": int(n0), "kept": int(len(d))}
        d = d.assign(sensor=sensor)
        frames.append(d)
    return pd.concat(frames, ignore_index=True), qa


def main(firms_dir, power_dir, ne_path, out_dir):
    os.makedirs(out_dir, exist_ok=True)
    reg = regions(ne_path)
    reg.to_file(os.path.join(out_dir, "regions.geojson"), driver="GeoJSON")

    fires, qa = load_firms(firms_dir)
    pts = gpd.GeoDataFrame(fires, geometry=gpd.points_from_xy(fires["longitude"], fires["latitude"]),
                           crs="EPSG:4326")
    j = gpd.sjoin(pts, reg, how="inner", predicate="within")
    qa["outside_regions"] = int(len(pts) - len(j))
    dt = pd.to_datetime(j["acq_date"])
    j = j.assign(year=dt.dt.year, month=dt.dt.month)
    fm = (j.groupby(["region", "sensor", "year", "month"])
            .agg(n=("frp", "size"), frp_sum=("frp", "sum")).reset_index())
    # complete grid so months with zero detections are explicit zeros
    years = {s: sorted(fm.loc[fm.sensor == s, "year"].unique()) for s in fm.sensor.unique()}
    full = pd.concat([
        pd.MultiIndex.from_product([reg.region, [s], ys, range(1, 13)],
                                   names=["region", "sensor", "year", "month"]).to_frame(index=False)
        for s, ys in years.items()])
    fm = full.merge(fm, how="left").fillna({"n": 0, "frp_sum": 0.0})
    fm.to_csv(os.path.join(out_dir, "fires_monthly.csv"), index=False)

    # weather: average MERRA-2 grid points inside each region (nearest grid point if none inside)
    wx = None
    for p in ["T2M_MAX", "PRECTOTCORR", "RH2M", "GWETTOP"]:
        w = pd.read_csv(os.path.join(power_dir, f"power_{p}.csv"))
        w = w[w["value"] > -900]
        grid = w[["lat", "lon"]].drop_duplicates()
        gg = gpd.GeoDataFrame(grid, geometry=gpd.points_from_xy(grid.lon, grid.lat), crs="EPSG:4326")
        inside = gpd.sjoin(gg, reg, how="inner", predicate="within")[["lat", "lon", "region"]]
        missing = set(reg.region) - set(inside.region)
        if missing:
            cent = reg[reg.region.isin(missing)].copy()
            cent["geometry"] = cent.geometry.representative_point()
            near = gpd.sjoin_nearest(cent.to_crs(3857), gg.to_crs(3857))[["region", "lat", "lon"]]
            inside = pd.concat([inside, near])
        m = w.merge(inside, on=["lat", "lon"])
        m["year"] = m["yyyymm"] // 100
        m["month"] = m["yyyymm"] % 100
        a = m.groupby(["region", "year", "month"])["value"].mean().rename(p).reset_index()
        wx = a if wx is None else wx.merge(a, on=["region", "year", "month"])
    wx.to_csv(os.path.join(out_dir, "weather_monthly.csv"), index=False)

    with open(os.path.join(out_dir, "qa.json"), "w", encoding="utf-8") as fh:
        json.dump(qa, fh, indent=1)
    print("regions", len(reg), "| fire rows", len(fm), "| weather rows", len(wx),
          "| outside", qa["outside_regions"])
    print(fm.groupby(["sensor", "year"])["n"].sum().unstack(0).to_string())


if __name__ == "__main__":
    main(*sys.argv[1:5])