
from core.flow_algorithms import FlowRouter, HAS_NUMBA
import sys
import os
import time
import numpy as np

# Add plugin directory to path
plugin_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.append(plugin_dir)


def test_flow_algorithms():
    print(f"Numba Available: {HAS_NUMBA}")
    if HAS_NUMBA:
        import numba
        print(f"Numba Version: {numba.__version__}")
    else:
        print("WARNING: Numba is NOT available. Algorithms will run in pure Python (slow).")

    print("Generating synthetic DEM (1000x1000)...")
    rows, cols = 1000, 1000
    # Create a simple slope
    y, x = np.mgrid[0:rows, 0:cols]
    dem = (rows - y) + (cols - x) + np.random.rand(rows, cols) * 0.1
    dem = dem.astype(np.float32)

    cellsize = 10.0

    router = FlowRouter(dem, cellsize, cellsize)

    print("\n--- Testing D8 Flow Direction (Parallelized) ---")
    start_time = time.time()
    flow_dir = router.d8_flow_direction()
    end_time = time.time()
    print(f"Time: {end_time - start_time:.4f} seconds")

    if np.all(flow_dir == 0):
        print("ERROR: Flow direction is all zeros!")
    else:
        print("Flow direction computed successfully.")

    print("\n--- Testing D8 Flow Accumulation (Optimized Queue) ---")
    start_time = time.time()
    acc = router.d8_flow_accumulation(flow_dir)
    end_time = time.time()
    print(f"Time: {end_time - start_time:.4f} seconds")

    if np.all(acc == 1):
        print("WARNING: Flow accumulation seems to be all 1s (no accumulation).")
    else:
        print(f"Max accumulation: {np.max(acc)}")

    # 3. Flow Length (Downstream)
    print("\n--- Testing Flow Length (Downstream / Distance to Outlet) ---")
    start_time = time.time()
    # Use the new convenience method
    dist_down = router.flow_distance_to_outlet(flow_dir)
    end_time = time.time()
    print(f"Time: {end_time - start_time:.4f} seconds")

    if np.all(dist_down == -1):
        print("ERROR: Flow length (downstream) is all -1!")
    else:
        print(f"Max length: {np.nanmax(dist_down)}")

    # 4. Flow Length (Upstream)
    print("\n--- Testing Flow Length (Upstream / Distance to Ridge) ---")
    start_time = time.time()
    # Use the new convenience method
    dist_up = router.flow_distance_upstream(flow_dir)
    end_time = time.time()
    print(f"Time: {end_time - start_time:.4f} seconds")

    if np.all(dist_up == 0):
        print("WARNING: Flow length (upstream) is all 0.")
    else:
        print(f"Max length: {np.nanmax(dist_up)}")

    # 5. Remove Short Streams
    print("\n--- Testing Remove Short Streams ---")
    # Create dummy streams (thresholding accumulation)
    streams = (acc > 100).astype(np.int8)
    print(f"Initial stream cells: {np.sum(streams)}")

    start_time = time.time()
    cleaned = router.remove_short_streams(flow_dir, streams, min_length=50.0)
    end_time = time.time()
    print(f"Time: {end_time - start_time:.4f} seconds")
    print(f"Remaining stream cells: {np.sum(cleaned)}")


if __name__ == "__main__":
    try:
        test_flow_algorithms()
        print("\nVerification PASSED.")
    except Exception as e:
        print(f"\nVerification FAILED: {e}")
        import traceback
        traceback.print_exc()
