code/data_preparing/src/satellite_radar.py

"""
Satellite and radar preprocessing


Revision History
-----------------
- [First]       2026-01-21
- [Updated]     2026-07-10

Purpose
-------
GK2A 위성(NetCDF) 채널 자료와 레이더(bin.gz) 강수 자료를 시간(dt_str, 10분 간격) 단위로
동일 격자/해상도로 맞춘 뒤, 하나의 npy(dict)로 저장하기 위한 전처리 스크립트이다.

Main Features
-------------
1. GK2A 위성 채널(NetCDF) 읽기 및 보정(LUT: Excel calibration table)
2. 레이더 자료(bin.gz) 읽기 및 결측 처리
3. 레이더(500m) -> GK2A(2km) 가우시안 리샘플링(pyresample)
4. 선택적으로 2km -> 6km 다운스케일링(3x 집계)
5. GK2A 채널이 모두 존재할 때 npy 저장
   (레이더 cappi/hsr/hsp는 존재하는 제품만 dict에 포함, 결측 제품은 키 생략)

Inputs
------
- metadata.json (CFG): 경로/채널/리샘플 파라미터 등 실행 설정
- GK2A NetCDF files: CFG["gk2a_base_dir"] 아래 구조에 존재
- Radar bin.gz files: CFG["radar_path_cappi"], CFG["radar_path_hsr"],
  CFG["radar_path_hsp"] 아래 존재
- Lat/Lon reference files: CFG["radar_latlon_file"], CFG["gk2a_ea020_latlon_file"]

Outputs
-------
- npy file: {save_dir}/res_{res}/L1B/{YYYYMMDD}/concat_gk2a_radar_{dt_str}.npy
  (dict 형태: {채널명: 2D array, "cappi": 2D array, "hsr": 2D array, "hsp": 2D array}
   단, 결측인 레이더 제품의 키는 생략될 수 있음)

Usage
-----
$ python -m src.preprocess --config CONFIG.yaml

Notes
-----
- 본 스크립트는 대용량 파일 I/O가 포함되므로, 예외 처리와 로그를 통해 누락/오류를 추적한다.
- 리샘플링/다운스케일 과정에서 결측값은 np.nan으로 유지한다.
"""

from __future__ import annotations

import argparse
import gzip
import json
import logging
import os
import sys
from datetime import datetime, timedelta
from functools import lru_cache
from pathlib import Path
from typing import Any, Dict, Optional, Tuple

import numpy as np
import pandas as pd
import pyresample
import xarray as xr
from netCDF4 import Dataset
from tqdm import tqdm


# =============================================================================
# Argument parser
# =============================================================================
def build_parser() -> argparse.ArgumentParser:
    """
    Build CLI argument parser.

    Returns
    -------
    argparse.ArgumentParser
        Parser with arguments:
        - --config : str, path to metadata.json configuration file
    """
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", type=str, required=False, default="../run/metadata.json")
    return parser


def parse_args_auto() -> argparse.Namespace:
    """
    Parse arguments for both interactive(Jupyter) and CLI execution.

    Returns
    -------
    argparse.Namespace
        Parsed arguments.
    """
    parser = build_parser()
    if hasattr(sys, "ps1") or "ipykernel" in sys.modules:
        args, _ = parser.parse_known_args([])
    else:
        args, _ = parser.parse_known_args()
    return args


def load_config(config_path: str) -> Dict[str, Any]:
    """
    Load JSON configuration.

    Parameters
    ----------
    config_path : str
        Path to JSON config file.

    Returns
    -------
    dict
        Configuration dictionary.

    Raises
    ------
    FileNotFoundError
        If config file does not exist.
    json.JSONDecodeError
        If config file is not a valid JSON.
    """
    with open(config_path, "r") as f:
        return json.load(f)


