#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Mature Cloud Removal (temporal descendant masking)

개선 개요
- step2의 label/link 결과를 따라가며 초기 구름 또는 성숙 구름으로 판정된 클러스터와 미래 클러스터를 제거한다.
- 제거 조건은 BT, radar 값, 과거 연결 여부, 클러스터 크기를 함께 사용한다.
- 최신 버전은 위경도/낙뢰 보조파일을 사용하지 않고, BT 입력의 ir105 및 RADAR_VAR(hsr)만 사용한다.

마스킹 조건
1. BT <= MATURE_THRESHOLD K이면 해당 시점 이후 클러스터 제거
2. RADAR_VAR >= RADAR_THRESHOLD dBZ이면 해당 시점 이후 클러스터 제거
3. radar 영역 밖이고 과거 연결이 없으면 전체 생애주기 제거
4. MIN_SIZE_THRESHOLD 픽셀 이상이고 과거 연결이 없으면 전체 생애주기 제거

* 입력
    - step2 label: YYYYMMDDHHMM_label.nc
    - step2 links: YYYYMMDDHHMM_links.pkl
    - BT/RADAR npy: concat_gk2a_radar_YYYYMMDDHHMM.npy (ir105, hsr/cappi)
* 결과물
    - YYYYMMDDHHMM_label.nc (제거 대상 클러스터가 NaN 처리된 최종 label)

