#!/usr/bin/env python3
"""把三份公开数据落到同一张 2 弧分网格上：重力（GGMplus）、磁（EMAG2v3）、地质图标注（USGS geo3al）。
只做数据对齐与特征计算，不训练任何模型。

用法: python3 00_build_grid.py <raw_dir> <out_npz> <out_summary_json>
"""
import sys, json, hashlib
from pathlib import Path
import numpy as np
import tifffile, shapefile, pyproj
from matplotlib.path import Path as MPath

RAW = Path(sys.argv[1]); OUT = Path(sys.argv[2]); SUMMARY = Path(sys.argv[3])
LON0, LON1, LAT0, LAT1 = 105.0, 120.0, 35.0, 50.0
STEP = 1.0 / 30.0                      # 2 弧分，EMAG2v3 原生分辨率
NX = int(round((LON1 - LON0) / STEP)); NY = int(round((LAT1 - LAT0) / STEP))
lon_c = LON0 + np.arange(NX) * STEP    # 格心落在 1/30 度整数倍上，与 EMAG2 像元中心重合
lat_c = LAT0 + np.arange(NY) * STEP
summary = {"region": {"lon": [LON0, LON1], "lat": [LAT0, LAT1], "step_deg": STEP, "nx": NX, "ny": NY}}

def sha(p):
    h = hashlib.sha256(); h.update(Path(p).read_bytes()); return h.hexdigest()

# ---------- 磁：EMAG2v3 UpCont（海拔 4 km 上延），2 弧分 ----------
tif = RAW / "EMAG2_V3_20170530_UpCont.tif"
t = tifffile.TiffFile(tif); page = t.pages[0]
scale = page.tags["ModelPixelScaleTag"].value; tie = page.tags["ModelTiepointTag"].value
gk = page.tags.get(34735)
geokeys = gk.value if gk is not None else ()
raster_type = None
for k in range(4, len(geokeys), 4):
    if geokeys[k] == 1025:              # GTRasterTypeGeoKey: 1=PixelIsArea 2=PixelIsPoint
        raster_type = geokeys[k + 3]
emag = page.asarray().astype("float64")
emag[emag < -1e30] = np.nan
x0, y0 = tie[3], tie[4]
if raster_type == 2:                    # tiepoint 就是像元中心
    cx0, cy0 = x0, y0
else:                                   # tiepoint 是像元左上角
    cx0, cy0 = x0 + scale[0] / 2, y0 - scale[1] / 2
cols = np.rint((lon_c - cx0) / scale[0]).astype(int)
rows = np.rint((cy0 - lat_c) / scale[1]).astype(int)
assert np.allclose(cx0 + cols * scale[0], lon_c, atol=1e-6), "EMAG2 像元中心与网格没对齐"
assert np.allclose(cy0 - rows * scale[1], lat_c, atol=1e-6), "EMAG2 像元中心与网格没对齐"
mag = emag[np.ix_(rows, cols)]          # [ny, nx]，行随纬度递增
summary["emag2"] = {"file": tif.name, "sha256": sha(tif), "shape": list(emag.shape),
                    "pixel_scale": list(scale[:2]), "tiepoint": list(tie[3:5]),
                    "raster_type_geokey": raster_type, "first_center": [cx0, cy0],
                    "nan_fraction_in_region": float(np.isnan(mag).mean())}
del emag

# ---------- 重力：GGMplus 地面重力扰动 dg，0.002 度，5x5 度分块 ----------
N = 2500
sum_g = np.zeros((NY, NX)); cnt_g = np.zeros((NY, NX))
tiles = {}
seam = {}
for la in (35, 40, 45):
    for lo in (105, 110, 115):
        name = f"N{la:02d}E{lo:03d}.dg"; p = RAW / name
        raw = np.fromfile(p, dtype=">i2").astype("float64")
        assert raw.size == N * N
        raw[raw == -32768] = np.nan
        raw /= 10.0                      # 0.1 mGal -> mGal
        tiles[(la, lo)] = raw
        summary.setdefault("ggmplus_tiles", {})[name] = sha(p)

def orient(raw, mode):
    # 读我文件：记录沿经线由南向北、列由西向东。两种读法都试，选拼接缝更连续的那一种，结果写进 summary。
    a = raw.reshape(N, N)
    return a.T if mode == "row_major" else a     # 返回 [lat_idx(南->北), lon_idx(西->东)]