# =============================================================================
# Logger
# =============================================================================
def setup_logger(log_path: str) -> logging.Logger:
    """
    Set up file + stdout logger.

    Parameters
    ----------
    log_path : str
        Log file path.

    Returns
    -------
    logging.Logger
        Configured logger instance.
    """
    logger = logging.getLogger("concat")
    logger.setLevel(logging.INFO)
    logger.handlers.clear()

    fmt = logging.Formatter("%(asctime)s | %(levelname)s | %(message)s")

    fh = logging.FileHandler(log_path)
    fh.setFormatter(fmt)
    logger.addHandler(fh)

    sh = logging.StreamHandler(sys.stdout)
    sh.setFormatter(fmt)
    logger.addHandler(sh)

    return logger


# =============================================================================
# Calibration table
# =============================================================================
_CHANNEL_PREFIX: Dict[str, str] = {
    "wv063": "IR 6.3",
    "wv073": "IR 7.3",
    "ir087": "IR 8.7",
    "ir105": "IR 10.5",
    "ir112": "IR 11.2",
    "ir123": "IR 12.3",
    "ir133": "IR 13.3",
}

_VARIABLE_KEYWORD: Dict[str, str] = {
    "Radiance": "Radiance",
    "Albedo": "Albedo",
    "BT": "Brightness Temperature",
}


@lru_cache(maxsize=1)
def _load_cal_table(calib_path: str) -> pd.DataFrame:
    """
    Load calibration LUT table from Excel once (cached).

    Parameters
    ----------
    cfg : dict
        Configuration dictionary which must include:
        - "Calibration_table_path": Excel file path

    Returns
    -------
    pandas.DataFrame
        Parsed calibration table.

    Raises
    ------
    FileNotFoundError
        If the Excel file does not exist.
    ValueError
        If required sheet/format is not found.
    """
    df = pd.read_excel(
        calib_path,
        sheet_name="Calibration Table",
        header=[0, 1],
        index_col=0,
        engine="openpyxl",
    )
    df.columns = [f"{l0} {l1}".strip() for l0, l1 in df.columns]
    return df


def calibrate_from_excel(
    dn: np.ndarray,
    channel: str,
    var: str,
    cfg: Dict[str, Any],
) -> np.ndarray:
    """
    Calibrate DN values using LUT from Excel.

    Parameters
    ----------
    dn : numpy.ndarray
        Digital number array.
    channel : str
        GK2A channel key (e.g., "ir105", "wv063").
    var : str
        Variable type. One of {"Radiance", "Albedo", "BT"}.
    cfg : dict
        Configuration dictionary.

    Returns
    -------
    numpy.ndarray
        Calibrated array (float32), same shape as `dn`.

    Raises
    ------
    ValueError
        If channel/variable mapping is not found in calibration table.
    """
    df = _load_cal_table(cfg["Calibration_table_path"])
    pref = _CHANNEL_PREFIX[channel]
    kw = _VARIABLE_KEYWORD[var]

    try:
        col = next(c for c in df.columns if c.startswith(pref) and kw in c)
    except StopIteration as e:
        raise ValueError(
            f"Channel {channel} not found in calibration table. "
            f"(prefix: {pref}, keyword: {kw})"
        ) from e

    max_dn = int(df.index.max())
    lut = np.full(max_dn + 1, np.nan, dtype=np.float32)
    lut[df.index.astype(int)] = df[col].values

    idx = np.clip(np.rint(dn).astype(int), 0, max_dn)
    return lut[idx]


# =============================================================================
# Resampling / Downscaling
# =============================================================================
def resampling(
    orig_grid: pyresample.geometry.GridDefinition,
    target: np.ndarray,
    targ_grid: pyresample.geometry.GridDefinition,
    resample_cfg: Dict[str, Any],
) -> np.ndarray:
    """
    Resample 2D field from `orig_grid` to `targ_grid` using Gaussian weighting.

    Parameters
    ----------
    orig_grid : pyresample.geometry.GridDefinition
        Source grid definition (lon/lat).
    target : numpy.ndarray
        Source data array (2D).
    targ_grid : pyresample.geometry.GridDefinition
        Target grid definition (lon/lat).
    resample_cfg : dict
        Resampling parameters with keys:
        - radius_of_influence
        - neighbours
        - sigmas

    Returns
    -------
    numpy.ndarray
        Resampled 2D array with shape of target grid.
    """
    return pyresample.kd_tree.resample_gauss(
        orig_grid,
        target,
        targ_grid,
        radius_of_influence=resample_cfg["radius_of_influence"],
        neighbours=resample_cfg["neighbours"],
        sigmas=resample_cfg["sigmas"],
        fill_value=np.nan,
    )


