code/data_preparing/src/l2.py

"""
Level-2 AII preprocessing


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

Purpose
-------
GK2A L2 AII(대기 불안정 지수, NetCDF) 자료를 시간(dt_str, 10분 간격) 단위로 읽어
GK2A 2km(EA020LC) 격자에 맞춘 뒤, 하나의 npy(dict)로 저장하기 위한 전처리 스크립트이다.
satellite_radar 모듈의 산출물과 동일한 관심영역(bbox) crop / 해상도(res) 격자를
사용하므로, 두 스크립트의 npy 산출물은 격자가 서로 일치한다.

Main Features
-------------
1. GK2A L2 AII NetCDF 읽기 및 물리값 복원(_FillValue/scale_factor/add_offset 적용)
2. L2 원본 6km(EA060LC) -> 2km(EA020LC) 최근접 업샘플링(LCC 투영 인덱스 매핑)
3. 선택적으로 2km -> 6km 다운스케일링(3x 집계)
4. 모든 L2 변수(CAPE, KI, LI, SI, TTI)가 존재할 때만 npy 저장

Inputs
------
- metadata.json (CFG): 경로/변수/영역 파라미터 등 실행 설정
- GK2A L2 NetCDF files: CFG["gk2a_l2_base_dir"] 아래 YYYYMM/DD/HH 구조에 존재
- Lat/Lon reference file: CFG["gk2a_ea020_latlon_file"]

Outputs
-------
- npy file: {save_dir}/res_{res}/L2/{YYYYMMDD}/l2_aii_{dt_str}.npy
  (dict 형태: {변수명: 2D array})

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

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

from __future__ import annotations

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

import numpy as np
import pyproj
import xarray as xr
from netCDF4 import Dataset
from tqdm import tqdm


# =============================================================================
# GK2A EA060LC (6km) grid specification
# =============================================================================
# GK2A AMI East Asia Lambert Conformal Conic(LC) 투영 파라미터 및
# 6km(EA060LC) 격자 원점/크기. L2 AII 파일(866 x 1000)이 이 격자에 정의되어 있다.
_LCC_PROJ_PARAMS: Dict[str, Any] = {
    "proj": "lcc",
    "lat_1": 30,
    "lat_2": 60,
    "lat_0": 38,
    "lon_0": 126,
    "ellps": "WGS84",
}
_X0_6KM = -2997000  # 6km 격자 좌상단 x 좌표 (m)
_Y0_6KM = 2595000   # 6km 격자 좌상단 y 좌표 (m)
_RES_6KM = 6000     # 6km 격자 간격 (m)
_NX_6KM = 1000      # 6km 격자 x 크기
_NY_6KM = 866       # 6km 격자 y 크기


# =============================================================================
# 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("l2_preprocess")
    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


# =============================================================================
# Downscaling
# =============================================================================
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 build_index_map(
    latlon_file: str,
    bbox: Dict[str, float],
) -> Tuple[np.ndarray, np.ndarray, Tuple[int, int, int, int]]:
    """
    Build nearest-neighbour index map from 2km(EA020LC) grid to 6km(EA060LC) grid.

    관심영역(bbox)으로 crop한 2km 격자의 각 화소 위경도를 LCC 투영 좌표로 변환한 뒤,
    해당 위치에 대응하는 6km 격자 인덱스(iy, ix)를 계산한다.
    L2 6km 자료를 arr[iy, ix]로 fancy-indexing 하면 2km 격자로 최근접 업샘플링된다.

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

    Returns
    -------
    iy : numpy.ndarray
        6km grid row indices, shape = cropped 2km grid.
    ix : numpy.ndarray
        6km grid column indices, shape = cropped 2km grid.
    crop_idx : tuple of int
        (row_min, row_max, col_min, col_max) indices used for cropping.
    """
    ds = xr.open_dataset(latlon_file)
    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]]

    lon_crop = x[row_min : row_max + 1, col_min : col_max + 1]
    lat_crop = y[row_min : row_max + 1, col_min : col_max + 1]
    ds.close()

    proj = pyproj.Proj(**_LCC_PROJ_PARAMS)
    px, py = proj(lon_crop.astype("f8"), lat_crop.astype("f8"))

    ix = np.round((px - _X0_6KM) / _RES_6KM).astype(int)
    iy = np.round((_Y0_6KM - py) / _RES_6KM).astype(int)
    np.clip(ix, 0, _NX_6KM - 1, out=ix)
    np.clip(iy, 0, _NY_6KM - 1, out=iy)

    return iy, ix, (row_min, row_max, col_min, col_max)


# =============================================================================
# Data readers
# =============================================================================
def decode_var(src: Dataset, vname: str, iy: np.ndarray, ix: np.ndarray) -> np.ndarray:
    """
    Decode one L2 variable and upsample to 2km grid.

    Parameters
    ----------
    src : netCDF4.Dataset
        Opened L2 NetCDF dataset (auto mask/scale disabled).
    vname : str
        Variable name (e.g., "CAPE").
    iy, ix : numpy.ndarray
        6km grid index map from `build_index_map`.

    Returns
    -------
    numpy.ndarray
        Decoded 2D array (float32) on the cropped 2km grid.
        Missing values are np.nan.

    Notes
    -----
    - Missing flag: variable `_FillValue` attribute
    - Scaling: raw * scale_factor + add_offset
    """
    sv = src.variables[vname]
    raw = sv[:][iy, ix]
    fill = int(sv._FillValue)
    scale = float(getattr(sv, "scale_factor", 1.0))
    offset = float(getattr(sv, "add_offset", 0.0))
    valid = raw != fill
    return np.where(valid, raw.astype("f4") * scale + offset, np.nan).astype("f4")


def read_gk2a_l2(
    path: str,
    variables: list,
    iy: np.ndarray,
    ix: np.ndarray,
    logger: Optional[logging.Logger] = None,
) -> Optional[Dict[str, np.ndarray]]:
    """
    Read and decode all L2 variables from one NetCDF file.

    Parameters
    ----------
    path : str
        NetCDF file path.
    variables : list of str
        Variable names to read (e.g., ["CAPE", "KI", "LI", "SI", "TTI"]).
    iy, ix : numpy.ndarray
        6km grid index map from `build_index_map`.
    logger : logging.Logger, optional
        Logger for error reporting.

    Returns
    -------
    dict or None
        {변수명: 2D array} if success, otherwise None.
    """
    try:
        with Dataset(path, "r") as src:
            src.set_auto_maskandscale(False)
            return {v: decode_var(src, v, iy, ix) for v in variables}
    except Exception as e:
        if logger:
            logger.error(f"READ_FAIL_GK2A_L2 | path={path} err={repr(e)}")
        return None


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

    Workflow
    --------
    For each dt_str (10-min step):
    1) Load GK2A L2 AII variables (6km) and upsample to cropped 2km grid
    2) Optionally downscale 2km -> 6km (res="6km")
    3) Save npy only if all L2 variables 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에 없으면 기본값

    if res not in ("2km", "6km"):
        raise ValueError(f"Invalid res: {res} (must be '2km' or '6km')")

    # Output directory
    save_dir = os.path.join(cfg["save_dir"], f"res_{res}", "L2")
    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
    l2_base_dir = cfg["gk2a_l2_base_dir"]
    l2_filename = cfg["l2_filename"]
    variables = cfg["l2_variables"]

    # Index map: cropped 2km grid -> 6km grid (nearest neighbour)
    iy, ix, crop_idx = build_index_map(cfg["gk2a_ea020_latlon_file"], cfg["bbox"])
    r0, r1, c0, c1 = crop_idx
    logger.info(
        f"GRID | 2km crop shape={iy.shape} (rows {r0}:{r1 + 1}, cols {c0}:{c1 + 1}) | "
        f"6km index range: iy {iy.min()}~{iy.max()}, ix {ix.min()}~{ix.max()}"
    )

    for i in tqdm(range(num_days), desc="Processing L2"):
        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}"

                nc_path = os.path.join(
                    l2_base_dir,
                    ymd[:6],
                    ymd[6:8],
                    f"{hour:02d}",
                    l2_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_L2 | dt={dt_str} path={nc_path}")
                    continue

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

                # 1) Read + decode all variables (upsampled to 2km grid)
                data_dict = read_gk2a_l2(nc_path, variables, iy, ix, logger=logger)

                # If reading/decoding fails (None): log and skip this dt
                if data_dict is None:
                    logger.warning(f"SKIP_DT_L2_INCOMPLETE | dt={dt_str}")
                    continue

                # 2) Optional 2km -> 6km downscaling
                if res == "6km":
                    data_dict = {v: downscale_3x(arr, agg="mean") for v, arr in data_dict.items()}

                # 3) Save npy only if ALL required variables exist
                required = set(variables)
                if not required.issubset(data_dict.keys()):
                    missing = sorted(required - set(data_dict.keys()))
                    logger.warning(f"SKIP_SAVE_INCOMPLETE | dt={dt_str} missing={missing}")
                    continue

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

                save_path = os.path.join(day_dir, f"l2_aii_{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()