import unittest

import numpy as np

from B05_wf2_Route.B05_wf2_Route_Engine_Skeleton import (
    d8_flow_accumulation_numpy,
    extract_skeleton_from_grid,
)


class D8FlowAccumulationTest(unittest.TestCase):
    def test_valley_concentrates_flow(self) -> None:
        # V자 계곡: 중앙 열(x=20)이 가장 낮다 → 흐름이 중앙에 모인다.
        coords = np.linspace(0, 40, 41)
        gx, gy = np.meshgrid(coords, coords)
        z = 10.0 + np.abs(gx - 20.0) * 0.3
        valid = np.ones_like(z, dtype=bool)

        acc = d8_flow_accumulation_numpy(z, valid)
        # 누적 최대 셀은 중앙 열(index 20) 근처여야 한다.
        _, max_col = np.unravel_index(acc.argmax(), acc.shape)
        self.assertTrue(18 <= max_col <= 22)
        self.assertGreater(acc.max(), 10.0)

    def test_flat_grid_uniform(self) -> None:
        z = np.full((10, 10), 5.0)
        valid = np.ones_like(z, dtype=bool)
        acc = d8_flow_accumulation_numpy(z, valid)
        # 평지는 하강 이웃이 없어 각 셀 누적이 1.
        self.assertTrue(np.allclose(acc, 1.0))


class ExtractSkeletonTest(unittest.TestCase):
    def test_valley_detected_with_low_threshold(self) -> None:
        coords = np.linspace(0, 40, 41)
        gx, gy = np.meshgrid(coords, coords)
        z = 10.0 + np.abs(gx - 20.0) * 0.3
        valid = np.ones_like(z, dtype=bool)
        thresholds = {
            "valley_acc": 5.0,
            "main_valley_acc": 30.0,
            "ridge_acc": 5.0,
            "main_ridge_acc": 30.0,
        }
        result = extract_skeleton_from_grid(
            coords, coords, z, valid, 1.0, thresholds=thresholds, use_whitebox=False
        )
        self.assertEqual(set(result), {"main_ridge", "minor_ridge", "main_valley", "minor_valley"})
        # 계곡 클래스에 최소 하나의 polyline이 검출되어야 한다.
        valley_total = len(result["main_valley"]) + len(result["minor_valley"])
        self.assertGreater(valley_total, 0)