def seam_score(mode):
    d = []
    for la in (35, 40, 45):
        for lo in (105, 110):
            A = orient(tiles[(la, lo)], mode); B = orient(tiles[(la, lo + 5)], mode)
            d.append(np.nanmean(np.abs(A[:, -1] - B[:, 0])))
    for la in (35, 40):
        for lo in (105, 110, 115):
            A = orient(tiles[(la, lo)], mode); B = orient(tiles[(la + 5, lo)], mode)
            d.append(np.nanmean(np.abs(A[-1, :] - B[0, :])))
    inner = []
    for key in tiles:
        A = orient(tiles[key], mode)
        inner.append(np.nanmean(np.abs(np.diff(A, axis=0)))); inner.append(np.nanmean(np.abs(np.diff(A, axis=1))))
    return float(np.nanmean(d)), float(np.nanmean(inner))

for mode in ("col_major", "row_major"):
    seam[mode] = seam_score(mode)
mode = min(seam, key=lambda m: seam[m][0])
summary["ggmplus_orientation"] = {"seam_mean_abs_jump_mGal_vs_inner_step": seam, "chosen": mode}

pix_off = 0.001 + np.arange(N) * 0.002
for (la, lo), raw in tiles.items():
    A = orient(raw, mode)
    plat = la + pix_off; plon = lo + pix_off
    iy = np.floor((plat - LAT0 + STEP / 2) / STEP).astype(int)
    ix = np.floor((plon - LON0 + STEP / 2) / STEP).astype(int)
    oky = (iy >= 0) & (iy < NY); okx = (ix >= 0) & (ix < NX)
    A = A[np.ix_(oky, okx)]; iy = iy[oky]; ix = ix[okx]
    IY, IX = np.meshgrid(iy, ix, indexing="ij")
    valid = ~np.isnan(A)
    np.add.at(sum_g, (IY[valid], IX[valid]), A[valid])
    np.add.at(cnt_g, (IY[valid], IX[valid]), 1)
grav = np.where(cnt_g > 0, sum_g / np.maximum(cnt_g, 1), np.nan)
summary["ggmplus"] = {"functional": "gravity disturbance at Earth's surface (dg), mGal",
                      "pixels_per_cell_median": float(np.median(cnt_g[cnt_g > 0])),
                      "nan_fraction_in_region": float(np.isnan(grav).mean())}
del tiles

# ---------- 特征：两种数据源用完全对称的 4 个特征 ----------
def box_mean(a, r):
    v = np.where(np.isnan(a), 0.0, a); m = (~np.isnan(a)).astype(float)
    def csum(x):
        c = np.pad(np.cumsum(np.cumsum(x, 0), 1), ((1, 0), (1, 0)))
        ny, nx = x.shape
        y0 = np.clip(np.arange(ny) - r, 0, ny); y1 = np.clip(np.arange(ny) + r + 1, 0, ny)
        x0 = np.clip(np.arange(nx) - r, 0, nx); x1 = np.clip(np.arange(nx) + r + 1, 0, nx)
        return (c[np.ix_(y1, x1)] - c[np.ix_(y0, x1)] - c[np.ix_(y1, x0)] + c[np.ix_(y0, x0)])
    s = csum(v); n = csum(m)
    return np.where(n > 0, s / np.maximum(n, 1), np.nan), n

def feats(a):
    gy, gx = np.gradient(a)
    grad = np.hypot(gx, gy)                        # 每格（约 3.7 km）的变化量
    mean3, _ = box_mean(a, 1); sq3, _ = box_mean(a * a, 1)
    std3 = np.sqrt(np.clip(sq3 - mean3 ** 2, 0, None))
    mean15, _ = box_mean(a, 7)                     # 15x15 格 ≈ 55 km 窗口的区域场
    resid = a - mean15
    return {"value": a, "grad": grad, "std3": std3, "resid15": resid}

FG = feats(grav); FM = feats(mag)
feat_names = [f"grav_{k}" for k in FG] + [f"mag_{k}" for k in FM]
X = np.stack([FG[k] for k in FG] + [FM[k] for k in FM], axis=-1)   # [ny, nx, 8]

