"""Area weights for the regional aggregation (item 6): fractional cell-polygon overlap x cos(lat).

Uses Vietnam7.shp (7 polygons named S1..S7 -> region ids 1..7, verified against the gridded reg
mask and the region names) with shapely for the fractional overlap of each 0.25-deg grid cell with
its region polygon, multiplied by cos(latitude), so boundary/coastal cells are partial and large
regions are not advantaged by cell count. Falls back to cos(lat) on whole-cell gridded `reg`
membership if the shapefile cannot be read (documented in output metadata). geopandas is absent in
the env; pyshp + shapely are used directly.
"""
from __future__ import annotations
import os
import numpy as np


def _load_polygons(shp_path):
    import shapefile
    from shapely.geometry import shape
    sf = shapefile.Reader(shp_path)
    names = [r["Name"] if isinstance(r, dict) else r[sf.fields.index(("Name", "C", 254, 0)) - 1]
             for r in sf.records()]
    # robust name fetch
    fld = [f[0] for f in sf.fields[1:]]
    ni = fld.index("Name")
    polys = {}
    for i, rec in enumerate(sf.records()):
        nm = rec[ni]                       # 'S1'..'S7'
        rid = int(str(nm).lstrip("Ss"))
        polys[rid] = shape(sf.shape(i).__geo_interface__)
    return polys


def cell_area_weights(lat, lon, region_of_cell, iy, ix, shp_path, res=0.25):
    """Return w[ncell] = fractional overlap of the cell with its region polygon x cos(lat),
    and a string describing the rule used. region_of_cell gives each land cell's reg id (1..7)."""
    ncell = iy.size
    coslat = np.cos(np.deg2rad(lat[iy]))
    try:
        from shapely.geometry import box
        polys = _load_polygons(shp_path)
        w = np.zeros(ncell)
        h = res / 2.0
        for k in range(ncell):
            rid = int(region_of_cell[k])
            cy, cx = lat[iy[k]], lon[ix[k]]
            cell = box(cx - h, cy - h, cx + h, cy + h)
            poly = polys.get(rid)
            frac = 1.0
            if poly is not None:
                inter = cell.intersection(poly).area
                frac = float(inter / cell.area) if cell.area > 0 else 1.0
            w[k] = frac * coslat[k]
        rule = ("fractional cell-polygon overlap (Vietnam7.shp S1..S7 -> R1..R7) x cos(lat)")
        # guard: if overlaps are degenerate (all ~0), fall back
        if np.nanmax(w) <= 0:
            raise ValueError("degenerate overlaps")
        return w, rule
    except Exception as e:
        return coslat.copy(), f"cos(lat) on gridded reg mask (fallback: {type(e).__name__})"