경로 설정
- CONCAT_BT_BASE/LABEL_BASE/LINKS_BASE/SAVE_DIR는 import 시점에는 빈 문자열이다.
- main에서 build/run/step3_mature_cloud_masking_config.json을 읽은 뒤 apply_config()가 실제 경로로 채운다.
- bt_root가 data_preprocess/result/res_2km처럼 L1B/L2 폴더이면, 실제 labeling 입력인 L1B/YYYYMMDD를 자동 선택한다.
"""
import os
import sys
import argparse
import numpy as np
import xarray as xr
from datetime import datetime, timedelta, timezone
from collections import deque
from tqdm import tqdm
import pickle
import pandas as pd

# ──────────────────────────────────────────────────────────────
# 경로/파라미터
# ──────────────────────────────────────────────────────────────
try:
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")
    sys.stderr.reconfigure(encoding="utf-8", errors="replace")
except AttributeError:
    pass


CONFIG_NAME = "step3_mature_cloud_masking_config.json"


def _find_package_root(config_name: str) -> str:
    cur = os.path.dirname(os.path.abspath(__file__))
    while True:
        if os.path.exists(os.path.join(cur, "build", "run", config_name)):
            return cur
        parent = os.path.dirname(cur)
        if parent == cur:
            return os.path.abspath(os.getcwd())
        cur = parent


ROOT = _find_package_root(CONFIG_NAME)
RUN_DIR = os.path.join(ROOT, "build", "run")
sys.path.insert(0, ROOT)

from .config_utils import load_config


def _resolve_path(path: str) -> str:
    if os.path.isabs(path):
        return path
    return os.path.abspath(os.path.join(RUN_DIR, path))


def _has_date_dirs(path: str) -> bool:
    if not os.path.isdir(path):
        return False
    return any(
        os.path.isdir(os.path.join(path, name)) and len(name) == 8 and name.isdigit()
        for name in os.listdir(path)
    )


def _resolve_data_root(path: str) -> str:
    if _has_date_dirs(path):
        return path
    for subdir in ("L1B", "l1b"):
        candidate = os.path.join(path, subdir)
        if _has_date_dirs(candidate):
            return candidate
    return path


# Config 적용 전 placeholder. 실제 값은 apply_config()에서 채운다.
CONCAT_BT_BASE = ""
LABEL_BASE = ""
LINKS_BASE = ""
SAVE_DIR = ""

MATURE_THRESHOLD = 233.0  # IR10.5 (K)
RADAR_VAR        = "hsr"  # 사용할 레이더 변수: "hsr" 또는 "cappi"
RADAR_THRESHOLD  = 35.0   # dBZ

TIME_STEP         = 10    # 분
MAX_REMOVE_HOURS  = 6     # 시간

# 초기 구름 필터링 설정
MIN_SIZE_THRESHOLD = 156   # 픽셀 수 임계값

START_TIME = "202508010000"  # YYYYMMDDHHMM
END_TIME   = "202510312350"  # YYYYMMDDHHMM


def apply_config(cfg: dict) -> None:
    global CONCAT_BT_BASE, LABEL_BASE, LINKS_BASE, SAVE_DIR
    global MATURE_THRESHOLD, RADAR_VAR, RADAR_THRESHOLD, TIME_STEP
    global MAX_REMOVE_HOURS, MIN_SIZE_THRESHOLD, START_TIME, END_TIME

    CONCAT_BT_BASE = _resolve_data_root(_resolve_path(cfg["bt_root"]))
    LABEL_BASE = _resolve_path(cfg["label_root"])
    LINKS_BASE = _resolve_path(cfg.get("links_root", cfg["label_root"]))
    SAVE_DIR = _resolve_path(cfg["output_dir"])
    os.makedirs(SAVE_DIR, exist_ok=True)

    MATURE_THRESHOLD = float(cfg.get("mature_threshold", MATURE_THRESHOLD))
    RADAR_VAR = cfg.get("radar_var", RADAR_VAR)
    RADAR_THRESHOLD = float(cfg.get("radar_threshold", RADAR_THRESHOLD))
    TIME_STEP = int(cfg.get("time_step_min", TIME_STEP))
    MAX_REMOVE_HOURS = int(cfg.get("max_remove_hours", MAX_REMOVE_HOURS))
    MIN_SIZE_THRESHOLD = int(cfg.get("min_size_threshold", MIN_SIZE_THRESHOLD))
    START_TIME = cfg.get("start_time", START_TIME)
    END_TIME = cfg.get("end_time", END_TIME)

# ──────────────────────────────────────────────────────────────
# 유틸/캐시
# ──────────────────────────────────────────────────────────────
DATA_CACHE  = {}  # dt -> (label, ir105, radar)
LINKS_CACHE = {}  # dt -> { "YYYY..._id": ["YYYY.. ._id", ...] }

def generate_timesteps(start_time, end_time, step_minutes=TIME_STEP):
    s = datetime.strptime(start_time, "%Y%m%d%H%M")
    e = datetime.strptime(end_time, "%Y%m%d%H%M")
    out, cur = [], s
    while cur <= e:
        out.append(cur.strftime("%Y%m%d%H%M"))
        cur += timedelta(minutes=step_minutes)
    return out

def safe_unique_labels(label_arr):
    """라벨에서 0/NaN 제거하고 정수 레이블 배열 반환"""
    vals = np.unique(label_arr)
    if vals.dtype.kind == "f":
        vals = vals[np.isfinite(vals)]
    vals = vals[vals != 0]
    return vals.astype(np.int64)

def load_data(dt_str):
    day = dt_str[:8]
    lp  = os.path.join(LABEL_BASE, day, f"{dt_str}_label.nc")
    bp  = os.path.join(CONCAT_BT_BASE, day, f"concat_gk2a_radar_{dt_str}.npy")
    
    label, ir105, radar = None, None, None
    
    # 1. 라벨 데이터 로드 (레이더 유무와 관계없이 로드 시도)
    if os.path.exists(lp):
        try:
            with xr.open_dataset(lp) as lds:
                label = lds["label"].values.copy()
        except Exception:
            pass

    # 2. BT/레이더(.npy) 데이터 로드 (없거나 키가 누락되면 None 유지)
    if os.path.exists(bp):
        try:
            data = np.load(bp, allow_pickle=True).item()
            if isinstance(data, dict):
                ir105 = data.get("ir105", None)
                radar = data.get(RADAR_VAR, None)
        except Exception:
            pass

    return label, ir105, radar

def get_data(dt):
    if dt not in DATA_CACHE:
        DATA_CACHE[dt] = load_data(dt)
    return DATA_CACHE[dt]

def load_links(dt_str):
    """엄격 형식: {"YYYYMMDDHHMM_id": ["YYYYMMDDHHMM_id", ...]}"""
    day = dt_str[:8]
    p   = os.path.join(LINKS_BASE, day, f"{dt_str}_links.pkl")
    if not os.path.exists(p):
        return {}
    try:
        with open(p, "rb") as f:
            d = pickle.load(f)
        return d if isinstance(d, dict) else {}
    except Exception:
        return {}

def get_links(dt):
    if dt not in LINKS_CACHE:
        LINKS_CACHE[dt] = load_links(dt)
    return LINKS_CACHE[dt]

def parse_cid_str(s):
    """'YYYYMMDDHHMM_id' → (dt:str, id:int). 형식 불일치 시 (None,None)"""
    if not isinstance(s, str) or "_" not in s:
        return (None, None)
    dt, nid = s.split("_", 1)
    try:
        return (dt, int(nid))
    except ValueError:
        return (None, None)

def make_cid(dt, cid_int):
    return f"{dt}_{int(cid_int)}"

# ──────────────────────────────────────────────────────────────
# 초기 구름 조건 체크
# ──────────────────────────────────────────────────────────────
def is_outside_radar_region(dt, cid_int):
    """구름의 모든 픽셀이 레이더 영역 밖에 있는지 확인"""
    label, _, radar_data = get_data(dt)
    if label is None or radar_data is None:
        return False
    
    mask = (label == cid_int)
    if not np.any(mask):
        return False
    
    # 클러스터 영역의 레이더 값 추출
    radar_values = radar_data[mask]
    valid_radar = radar_values[np.isfinite(radar_values)]
    
    # 유효한 레이더 값이 하나도 없으면 레이더 영역 밖
    return len(valid_radar) == 0

def has_no_past_connection(dt, cid_int):
    """해당 구름이 과거와 연결이 없는지 확인 (최초 생성 구름)"""
    links = get_links(dt)
    cid_key = make_cid(dt, cid_int)
    
    if cid_key not in links:
        return True
    
    past_connections = links[cid_key]
    return not past_connections or len(past_connections) == 0

def get_cluster_size(dt, cid_int):
    """클러스터의 픽셀 수 반환"""
    label, _, _ = get_data(dt)
    if label is None:
        return 0
    
    mask = (label == cid_int)
    return np.sum(mask)

def is_initial_cloud_to_remove(dt, cid_int):
    """제거 대상 초기 구름인지 확인"""
    if not has_no_past_connection(dt, cid_int):
        return False
    
    cluster_size = get_cluster_size(dt, cid_int)
    
    if is_outside_radar_region(dt, cid_int):
        return True
    if cluster_size >= MIN_SIZE_THRESHOLD:
        return True
    
    return False

# ──────────────────────────────────────────────────────────────
# 성숙 조건 체크
# ──────────────────────────────────────────────────────────────
def is_mature(dt, cid_int):
    """성숙 구름 조건 체크 (단순화)"""
    label, ir, radar_data = get_data(dt)
    if label is None or ir is None or radar_data is None:
        return False
        
    m = (label == cid_int)
    if not np.any(m):
        return False
    
    # 조건 1: BT 조건 체크
    v_ir = ir[m]
    v_ir = v_ir[np.isfinite(v_ir)]
    if len(v_ir) > 0 and np.nanmin(v_ir) <= MATURE_THRESHOLD:
        return True
    
    # 조건 2: 레이더 임계값 체크
    v_radar = radar_data[m]
    v_radar = v_radar[np.isfinite(v_radar)]
    if len(v_radar) > 0 and np.nanmax(v_radar) >= RADAR_THRESHOLD:
        return True
    
    return False

# ──────────────────────────────────────────────────────────────
# 후속(미래) 구름 찾기 / 수집
# ──────────────────────────────────────────────────────────────
def successors(prev_dt, prev_id, next_dt):
    links_next = get_links(next_dt)
    if not links_next:
        return []
    out = []
    for now_key, prev_list in links_next.items():
        now_dt, now_id = parse_cid_str(now_key)
        if now_dt is None or now_id is None:
            continue
        for pv in (prev_list or []):
            pv_dt, pv_id = parse_cid_str(pv)
            if pv_dt == prev_dt and pv_id == prev_id:
                out.append(now_id)
                break
    return out

def check_mature_cloud(dt, cid_int):
    return is_mature(dt, cid_int)

def collect_descendants_from(dt, cid_int, timesteps, max_minutes):
    out = set([make_cid(dt, cid_int)])
    if dt not in timesteps:
        return out
    start_idx = timesteps.index(dt)
    last_idx  = len(timesteps)-1

    q = deque()
    q.append((start_idx, cid_int))
    visited = set([(start_idx, cid_int)])

    while q:
        idx, cur_id = q.popleft()
        cur_dt = timesteps[idx]
        out.add(make_cid(cur_dt, cur_id))
        if idx == last_idx:
            continue
        next_dt = timesteps[idx+1]
        for nx_id in successors(cur_dt, cur_id, next_dt):
            st = (idx+1, nx_id)
            if st in visited:
                continue
            visited.add(st)
            q.append(st)
    return out

# ──────────────────────────────────────────────────────────────
# 메인
# ──────────────────────────────────────────────────────────────
def process_mature_cloud_masking(start_time=START_TIME, end_time=END_TIME):
    print(f"🚀 구름 마스킹 (단순화된 조건): {start_time}~{end_time}")

    timesteps = generate_timesteps(start_time, end_time, TIME_STEP)
    print(f"\n🔍 타임스텝: {len(timesteps)}개, 제거 범위: {MAX_REMOVE_HOURS}h")
    print(f"📋 마스킹 조건:")
    print(f"   1. BT ≤ {MATURE_THRESHOLD}K (최초 시점부터 미래 전체)")
    print(f"   2. {RADAR_VAR.upper()} ≥ {RADAR_THRESHOLD}dBZ (최초 시점부터 미래 전체)")
    print(f"   3. {RADAR_VAR.upper()} 영역 밖 + 과거 연결 없음 (전체 생애주기)")
    print(f"   4. {MIN_SIZE_THRESHOLD}픽셀 이상 + 과거 연결 없음 (전체 생애주기)")

    remove_cids = set()
    initial_remove_count = 0
    mature_remove_count = 0

    # 1) 통합 구름 필터링 (초기 + 성숙)
    for dt in tqdm(timesteps, desc="구름 필터링"):
        label, ir, radar_data = get_data(dt)
        
        # 라벨이 아예 없으면 건너뜀
        if label is None:
            continue
            
        # 데이터(IR 또는 HSR)가 없으면 해당 타임스텝의 "새로운 성숙 구름 판별"은 패스
        # 단, 다른 타임스텝에서 수집된 remove_cids 추적에는 영향 없음
        if ir is None or radar_data is None:
            continue

        for cid in safe_unique_labels(label):
            start_key = make_cid(dt, cid)
            if start_key in remove_cids:
                continue
                
            # 1) 초기 구름 체크 (우선순위)
            if is_initial_cloud_to_remove(dt, cid):
                descendants = collect_descendants_from(dt, cid, timesteps, MAX_REMOVE_HOURS*60)
                remove_cids |= descendants
                initial_remove_count += len(descendants)
                
                cluster_size = get_cluster_size(dt, cid)
                is_outside_radar = is_outside_radar_region(dt, cid)
                reason = f"{RADAR_VAR.upper()}영역밖" if is_outside_radar else f"대형구름({cluster_size}픽셀)"
                # print(f"   📝 초기 구름 제거: {dt}_{cid} ({reason}) → {len(descendants)}개 후손")
                continue
                
            # 2) 성숙 구름 체크
            if check_mature_cloud(dt, cid):
                mature_descendants = collect_descendants_from(dt, cid, timesteps, MAX_REMOVE_HOURS*60)
                remove_cids |= mature_descendants
                mature_remove_count += len(mature_descendants)
                
                # print(f"   ⚡ 성숙 구름 제거: {dt}_{cid} (성숙조건) → {len(mature_descendants)}개 미래 클러스터")

    print(f"\n📊 필터링 결과:")
    print(f"   - 초기 구름 제거: {initial_remove_count}개")
    print(f"   - 성숙 구름 제거: {mature_remove_count}개")
    print(f"   - 총 제거 대상: {len(remove_cids)}개")

    # 2) 라벨 파일 업데이트/저장
    for dt in tqdm(timesteps, desc="라벨 업데이트"):
        label, _, _ = get_data(dt)
        
        # HSR 데이터 누락 여부와 상관없이 라벨 파일이 존재하면 마스킹 진행
        if label is None:
            continue

        # 이 시각의 제거 대상 id들 추출
        to_remove = []
        for key in remove_cids:
            kdt, kid = parse_cid_str(key)
            if kdt == dt and kid is not None:
                to_remove.append(kid)

        new_label = label.astype(float, copy=True)
        removed = 0
        for cid in to_remove:
            m = (label == cid)
            if np.any(m):
                new_label[m] = np.nan
                removed += 1

        # 모두 제거되면 저장 스킵
        keep = new_label[np.isfinite(new_label)]
        keep = keep[keep != 0]
        if keep.size == 0:
            continue

        day = dt[:8]
        out_dir = os.path.join(SAVE_DIR, day)
        os.makedirs(out_dir, exist_ok=True)
        out_path = os.path.join(out_dir, f"{dt}_label.nc")
        xr.Dataset({"label": (["y","x"], new_label)}).to_netcdf(out_path)
        # print(f"   ✅ {dt}: {removed}개 제거 → {out_path}")

    print("\n🎉 완료: 단순화된 구름 마스킹 (즉시 조건 + 미래 후손 제거)")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Remove mature-cloud objects from CI labels.")
    parser.add_argument("--config", required=True, help="JSON configuration path")
    parser.add_argument("--device", default=None, help="Accepted for a common CLI; this stage runs on CPU")
    parser.add_argument("--output-dir", default=None, help="Override save_dir")
    args = parser.parse_args()
    config_path = os.path.abspath(args.config)
    RUN_DIR = os.path.dirname(config_path)
    cfg = load_config(config_path)
    if args.output_dir:
        cfg["output_dir"] = os.path.abspath(args.output_dir)
    apply_config(cfg)
    process_mature_cloud_masking(START_TIME, END_TIME)