# ---------- 标注：USGS geo3al 概化地质图 → 4 类 ----------
CENOZOIC = {"Q", "N", "Pg", "T", "Cz"}
MESO_PALEO = {"TK", "Mz", "K", "KJ", "J", "JTr", "Tr", "TrP", "MzPz", "Pz", "Pzu", "P", "PC", "C",
              "CD", "D", "DS", "Pzl", "S", "SO", "O", "OCm", "Cm"}
PRECAMB = {"pCm", "Pt", "A"}
CLASSES = ["cenozoic_cover", "mesozoic_paleozoic_strata", "precambrian", "igneous"]
def label_of(typ, glg):
    typ = (typ or "").strip(); glg = (glg or "").strip()
    if typ in ("i", "v", "w"): return 3      # i 侵入岩 / v 喷出岩 / w 超基性岩与蛇绿岩（见 geo3al.xml）
    if typ == "e": return 0                  # e 风成沉积（黄土等），归入新生代覆盖层
    if typ != "": return -1                  # x：元数据未定义，排除
    if glg in CENOZOIC: return 0
    if glg in MESO_PALEO: return 1
    if glg in PRECAMB: return 2
    return -1                                # PzpCm / CzMz / oth / und / H2O：跨界或未定，排除

r = shapefile.Reader(str(RAW / "geo3al"))
wkt = (RAW / "geo3al.prj").read_text()
tr = pyproj.Transformer.from_crs("EPSG:4326", pyproj.CRS.from_wkt(wkt), always_xy=True)
LON, LAT = np.meshgrid(lon_c, lat_c)
PX, PY = tr.transform(LON.ravel(), LAT.ravel())
PX = np.asarray(PX); PY = np.asarray(PY)
lab = np.full(PX.size, -2, dtype=int)        # -2 = 不在任何多边形里
glg_code = np.full(PX.size, "", dtype=object)
for shp, rec in zip(r.iterShapes(), r.iterRecords()):
    bx0, by0, bx1, by1 = shp.bbox
    cand = np.nonzero((PX >= bx0) & (PX <= bx1) & (PY >= by0) & (PY <= by1))[0]
    if cand.size == 0: continue
    pts = np.column_stack([PX[cand], PY[cand]])
    parts = list(shp.parts) + [len(shp.points)]
    inside = np.zeros(cand.size, dtype=int)
    for a, b in zip(parts[:-1], parts[1:]):
        inside += MPath(np.asarray(shp.points[a:b])).contains_points(pts)
    hit = cand[(inside % 2) == 1]           # 奇偶规则处理洞
    lab[hit] = label_of(rec["TYPE"], rec["GEN_GLG"])
    glg_code[hit] = f'{rec["TYPE"]}|{rec["GEN_GLG"]}'
lab = lab.reshape(NY, NX); glg_code = glg_code.reshape(NY, NX)

usable = (lab >= 0) & np.all(np.isfinite(X), axis=-1)
summary["labels"] = {"source": "USGS OFR 97-470F geo3al", "classes": CLASSES,
                     "rule": {"igneous": "TYPE in {i,v,w} (任何时代)", "eolian_e": "归入 cenozoic_cover", "cenozoic_cover": sorted(CENOZOIC),
                              "mesozoic_paleozoic_strata": sorted(MESO_PALEO), "precambrian": sorted(PRECAMB),
                              "excluded": "TYPE x；GEN_GLG 为 PzpCm/CzMz/oth/und/H2O；不在任何多边形内"},
                     "cells_total": int(lab.size),
                     "cells_outside_polygons": int((lab == -2).sum()),
                     "cells_excluded_label": int((lab == -1).sum()),
                     "cells_usable": int(usable.sum()),
                     "class_counts_usable": {c: int(((lab == i) & usable).sum()) for i, c in enumerate(CLASSES)}}
summary["features"] = feat_names
iy, ix = np.nonzero(usable)
np.savez_compressed(OUT, X=X[usable].astype("float32"), y=lab[usable].astype("int8"),
                    iy=iy.astype("int16"), ix=ix.astype("int16"),
                    lon=lon_c[ix].astype("float32"), lat=lat_c[iy].astype("float32"),
                    feat_names=np.array(feat_names), classes=np.array(CLASSES),
                    label_grid=lab.astype("int8"), grav_grid=grav.astype("float32"), mag_grid=mag.astype("float32"))
SUMMARY.write_text(json.dumps(summary, ensure_ascii=False, indent=1))
print(json.dumps(summary, ensure_ascii=False, indent=1))