def downscale_3x(arr2d: np.ndarray, agg: str = "mean", f: int = 3) -> np.ndarray:
    """
    Downscale 2D array by integer factor `f` using block aggregation.

    Parameters
    ----------
    arr2d : numpy.ndarray
        2D array (ny, nx).
    agg : str, default="mean"
        Aggregation method. One of {"mean", "max", "min", "median"}.
    f : int, default=3
        Downscale factor (e.g., f=3 for 2km->6km).

    Returns
    -------
    numpy.ndarray
        Downscaled 2D array with shape (ny//f, nx//f).

    Raises
    ------
    ValueError
        If `agg` is not supported.
    """
    ny, nx = arr2d.shape
    ny2 = (ny // f) * f
    nx2 = (nx // f) * f

    a = arr2d[:ny2, :nx2]
    a = a.reshape(ny2 // f, f, nx2 // f, f)

    if agg == "mean":
        s = np.nansum(a, axis=(1, 3))
        c = np.sum(~np.isnan(a), axis=(1, 3))
        out = s / np.where(c == 0, 1, c)
        out[c == 0] = np.nan
        return out.astype(np.float32, copy=False)

    if agg == "max":
        out = np.nanmax(np.where(np.isnan(a), -np.inf, a), axis=(1, 3))
        out[np.isneginf(out)] = np.nan
        return out.astype(np.float32, copy=False)

    if agg == "min":
        out = np.nanmin(np.where(np.isnan(a), np.inf, a), axis=(1, 3))
        out[np.isposinf(out)] = np.nan
        return out.astype(np.float32, copy=False)

    if agg == "median":
        out = np.nanmedian(a, axis=(1, 3))
        return out.astype(np.float32, copy=False)

    raise ValueError("agg must be 'mean'|'max'|'min'|'median'")


# =============================================================================
# Grid utilities
# =============================================================================
def get_orig_grid(file_name: str) -> pyresample.geometry.GridDefinition:
    """
    Build source grid definition from a NetCDF lat/lon reference file.

    Parameters
    ----------
    file_name : str
        Path to NetCDF file containing `lon` and `lat` variables.

    Returns
    -------
    pyresample.geometry.GridDefinition
        Grid definition created from full lon/lat arrays.
    """
    ds = xr.open_dataset(file_name)
    x = ds["lon"][:].data
    y = ds["lat"][:].data
    orig_grid = pyresample.geometry.GridDefinition(lons=x, lats=y)
    ds.close()
    return orig_grid


def get_targ_grid(
    file_name: str,
    bbox: Dict[str, float],
) -> Tuple[pyresample.geometry.GridDefinition, Tuple[int, int, int, int]]:
    """
    Build target grid definition by cropping lon/lat based on bounding box.

    Parameters
    ----------
    file_name : str
        Path to NetCDF file containing `lon` and `lat` variables.
    bbox : dict
        Bounding box with keys:
        - lon_min, lon_max, lat_min, lat_max

    Returns
    -------
    targ_grid : pyresample.geometry.GridDefinition
        Cropped target grid definition.
    crop_idx : tuple of int
        (row_min, row_max, col_min, col_max) indices used for cropping.
    """
    ds = xr.open_dataset(file_name)
    x = ds["lon"][:].data
    y = ds["lat"][:].data

    lon_min = bbox["lon_min"]
    lon_max = bbox["lon_max"]
    lat_min = bbox["lat_min"]
    lat_max = bbox["lat_max"]

    mask = (x >= lon_min) & (x <= lon_max) & (y >= lat_min) & (y <= lat_max)
    rows = np.any(mask, axis=1)
    cols = np.any(mask, axis=0)

    row_min, row_max = np.where(rows)[0][[0, -1]]
    col_min, col_max = np.where(cols)[0][[0, -1]]

    x_crop = x[row_min : row_max + 1, col_min : col_max + 1]
    y_crop = y[row_min : row_max + 1, col_min : col_max + 1]

    targ_grid = pyresample.geometry.GridDefinition(lons=x_crop, lats=y_crop)
    ds.close()

    return targ_grid, (row_min, row_max, col_min, col_max)


# =============================================================================
# Data readers
# =============================================================================
def radar_filename_time_from_gk2a_utc(dt_str_utc: str):
    """
    Convert GK2A UTC datetime string to the corresponding radar filename timestamp in KST.

    Parameters
    ----------
    dt_str_utc : str
        Datetime string in UTC, format 'YYYYMMDDHHMM' (e.g., '202201010500' for 05:00 UTC).

    Returns
    -------
    radar_ymd_kst : str
        Date string in KST, format 'YYYYMMDD'.
    radar_dt_str_kst : str
        Datetime string in KST, format 'YYYYMMDDHHMM'.

    Notes
    -----
    This function shifts the given UTC time by +9 hours to KST
    to match the radar file naming convention.
    Example: '202201010500' (05:00 UTC) -> '202201011400' (14:00 KST)
    """
    dt_utc = datetime.strptime(dt_str_utc, "%Y%m%d%H%M")
    dt_kst = dt_utc + timedelta(hours=9)
    radar_ymd_kst = dt_kst.strftime("%Y%m%d")
    radar_dt_str_kst = dt_kst.strftime("%Y%m%d%H%M")
    return radar_ymd_kst, radar_dt_str_kst


def read_gk2a(
    path: str,
    ch: str,
    cfg: Dict[str, Any],
    logger: Optional[logging.Logger] = None,
) -> Optional[np.ndarray]:
    """
    Read and calibrate GK2A channel data from NetCDF.

    Parameters
    ----------
    path : str
        NetCDF file path.
    ch : str
        Channel key.
    cfg : dict
        Configuration dictionary.
    logger : logging.Logger, optional
        Logger for error reporting.

    Returns
    -------
    numpy.ndarray or None
        Calibrated 2D array if success, otherwise None.

    Notes
    -----
    - NetCDF variable: 'image_pixel_values'
    - Calibration rule:
      * vi* : Radiance
      * ir*/wv* : Brightness Temperature (BT)
      * else : raw data
    """
    try:
        nc = Dataset(path, "r")
        data = nc["image_pixel_values"][:].data
        nc.close()
    except Exception as e:
        if logger:
            logger.error(f"READ_FAIL_GK2A | ch={ch} path={path} err={repr(e)}")
        return None

    try:
        if ch.startswith("vi"):
            return calibrate_from_excel(data, ch, "Radiance", cfg)
        if ch.startswith(("ir", "wv")):
            return calibrate_from_excel(data, ch, "BT", cfg)
        return data
    except Exception as e:
        if logger:
            logger.error(f"CAL_FAIL_GK2A | ch={ch} path={path} err={repr(e)}")
        return None


def read_radar(
    file_name: str,
    nx: int = 2305,
    ny: int = 2881,
    header_bytes: int = 1024,
) -> np.ndarray:
    """
    Read radar bin.gz file and convert to rain rate.

    Parameters
    ----------
    file_name : str
        Path to compressed radar binary file (*.bin.gz).
    nx : int, default=2305
        X dimension.
    ny : int, default=2881
        Y dimension.
    header_bytes : int, default=1024
        Header size to skip before the int16 data block.
        - CMP products (CAPPI/HSR, RDR_CMP_*): 1024
        - HSP product (RDR_HSP_EXT_*): 4

    Returns
    -------
    numpy.ndarray
        2D rain rate array (float32) with shape (ny, nx).
        Missing values are np.nan.

    Notes
    -----
    - Data starts after `header_bytes` header.
    - Raw dtype: int16
    - Missing flag: <= -30000
    - Scaling: / 100
    """
    with open(file_name, "rb") as f:
        decompressed_bytes = gzip.decompress(f.read())

    rain_rate = (
        np.frombuffer(decompressed_bytes, dtype=np.int16, offset=header_bytes)
        .astype(np.float32)
        .reshape(ny, nx)
    )

    null_mask = rain_rate <= -30000
    rain_rate[null_mask] = np.nan
    rain_rate /= 100.0

    return rain_rate


# =============================================================================
# Main
# =============================================================================
def main() -> None:
    """
    Run preprocessing pipeline for given date range.

    Workflow
    --------
    For each dt_str (10-min step):
    1) Load GK2A channels (cropped to bbox, res=2km/6km)
    2) Load radar products (CAPPI / HSR / HSP) and resample to GK2A grid
    3) Save npy only if all GK2A channels + all radar fields exist
    """
    args = parse_args_auto()

    # Script working directory: script location (for relative config path)
    script_dir = Path(__file__).resolve().parent
    os.chdir(script_dir)

    # ✅ config 경로도 CLI에서 바꿀 수 있게
    cfg = load_config(args.config)

    # ✅ 날짜/해상도는 json에서 읽음
    start_date = cfg["start_date"]  # e.g., "20210701"
    end_date = cfg["end_date"]      # e.g., "20210703"
    res = cfg.get("res", "2km")     # json에 없으면 기본값

    # Output directory
    save_dir = os.path.join(cfg["save_dir"], f"res_{res}", "L1B")
    os.makedirs(save_dir, exist_ok=True)

    # Date settings
    start_dt = datetime.strptime(start_date, "%Y%m%d")
    end_dt = datetime.strptime(end_date, "%Y%m%d")
    num_days = (end_dt - start_dt).days + 1

    # Logger
    log_dir = os.path.join(save_dir, "_logs")
    os.makedirs(log_dir, exist_ok=True)
    log_path = os.path.join(log_dir, f"log_{start_date}_{end_date}.log")
    logger = setup_logger(log_path)

    logger.info(f"START | {start_date} ~ {end_date}")
    logger.info(f"save_dir={save_dir}")

    # Paths / configs
    gk2a_base_dir = cfg["gk2a_base_dir"]
    channels = cfg["channels"]

    # Radar products: (key, base_path, filename_prefix, header_bytes)
    radar_products = [
        ("cappi", cfg["radar_path_cappi"], cfg["radar_prefix_cappi"], 1024),
        ("hsr", cfg["radar_path_hsr"], cfg["radar_prefix_hsr"], 1024),
        ("hsp", cfg["radar_path_hsp"], cfg["radar_prefix_hsp"], 4),
    ]

    # Grids
    orig_grid_radar = get_orig_grid(cfg["radar_latlon_file"])
    orig_grid_gk2a = get_orig_grid(cfg["gk2a_ea020_latlon_file"])  # reserved
    targ_grid_gk2a, crop_idx = get_targ_grid(cfg["gk2a_ea020_latlon_file"], cfg["bbox"])
    r0, r1, c0, c1 = crop_idx

    _ = orig_grid_gk2a  # avoid unused warnings in some linters

    for i in tqdm(range(num_days), desc="Processing concat"):
        current_dt = start_dt + timedelta(days=i)
        ymd = current_dt.strftime("%Y%m%d")
        print(f"Processing date: {ymd}")

        for hour in range(0, 24):
            for minute in range(0, 60, 10):
                dt_str = f"{ymd}{hour:02d}{minute:02d}"
                data_dict: Dict[str, np.ndarray] = {}

                # If any GK2A channel fails for this dt_str, skip the entire timestamp
                skip_dt = False

                # 1) GK2A channels
                for ch, info in channels.items():
                    nc_path = os.path.join(
                        gk2a_base_dir,
                        ymd[:6],
                        ymd[6:8],
                        f"{hour:02d}",
                        info["filename"].format(dt=dt_str),
                    )

                    # If file does not exist: log and skip this dt
                    if not os.path.exists(nc_path):
                        logger.warning(f"MISS_GK2A | dt={dt_str} ch={ch} path={nc_path}")
                        skip_dt = True
                        break

                    # (Optional) If file is too small, treat as corrupted and skip this dt
                    min_size = cfg.get("min_gk2a_nc_size_bytes", 0)
                    if min_size and os.path.getsize(nc_path) < min_size:
                        logger.warning(
                            f"CORRUPT_GK2A_SMALLFILE | dt={dt_str} ch={ch} "
                            f"size={os.path.getsize(nc_path)} path={nc_path}"
                        )
                        skip_dt = True
                        break

                    gk2a_data = read_gk2a(nc_path, ch, cfg, logger=logger)

                    # If reading/calibration fails (None): log and skip this dt
                    if gk2a_data is None:
                        logger.warning(f"READ_FAIL_GK2A | dt={dt_str} ch={ch} path={nc_path}")
                        skip_dt = True
                        break

                    gk2a_2km = gk2a_data[r0 : r1 + 1, c0 : c1 + 1]
                    if res == "2km":
                        gk2a_resampled = gk2a_2km
                    elif res == "6km":
                        gk2a_resampled = downscale_3x(gk2a_2km, agg="mean")
                    else:
                        logger.error(f"INVALID_RES | res={res}")
                        skip_dt = True
                        break

                    data_dict[ch] = gk2a_resampled

                # If any GK2A problem occurs, do not read radar and skip to next dt
                if skip_dt:
                    logger.warning(f"SKIP_DT_GK2A_INCOMPLETE | dt={dt_str}")
                    continue

                # Convert GK2A UTC datetime string to the corresponding radar filename timestamp in KST
                radar_ymd_kst, radar_dt_str_kst = radar_filename_time_from_gk2a_utc(dt_str)

                logger.info(f"TIME_MATCH | gk2a_utc={dt_str} <-> radar_fname_kst={radar_dt_str_kst}")

                # 2) Radar products (CAPPI / HSR / HSP)
                for rkey, rpath, rprefix, rheader in radar_products:
                    radar_file = os.path.join(
                        rpath,
                        radar_ymd_kst,
                        f"{rprefix}{radar_dt_str_kst}.bin.gz",
                    )
                    if not os.path.exists(radar_file):
                        logger.warning(
                            f"MISS_RADAR_{rkey.upper()} | dt={dt_str} path={radar_file}"
                        )
                        continue

                    try:
                        radar_data = read_radar(radar_file, header_bytes=rheader)
                    except Exception as e:
                        logger.error(
                            f"READ_FAIL_RADAR_{rkey.upper()} | "
                            f"dt={dt_str} path={radar_file} err={repr(e)}"
                        )
                        continue

                    radar_2km = resampling(
                        orig_grid_radar,
                        radar_data,
                        targ_grid_gk2a,
                        cfg["resample"],
                    )
                    if res == "2km":
                        data_dict[rkey] = radar_2km
                    else:
                        data_dict[rkey] = downscale_3x(radar_2km, agg="mean")

                # 3) Save npy
                #    - GK2A 채널: 모두 필수 (하나라도 없으면 위에서 이미 skip됨)
                #    - 레이더: 존재하는 제품만 dict에 포함하여 저장 (결측 제품은 키 생략)
                required_gk2a = set(channels.keys())
                if not required_gk2a.issubset(data_dict.keys()):
                    missing_gk2a = sorted(required_gk2a - set(data_dict.keys()))
                    logger.warning(
                        f"SKIP_SAVE_INCOMPLETE | dt={dt_str} missing_gk2a={missing_gk2a}"
                    )
                    continue

                missing_radar = sorted(
                    {rkey for rkey, _, _, _ in radar_products} - set(data_dict.keys())
                )
                if missing_radar:
                    logger.warning(
                        f"SAVE_PARTIAL_RADAR | dt={dt_str} missing_radar={missing_radar}"
                    )

                day_dir = os.path.join(save_dir, ymd)
                os.makedirs(day_dir, exist_ok=True)

                save_path = os.path.join(day_dir, f"concat_gk2a_radar_{dt_str}.npy")
                if os.path.exists(save_path):
                    print(f"[ Skip ]: {save_path}")
                    continue

                np.save(save_path, data_dict)
                print(f"[Saved]: {save_path}")

    print(" Done!")


if __name__ == "__main__":
    main()