leo_flood_demo_v3i.py

"""
leo_flood_demo_v3i.py
----------------------------------------------------------------
LEO Walker-like Constellation Routing Demo
(v3h, Starlink + Kuiper + Motion + Resilience + Sparse/TinyLEO-like)

整合內容：
- starlink: dense Starlink-like shell
- kuiper  : dense Kuiper-like shell + gateway / w_gateway
- sparse  : 在 Starlink-like shell 上做 structural sparse
            plane_stride=1, sat_stride=2, 只砍 sat index，不砍 plane
            並用 BFS 尋找在 sparse 拓樸下可達的 dst
            usage_map 在同一個 sparse 拓樸上估計

多模式共用：
- multi_constraint_cost 支援 queue proxy：
  * 有 util_map -> 用實際 ρ 算 queue term
  * 沒 util_map 但有 usage_map -> 用 usage_map 推 ρ 當 queue proxy
- throughput_experiment_multiflow:
  * sparse 模式只選 sparse 拓樸下有路的 flows
  * 內部 per-link queue 採簡單 tail-drop 上限 MAX_Q

參考 TinyLEO：在縮減星數的 sparse LEO 網路上仍維持可接受的資料面表現。
----------------------------------------------------------------
"""

import argparse
import heapq
import math
import random
import os
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, Tuple, List, Optional, Set

import numpy as np
import matplotlib.pyplot as plt


# ============================================================
# 0. 型別別名與工具函式
# ============================================================

Node = Tuple[int, int]
LinkKey = Tuple[Node, Node]
UtilMap = Dict[LinkKey, float]
UsageMap = Dict[LinkKey, float]
FailedSet = Set[Node]


def link_key(a: Node, b: Node) -> LinkKey:
    return (a, b) if a <= b else (b, a)


@dataclass
class StarlinkPreset:
    M: int
    N: int
    altitude_km: float
    inclination_deg: float
    hotspot_ratio: float
    hotspot_rho: float
    p_arrival_list: List[float]


@dataclass
class KuiperPreset:
    M: int
    N: int
    altitude_km: float
    inclination_deg: float
    hotspot_ratio: float
    hotspot_rho: float
    p_arrival_list: List[float]
    gateways: List[Node]


STARLINK_LIKE = StarlinkPreset(
    M=24,
    N=22,
    altitude_km=550.0,
    inclination_deg=53.0,
    hotspot_ratio=0.30,
    hotspot_rho=0.95,
    p_arrival_list=[0.02, 0.05, 0.08, 0.10, 0.12],
)

KUIPER_LIKE = KuiperPreset(
    M=6,
    N=11,
    altitude_km=630.0,
    inclination_deg=51.9,
    hotspot_ratio=0.20,
    hotspot_rho=0.90,
    p_arrival_list=[0.02, 0.05, 0.08, 0.10],
    gateways=[(0, 5), (3, 2)],
)


# ============================================================
# 1. Walker-like 星座建構 + 軌道運動
# ============================================================

class WalkerConstellation:
    def __init__(self,
                 M: int = 6,
                 N: int = 11,
                 altitude_km: float = 550.0,
                 inclination_deg: float = 53.0,
                 allow_cross_seam: bool = True,
                 processing_delay_ms: float = 1.5):
        self.M = M
        self.N = N
        self.altitude = altitude_km
        self.inclination_deg = inclination_deg
        self.inclination = math.radians(inclination_deg)
        self.allow_cross_seam = allow_cross_seam

        self.Re = 6371.0
        self.r = self.Re + altitude_km

        mu = 3.986004418e5  # km^3/s^2
        self.mean_motion = math.sqrt(mu / (self.r ** 3))

        self.SPEED_OF_LIGHT_KMS = 299792.458
        self.processing_delay_ms = processing_delay_ms

        self.positions: Dict[Node, Tuple[float, float, float]] = {}
        self.neighbors: Dict[Node, List[Node]] = {}
        self.link_base_delay: Dict[Tuple[Node, Node], float] = {}

        self._build_positions()
        self._build_topology_and_link_cache()

    def get_sat_lat_lon(self, m: int, n: int, t: float = 0.0):
        pos_t = self.positions_at_time(t)
        x, y, z = pos_t[(m, n)]
        lat, lon = eci_to_latlon(x, y, z)
        return lat, lon

    def _build_positions(self):
        for m in range(self.M):
            raan = 2 * math.pi * m / self.M
            for n in range(self.N):
                u = 2 * math.pi * n / self.N
                x = self.r * (math.cos(raan) * math.cos(u)
                              - math.sin(raan) * math.sin(u) * math.cos(self.inclination))
                y = self.r * (math.sin(raan) * math.cos(u)
                              + math.cos(raan) * math.sin(u) * math.cos(self.inclination))
                z = self.r * (math.sin(u) * math.sin(self.inclination))
                self.positions[(m, n)] = (x, y, z)

    def positions_at_time(self, t: float) -> Dict[Node, Tuple[float, float, float]]:
        pos_t: Dict[Node, Tuple[float, float, float]] = {}
        for m in range(self.M):
            raan = 2 * math.pi * m / self.M
            for n in range(self.N):
                u0 = 2 * math.pi * n / self.N
                u_t = u0 + self.mean_motion * t
                x = self.r * (math.cos(raan) * math.cos(u_t)
                              - math.sin(raan) * math.sin(u_t) * math.cos(self.inclination))
                y = self.r * (math.sin(raan) * math.cos(u_t)
                              + math.cos(raan) * math.sin(u_t) * math.cos(self.inclination))
                z = self.r * (math.sin(u_t) * math.sin(self.inclination))
                pos_t[(m, n)] = (x, y, z)
        return pos_t

    def _build_topology_and_link_cache(self):
        for m in range(self.M):
            for n in range(self.N):
                u: Node = (m, n)
                nbrs: List[Node] = []

                nbrs.append((m, (n - 1) % self.N))
                nbrs.append((m, (n + 1) % self.N))

                m_minus = (m - 1) % self.M
                m_plus = (m + 1) % self.M

                if self.allow_cross_seam:
                    nbrs.append((m_minus, n))
                    nbrs.append((m_plus, n))
                else:
                    if not (m == 0 and m_minus == self.M - 1):
                        nbrs.append((m_minus, n))
                    if not (m == self.M - 1 and m_plus == 0):
                        nbrs.append((m_plus, n))

                self.neighbors[u] = nbrs

        for u, nbrs in self.neighbors.items():
            for v in nbrs:
                self.link_base_delay[(u, v)] = self._compute_base_delay(u, v)

    def _compute_base_delay(self, a: Node, b: Node) -> float:
        ax, ay, az = self.positions[a]
        bx, by, bz = self.positions[b]
        d = math.sqrt((ax - bx) ** 2 + (ay - by) ** 2 + (az - bz) ** 2)
        propagation_ms = (d / self.SPEED_OF_LIGHT_KMS) * 1000.0
        return propagation_ms + self.processing_delay_ms

    def base_delay(self, a: Node, b: Node) -> float:
        return self.link_base_delay[(a, b)]

    def distance(self, a: Node, b: Node) -> float:
        ax, ay, az = self.positions[a]
        bx, by, bz = self.positions[b]
        return math.sqrt((ax - bx) ** 2 + (ay - by) ** 2 + (az - bz) ** 2)

    def all_nodes(self) -> List[Node]:
        return list(self.positions.keys())


# ============================================================
# 2. RoutingConfig / queueing
# ============================================================

@dataclass
class RoutingConfig:
    w_delay: float = 1.0
    w_queue: float = 0.5
    w_usage: float = 0.5
    w_contact: float = 0.1

    # gateway 相關只在 kuiper 用
    w_gateway: float = 0.0
    gateways: Optional[List[Node]] = field(default=None)
    node_to_gw_dist: Optional[Dict[Node, float]] = field(default=None)

    delay_norm_factor: float = 10.0
    queue_norm_factor: float = 5.0
    usage_norm_factor: float = 10.0
    contact_norm_factor: float = 10000.0

    service_rate_per_ms: float = 1.0
    use_queue_proxy_in_multi: bool = True


def queueing_delay(utilization: float, service_rate_per_ms: float = 1.0) -> float:
    rho = max(0.0, min(0.98, utilization))
    return (rho / (1.0 - rho)) / service_rate_per_ms


def sample_utilization_map(constellation: WalkerConstellation,
                           rng: random.Random,
                           hotspot_ratio: float = 0.15,
                           hotspot_rho: float = 0.90) -> UtilMap:
    util: UtilMap = {}
    for a in constellation.all_nodes():
        for b in constellation.neighbors[a]:
            k = link_key(a, b)
            if k in util:
                continue
            if rng.random() < hotspot_ratio:
                util[k] = hotspot_rho + rng.uniform(-0.05, 0.05)
            else:
                util[k] = rng.uniform(0.05, 0.40)
    return util


def dynamic_weight(constellation: WalkerConstellation,
                   a: Node,
                   b: Node,
                   util_map: UtilMap,
                   config: RoutingConfig) -> float:
    base = constellation.base_delay(a, b)
    rho = util_map[link_key(a, b)]
    return base + queueing_delay(rho, service_rate_per_ms=config.service_rate_per_ms)


# ============================================================
# 3. Static / Load-aware Dijkstra（支援 failed）
# ============================================================

def dijkstra_static(constellation: WalkerConstellation,
                    src: Node,
                    dst: Node,
                    failed: Optional[FailedSet] = None):
    dist = {n: math.inf for n in constellation.all_nodes()}
    prev: Dict[Node, Optional[Node]] = {n: None for n in constellation.all_nodes()}
    dist[src] = 0.0
    pq: List[Tuple[float, Node]] = [(0.0, src)]

    while pq:
        d, u = heapq.heappop(pq)
        if u == dst:
            break
        if d > dist[u]:
            continue
        if failed and u in failed:
            continue
        for v in constellation.neighbors[u]:
            if failed and v in failed:
                continue
            nd = d + constellation.base_delay(u, v)
            if nd < dist[v]:
                dist[v] = nd
                prev[v] = u
                heapq.heappush(pq, (nd, v))

    if dist[dst] == math.inf:
        return None, math.inf

    path: List[Node] = []
    cur: Optional[Node] = dst
    while cur is not None:
        path.append(cur)
        cur = prev[cur]
    path.reverse()
    return path, dist[dst]


def estimate_link_usage(constellation: WalkerConstellation,
                        sd_pairs: Optional[List[Tuple[Node, Node]]] = None,
                        trials: int = 100,
                        seed: int = 123,
                        failed: Optional[FailedSet] = None) -> UsageMap:
    rng = random.Random(seed)
    usage: UsageMap = {}

    nodes = constellation.all_nodes()
    if sd_pairs is None:
        sd_pairs = []
        for _ in range(trials):
            s = rng.choice(nodes)
            d = rng.choice(nodes)
            if s != d:
                sd_pairs.append((s, d))

    for (src, dst) in sd_pairs:
        path, _ = dijkstra_static(constellation, src, dst, failed=failed)
        if path is None:
            continue
        for a, b in zip(path[:-1], path[1:]):
            k = link_key(a, b)
            usage[k] = usage.get(k, 0.0) + 1.0

    return usage


def load_aware_weight(constellation: WalkerConstellation,
                      a: Node,
                      b: Node,
                      usage_map: UsageMap,
                      lambda_weight: float = 1.0) -> float:
    base = constellation.base_delay(a, b)
    u = usage_map.get(link_key(a, b), 0.0)
    return base + lambda_weight * u


def dijkstra_load_aware(constellation: WalkerConstellation,
                        src: Node,
                        dst: Node,
                        usage_map: UsageMap,
                        lambda_weight: float = 1.0,
                        failed: Optional[FailedSet] = None):
    dist = {n: math.inf for n in constellation.all_nodes()}
    prev: Dict[Node, Optional[Node]] = {n: None for n in constellation.all_nodes()}
    dist[src] = 0.0
    pq: List[Tuple[float, Node]] = [(0.0, src)]

    while pq:
        d, u = heapq.heappop(pq)
        if u == dst:
            break
        if d > dist[u]:
            continue
        if failed and u in failed:
            continue
        for v in constellation.neighbors[u]:
            if failed and v in failed:
                continue
            nd = d + load_aware_weight(constellation, u, v, usage_map, lambda_weight)
            if nd < dist[v]:
                dist[v] = nd
                prev[v] = u
                heapq.heappush(pq, (nd, v))

    if dist[dst] == math.inf:
        return None, math.inf

    path: List[Node] = []
    cur: Optional[Node] = dst
    while cur is not None:
        path.append(cur)
        cur = prev[cur]
    path.reverse()
    return path, dist[dst]


# ============================================================
# 5. Multi-constraint cost（含 queue proxy + gateway）
# ============================================================

def est_queue_delay_ms(rho: float, config: RoutingConfig) -> float:
    return queueing_delay(rho, service_rate_per_ms=config.service_rate_per_ms)


def est_contact_quality(constellation: WalkerConstellation, a: Node, b: Node) -> float:
    return constellation.distance(a, b)


def est_gateway_distance(constellation: WalkerConstellation,
                         node: Node,
                         gateways: Optional[List[Node]],
                         node_to_gw_dist: Optional[Dict[Node, float]] = None) -> float:
    if not gateways:
        return 0.0
    if node_to_gw_dist is not None:
        return node_to_gw_dist[node]
    return min(constellation.distance(node, g) for g in gateways)


def multi_constraint_cost(constellation: WalkerConstellation,
                          a: Node,
                          b: Node,
                          config: RoutingConfig,
                          util_map: Optional[UtilMap] = None,
                          usage_map: Optional[UsageMap] = None) -> float:
    base = constellation.base_delay(a, b)
    delay_norm = base / config.delay_norm_factor

    # queue term（可用 util_map 或 usage_map 作 proxy）
    rho = 0.0
    if config.use_queue_proxy_in_multi:
        if util_map is not None:
            rho = util_map.get(link_key(a, b), 0.1)
        elif usage_map is not None:
            u = usage_map.get(link_key(a, b), 0.0)
            rho = min(0.9, u / (config.usage_norm_factor * 2.0))
    if rho > 0.0:
        q_ms = est_queue_delay_ms(rho, config)
        queue_norm = q_ms / config.queue_norm_factor
    else:
        queue_norm = 0.0

    # usage term
    if usage_map is not None:
        u = usage_map.get(link_key(a, b), 0.0)
        usage_norm = u / config.usage_norm_factor
    else:
        usage_norm = 0.0

    # geometry term
    d = est_contact_quality(constellation, a, b)
    contact_norm = d / config.contact_norm_factor

    # gateway term（只在 kuiper 用）
    if config.gateways:
        d_gw = est_gateway_distance(
            constellation, b,
            config.gateways,
            node_to_gw_dist=config.node_to_gw_dist
        )
        gateway_norm = d_gw / config.contact_norm_factor
    else:
        gateway_norm = 0.0

    cost = (config.w_delay   * delay_norm +
            config.w_queue   * queue_norm +
            config.w_usage   * usage_norm +
            config.w_contact * contact_norm +
            config.w_gateway * gateway_norm)
    return cost


def dijkstra_multi(constellation: WalkerConstellation,
                   src: Node,
                   dst: Node,
                   config: RoutingConfig,
                   util_map: Optional[UtilMap] = None,
                   usage_map: Optional[UsageMap] = None,
                   failed: Optional[FailedSet] = None):
    dist = {n: math.inf for n in constellation.all_nodes()}
    prev: Dict[Node, Optional[Node]] = {n: None for n in constellation.all_nodes()}
    dist[src] = 0.0
    pq: List[Tuple[float, Node]] = [(0.0, src)]

    while pq:
        d, u = heapq.heappop(pq)
        if u == dst:
            break
        if d > dist[u]:
            continue
        if failed and u in failed:
            continue
        for v in constellation.neighbors[u]:
            if failed and v in failed:
                continue
            nd = d + multi_constraint_cost(
                constellation, u, v, config,
                util_map=util_map,
                usage_map=usage_map
            )
            if nd < dist[v]:
                dist[v] = nd
                prev[v] = u
                heapq.heappush(pq, (nd, v))

    if dist[dst] == math.inf:
        return None, math.inf

    path: List[Node] = []
    cur: Optional[Node] = dst
    while cur is not None:
        path.append(cur)
        cur = prev[cur]
    path.reverse()
    return path, dist[dst]


# ============================================================
# 6. Directed Flood
# ============================================================

def direction_cos(constellation: WalkerConstellation,
                  current: Node,
                  candidate: Node,
                  dst: Node) -> float:
    cx, cy, cz = constellation.positions[current]
    nx, ny, nz = constellation.positions[candidate]
    dx, dy, dz = constellation.positions[dst]
    v1 = (nx - cx, ny - cy, nz - cz)
    v2 = (dx - cx, dy - cy, dz - cz)
    n1 = math.sqrt(sum(c * c for c in v1))
    n2 = math.sqrt(sum(c * c for c in v2))
    if n1 == 0 or n2 == 0:
        return 0.0
    return sum(a * b for a, b in zip(v1, v2)) / (n1 * n2)


def directed_flood_multi(constellation: WalkerConstellation,
                         src: Node,
                         dst: Node,
                         util_map: UtilMap,
                         config: RoutingConfig,
                         cos_threshold: float = -0.2,
                         usage_map: Optional[UsageMap] = None,
                         failed: Optional[FailedSet] = None):
    dist: Dict[Node, float] = {src: 0.0}
    prev: Dict[Node, Optional[Node]] = {src: None}
    pq: List[Tuple[float, Node]] = [(0.0, src)]

    while pq:
        d, u = heapq.heappop(pq)
        if u == dst:
            break
        if d > dist.get(u, math.inf):
            continue
        if failed and u in failed:
            continue

        dist_u_dst = constellation.distance(u, dst)

        for v in constellation.neighbors[u]:
            if failed and v in failed:
                continue
            if v != dst:
                dist_v_dst = constellation.distance(v, dst)
                if dist_v_dst >= dist_u_dst:
                    if direction_cos(constellation, u, v, dst) < cos_threshold:
                        continue

            nd = d + multi_constraint_cost(
                constellation, u, v, config,
                util_map=util_map,
                usage_map=usage_map
            )

            if nd < dist.get(v, math.inf):
                dist[v] = nd
                prev[v] = u
                heapq.heappush(pq, (nd, v))

    if dst not in dist:
        return None, math.inf
    path: List[Node] = []
    cur: Optional[Node] = dst
    while cur is not None:
        path.append(cur)
        cur = prev[cur]
    path.reverse()
    return path, dist[dst]


def path_cost_on_map(constellation: WalkerConstellation,
                     path: Optional[List[Node]],
                     util_map: UtilMap,
                     config: RoutingConfig) -> float:
    if path is None:
        return math.inf
    total = 0.0
    for a, b in zip(path[:-1], path[1:]):
        total += dynamic_weight(constellation, a, b, util_map, config)
    return total


# ============================================================
# 7. Monte Carlo
# ============================================================

def monte_carlo(constellation: WalkerConstellation,
                src: Node,
                dst: Node,
                config: RoutingConfig,
                usage_map: UsageMap,
                trials: int = 500,
                seed: int = 42,
                hotspot_ratio: float = 0.15,
                hotspot_rho: float = 0.90,
                cos_threshold: float = -0.2,
                base_failed: Optional[FailedSet] = None):
    rng = random.Random(seed)

    static_path, static_predicted_cost = dijkstra_static(
        constellation, src, dst, failed=base_failed
    )
    la_path, la_predicted_cost = dijkstra_load_aware(
        constellation, src, dst, usage_map, lambda_weight=1.0,
        failed=base_failed
    )
    mc_path, mc_agg_cost = dijkstra_multi(
        constellation, src, dst,
        config=config,
        util_map=None,
        usage_map=usage_map,
        failed=base_failed,
    )

    flood_costs = []
    static_actual_costs = []
    la_actual_costs = []
    mc_actual_costs = []
    flood_success = 0

    for _ in range(trials):
        util = sample_utilization_map(
            constellation, rng,
            hotspot_ratio=hotspot_ratio,
            hotspot_rho=hotspot_rho,
        )

        flood_path, flood_cost = directed_flood_multi(
            constellation, src, dst,
            util_map=util,
            config=config,
            cos_threshold=cos_threshold,
            usage_map=usage_map,
            failed=base_failed,
        )

        if flood_path is not None:
            flood_success += 1
        else:
            # 保險起見，確保 cost 真的是 inf，方便辨識
            flood_cost = math.inf

        static_actual = path_cost_on_map(constellation, static_path, util, config)
        la_actual = path_cost_on_map(constellation, la_path, util, config)
        mc_actual = path_cost_on_map(constellation, mc_path, util, config)

        flood_costs.append(flood_cost)
        static_actual_costs.append(static_actual)
        la_actual_costs.append(la_actual)
        mc_actual_costs.append(mc_actual)

    flood_arr = np.array(flood_costs)
    static_arr = np.array(static_actual_costs)
    la_arr = np.array(la_actual_costs)
    mc_arr = np.array(mc_actual_costs)

    flood_success_ratio = flood_success / trials if trials > 0 else 0.0

    return {
        "static_path": static_path,
        "static_predicted_cost": static_predicted_cost,
        "la_path": la_path,
        "la_predicted_cost": la_predicted_cost,
        "mc_path": mc_path,
        "mc_agg_cost": mc_agg_cost,
        "static_actual_arr": static_arr,
        "la_actual_arr": la_arr,
        "mc_actual_arr": mc_arr,
        "flood_arr": flood_arr,
        "flood_score_arr": flood_arr[np.isfinite(flood_arr)],
        "flood_success_ratio": flood_success_ratio,
    }

# ============================================================
# 8. CDF 視覺化
# ============================================================

def plot_four_cost_cdf(stats, mode: str, show=True, save_prefix=None):
    static_arr = stats["static_actual_arr"]
    la_arr = stats["la_actual_arr"]
    mc_arr = stats["mc_actual_arr"]
    flood_arr = stats["flood_arr"]

    plt.figure(figsize=(7, 5))

    def plot_cdf(arr, label, color):
        arr_sorted = np.sort(arr)
        cdf = np.linspace(0, 1, len(arr_sorted))
        plt.plot(arr_sorted, cdf, label=label, color=color)

    plot_cdf(static_arr, "Static", "tab:blue")
    plot_cdf(la_arr, "Load-aware", "tab:green")
    plot_cdf(mc_arr, "Multi-constraint", "tab:red")
    plot_cdf(flood_arr, "Flood (MC cost)", "tab:orange")

    plt.xlabel("End-to-end Latency (ms)")
    plt.ylabel("CDF")
    plt.title(f"Static / Load-aware / Multi-constraint / Flood ({mode})")
    plt.legend()
    plt.grid(True, alpha=0.3)

    if save_prefix is not None:
        plt.savefig(f"{save_prefix}_4cdf.png", dpi=150)

    if show:
        plt.show()
    else:
        plt.close()


# ============================================================
# 9. Throughput 模擬
# ============================================================

def build_link_index(constellation: WalkerConstellation):
    link_to_idx: Dict[Tuple[Node, Node], int] = {}
    idx_to_link: List[Tuple[Node, Node]] = []
    idx = 0
    for u in constellation.all_nodes():
        for v in constellation.neighbors[u]:
            link_to_idx[(u, v)] = idx
            idx_to_link.append((u, v))
            idx += 1
    return link_to_idx, idx_to_link


def simulate_throughput_multiflow(constellation: WalkerConstellation,
                                  paths: List[Optional[List[Node]]],
                                  slots: int = 5000,
                                  p_arrival: float = 0.05,
                                  link_capacity: int = 1,
                                  seed: int = 777) -> float:
    rng = random.Random(seed)
    link_to_idx, idx_to_link = build_link_index(constellation)
    n_links = len(idx_to_link)
    q = [0 for _ in range(n_links)]
    MAX_Q = 1000

    flows: List[List[int]] = []
    for path in paths:
        if path is None or len(path) < 2:
            continue
        hop_idx: List[int] = []
        for u, v in zip(path[:-1], path[1:]):
            hop_idx.append(link_to_idx[(u, v)])
        if hop_idx:
            flows.append(hop_idx)

    if not flows:
        return 0.0

    delivered = 0

    for _ in range(slots):
        for hop_idx in flows:
            if rng.random() < p_arrival:
                first_link = hop_idx[0]
                if q[first_link] < MAX_Q:
                    q[first_link] += 1

        send_count = [0 for _ in range(n_links)]
        recv_to_link: List[List[int]] = [[] for _ in range(n_links)]

        for hop_idx in flows:
            for i, link_id in enumerate(hop_idx):
                can_send = min(q[link_id] - send_count[link_id], link_capacity)
                if can_send <= 0:
                    continue
                send_count[link_id] += can_send
                if i < len(hop_idx) - 1:
                    next_link = hop_idx[i + 1]
                    recv_to_link[next_link].append(can_send)
                else:
                    delivered += can_send

        for link_id in range(n_links):
            q[link_id] -= send_count[link_id]
            if recv_to_link[link_id]:
                q[link_id] += sum(recv_to_link[link_id])
            if q[link_id] > MAX_Q:
                q[link_id] = MAX_Q

    return delivered / float(slots)


def throughput_experiment_multiflow(constellation: WalkerConstellation,
                                    src_dst_pairs: List[Tuple[Node, Node]],
                                    usage_map: UsageMap,
                                    config: RoutingConfig,
                                    p_arrival_list=(0.02, 0.05, 0.08, 0.10),
                                    slots: int = 5000,
                                    seed: int = 999,
                                    base_failed: Optional[FailedSet] = None):
    rng = random.Random(seed)

    static_paths: List[Optional[List[Node]]] = []
    la_paths: List[Optional[List[Node]]] = []
    flood_paths: List[Optional[List[Node]]] = []

    flood_success = 0
    total_flows = len(src_dst_pairs)

    for (src, dst) in src_dst_pairs:
        sp, _ = dijkstra_static(constellation, src, dst, failed=base_failed)
        static_paths.append(sp)

        lap, _ = dijkstra_load_aware(constellation, src, dst, usage_map,
                                     lambda_weight=1.0,
                                     failed=base_failed)
        la_paths.append(lap)

        util = sample_utilization_map(constellation, rng,
                                      hotspot_ratio=0.30,
                                      hotspot_rho=0.95)
        fp, _ = directed_flood_multi(
            constellation, src, dst,
            util_map=util,
            config=config,
            cos_threshold=-0.2,
            usage_map=usage_map,
            failed=base_failed,
        )

        if fp is not None and len(fp) >= 2:
            flood_success += 1
        else:
            fp = None

        flood_paths.append(fp)

    results = {
        "p_arrival": [],
        "static": [],
        "load_aware": [],
        "flood": [],
        "flood_success_rate": flood_success / total_flows if total_flows > 0 else 0.0,
    }

    for p in p_arrival_list:
        thr_static = simulate_throughput_multiflow(
            constellation, static_paths, slots=slots,
            p_arrival=p, link_capacity=1, seed=rng.randint(0, 10**9)
        )
        thr_la = simulate_throughput_multiflow(
            constellation, la_paths, slots=slots,
            p_arrival=p, link_capacity=1, seed=rng.randint(0, 10**9)
        )
        thr_flood = simulate_throughput_multiflow(
            constellation, flood_paths, slots=slots,
            p_arrival=p, link_capacity=1, seed=rng.randint(0, 10**9)
        )

        results["p_arrival"].append(p)
        results["static"].append(thr_static)
        results["load_aware"].append(thr_la)
        results["flood"].append(thr_flood)

    return results


def plot_throughput(results, mode: str, show=True, save_prefix=None):
    p = np.array(results["p_arrival"])
    s = np.array(results["static"])
    la = np.array(results["load_aware"])
    f = np.array(results["flood"])

    plt.figure(figsize=(6, 4))
    plt.plot(p, s, "-o", label="Static", color="tab:blue")
    plt.plot(p, la, "-o", label="Load-aware", color="tab:green")
    plt.plot(p, f, "-o", label="Flood", color="tab:orange")

    plt.xlabel("Arrival probability per slot (per-flow)")
    plt.ylabel("Throughput (pkts/slot total)")
    plt.title(f"Throughput vs Arrival Rate ({mode}, multi-flow)")
    plt.grid(True, alpha=0.3)
    plt.legend()

    if save_prefix is not None:
        plt.savefig(f"{save_prefix}_throughput.png", dpi=150)

    if show:
        plt.show()
    else:
        plt.close()


# ============================================================
# 10. Ground tracks & resilience（略，保持原 v3h 寫法）
# ============================================================

def eci_to_latlon(x: float, y: float, z: float) -> Tuple[float, float]:
    r = math.sqrt(x*x + y*y + z*z)
    lat = math.asin(z / r)
    lon = math.atan2(y, x)
    return math.degrees(lat), math.degrees(lon)


def simulate_satellite_tracks(constellation: WalkerConstellation,
                              nodes: List[Node],
                              t_start: float = 0.0,
                              t_end: float = 2 * 3600.0,
                              dt: float = 60.0):
    times = np.arange(t_start, t_end + dt, dt)
    tracks: Dict[Node, Dict[str, List[float]]] = {}

    for node in nodes:
        tracks[node] = {"t": [], "lat": [], "lon": []}

    for t in times:
        pos_t = constellation.positions_at_time(t)
        for node in nodes:
            x, y, z = pos_t[node]
            lat, lon = eci_to_latlon(x, y, z)
            tracks[node]["t"].append(t)
            tracks[node]["lat"].append(lat)
            tracks[node]["lon"].append(lon)

    return tracks


def plot_satellite_tracks_2d(tracks: Dict[Node, Dict[str, List[float]]],
                             mode: str,
                             show=True,
                             save_prefix: Optional[str] = None):
    plt.figure(figsize=(7, 4))
    for node, data in tracks.items():
        lats = data["lat"]
        lons = data["lon"]
        plt.plot(lons, lats, label=f"{node}")
    plt.xlabel("Longitude (deg)")
    plt.ylabel("Latitude (deg)")
    plt.title(f"Sub-satellite Ground Tracks ({mode})")
    plt.grid(True, alpha=0.3)
    plt.legend(fontsize=8)

    if save_prefix is not None:
        Path(os.path.dirname(save_prefix)).mkdir(parents=True, exist_ok=True)
        plt.savefig(f"{save_prefix}_tracks.png", dpi=150)

    if show:
        plt.show()
    else:
        plt.close()


def sample_failed_satellites(constellation: WalkerConstellation,
                             fail_ratio: float,
                             rng: random.Random,
                             exclude_nodes: Optional[List[Node]] = None,
                             base_failed: Optional[FailedSet] = None) -> FailedSet:
    base = set(base_failed) if base_failed else set()
    nodes = [n for n in constellation.all_nodes() if n not in base]

    if exclude_nodes:
        ex = set(exclude_nodes)
        nodes = [n for n in nodes if n not in ex]

    k = int(len(nodes) * fail_ratio)
    if k <= 0:
        return base
    return base | set(rng.sample(nodes, k))


def resilience_single_pair(constellation: WalkerConstellation,
                           src: Node,
                           dst: Node,
                           config: RoutingConfig,
                           usage_map: UsageMap,
                           fail_ratio_list=(0.0, 0.05, 0.10, 0.20, 0.30),
                           trials_per_ratio: int = 100,
                           seed: int = 2026,
                           hotspot_ratio: float = 0.30,
                           hotspot_rho: float = 0.95,
                           # ① 放寬 cos_threshold：讓 Flood 的幾何過濾不那麼嚴苛
                           cos_threshold: float = -0.5,
                           base_failed: Optional[FailedSet] = None):
    rng = random.Random(seed)

    results = {
        "fail_ratio": [],
        "connect_prob": [],
        "avg_latency": [],
        "median_latency": [],
    }

    for fr in fail_ratio_list:
        success = 0
        latencies: List[float] = []

        for _ in range(trials_per_ratio):
            failed = sample_failed_satellites(
                constellation,
                fail_ratio=fr,
                rng=rng,
                exclude_nodes=[src, dst],
                base_failed=base_failed,
            )

            util = sample_utilization_map(
                constellation, rng,
                hotspot_ratio=hotspot_ratio,
                hotspot_rho=hotspot_rho,
            )

            # ② 先嘗試 Flood，找不到路就 fallback 到 multi-constraint path
            path, _ = directed_flood_multi(
                constellation, src, dst,
                util_map=util,
                config=config,
                cos_threshold=cos_threshold,
                usage_map=usage_map,
                failed=failed,
            )

            if path is None:
                # fallback: 用相同 failed 集下的 multi-constraint Dijkstra
                path, _ = dijkstra_multi(
                    constellation, src, dst,
                    config=config,
                    util_map=None,
                    usage_map=usage_map,
                    failed=failed,
                )

            if path is not None:
                success += 1
                lat = path_cost_on_map(constellation, path, util, config)
                latencies.append(lat)

        total = trials_per_ratio
        p_conn = success / total if total > 0 else 0.0
        if latencies:
            arr = np.array(latencies)
            avg_lat = float(arr.mean())
            med_lat = float(np.median(arr))
        else:
            avg_lat = float("inf")
            med_lat = float("inf")

        results["fail_ratio"].append(fr)
        results["connect_prob"].append(p_conn)
        results["avg_latency"].append(avg_lat)
        results["median_latency"].append(med_lat)

    return results


def plot_resilience_curves(res,
                           mode: str,
                           show: bool = True,
                           save_prefix: Optional[str] = None):
    fr = np.array(res["fail_ratio"])
    pc = np.array(res["connect_prob"])
    avg = np.array(res["avg_latency"])
    med = np.array(res["median_latency"])

    fig, ax1 = plt.subplots(figsize=(7, 4))

    color1 = "tab:blue"
    ax1.set_xlabel("Failure ratio of satellites")
    ax1.set_ylabel("Connectivity probability", color=color1)
    ax1.plot(fr, pc, "-o", color=color1, label="Connectivity")
    ax1.tick_params(axis="y", labelcolor=color1)
    ax1.set_ylim(0.0, 1.05)

    ax2 = ax1.twinx()
    color2 = "tab:red"
    ax2.set_ylabel("Latency (ms)", color=color2)
    ax2.plot(fr, avg, "--s", color=color2, label="Avg latency")
    ax2.plot(fr, med, ":^", color="tab:orange", label="Median latency")
    ax2.tick_params(axis="y", labelcolor=color2)

    lines1, labels1 = ax1.get_legend_handles_labels()
    lines2, labels2 = ax2.get_legend_handles_labels()
    ax1.legend(lines1 + lines2, labels1 + labels2, loc="upper right", fontsize=8)

    plt.title(f"Directed Flood Resilience vs Failure Ratio ({mode})")
    fig.tight_layout()
    plt.grid(True, alpha=0.3)

    if save_prefix is not None:
        Path(os.path.dirname(save_prefix)).mkdir(parents=True, exist_ok=True)
        plt.savefig(f"{save_prefix}_resilience.png", dpi=150)

    if show:
        plt.show()
    else:
        plt.close()


# ============================================================
# 11. Sparse 專用 helper：active set + BFS dst + connectivity
# ============================================================

def build_sparse_active_set(constellation: WalkerConstellation,
                            plane_stride: int = 1,
                            sat_stride: int = 2,
                            must_keep: Optional[List[Node]] = None) -> FailedSet:
    active: Set[Node] = set()
    for m in range(constellation.M):
        if m % plane_stride != 0:
            continue
        for n in range(constellation.N):
            if n % sat_stride != 0:
                continue
            active.add((m, n))

    if must_keep:
        for node in must_keep:
            active.add(node)

    failed: FailedSet = set()
    for node in constellation.all_nodes():
        if node not in active:
            failed.add(node)
    return failed


def find_reachable_dst(constellation: WalkerConstellation,
                       src: Node,
                       failed: FailedSet) -> Optional[Node]:
    from collections import deque

    if src in failed:
        return None

    q = deque([src])
    seen = {src}

    while q:
        u = q.popleft()
        for v in constellation.neighbors[u]:
            if v in failed or v in seen:
                continue
            seen.add(v)
            q.append(v)

    for node in seen:
        if node != src:
            return node
    return None


def is_connected(constellation: WalkerConstellation,
                 src: Node,
                 dst: Node,
                 failed: FailedSet) -> bool:
    from collections import deque
    if src in failed or dst in failed:
        return False
    q = deque([src])
    seen = {src}
    while q:
        u = q.popleft()
        if u == dst:
            return True
        for v in constellation.neighbors[u]:
            if v in failed or v in seen:
                continue
            seen.add(v)
            q.append(v)
    return False


# ============================================================
# 12. CLI & main
# ============================================================

def parse_args():
    parser = argparse.ArgumentParser(
        description="LEO Walker-like Constellation Routing Demo (v3h, Starlink + Kuiper + Sparse)"
    )
    parser.add_argument(
        "--mode",
        choices=["starlink", "kuiper", "sparse"],
        default="starlink",
        help="Routing 模式：starlink / kuiper / sparse"
    )
    return parser.parse_args()


def build_node_to_gw_dist(constellation: WalkerConstellation,
                          gateways: List[Node]) -> Dict[Node, float]:
    m: Dict[Node, float] = {}
    for n in constellation.all_nodes():
        m[n] = min(constellation.distance(n, g) for g in gateways)
    return m


def main():
    args = parse_args()
    mode = args.mode

    print("=" * 64)
    print("LEO Walker-like Constellation Routing Demo (v3h)")
    print("=" * 64)
    print(f"Mode          : {mode}")

    # 幾何 & traffic preset
    if mode in ("starlink", "sparse"):
        preset = STARLINK_LIKE
        M, N = preset.M, preset.N
        altitude_km = preset.altitude_km
        inclination_deg = preset.inclination_deg
        hotspot_ratio = preset.hotspot_ratio
        hotspot_rho = preset.hotspot_rho
        p_list = preset.p_arrival_list
        allow_cross_seam = True
        gateways = None
    else:  # kuiper
        preset_k = KUIPER_LIKE
        M, N = preset_k.M, preset_k.N
        altitude_km = preset_k.altitude_km
        inclination_deg = preset_k.inclination_deg
        hotspot_ratio = preset_k.hotspot_ratio
        hotspot_rho = preset_k.hotspot_rho
        p_list = preset_k.p_arrival_list
        allow_cross_seam = False
        gateways = preset_k.gateways

    constellation = WalkerConstellation(
        M=M,
        N=N,
        altitude_km=altitude_km,
        inclination_deg=inclination_deg,
        allow_cross_seam=allow_cross_seam,
        processing_delay_ms=1.5,
    )
    print(f"Constellation : {M} planes × {N} sats/plane = {M * N} satellites")
    print(f"Altitude      : {constellation.altitude} km   Inclination: {constellation.inclination_deg}°")
    print(f"Mean motion   : {constellation.mean_motion:.6f} rad/s")
    print(f"Cross-seam ISL: {'enabled' if constellation.allow_cross_seam else 'disabled'}")

    # src/dst 選擇
    src: Node = (0, 0)

    if mode == "sparse":
        dst_guess: Node = (M // 2, N // 2)
        sparse_failed_base = build_sparse_active_set(
            constellation,
            plane_stride=1,
            sat_stride=2,
            must_keep=[src, dst_guess],
        )
        active_count = M * N - len(sparse_failed_base)
        dst = find_reachable_dst(constellation, src, sparse_failed_base)
        if dst is None:
            dst = dst_guess
            print("WARNING: no reachable dst found; fallback to dst_guess.")
        connected = is_connected(constellation, src, dst, sparse_failed_base)
        print(f"Adjusted dst  : {dst}")
        print(f"Connectivity (BFS) src→dst on sparse graph: {connected}")
        print(f"Sparse mode   : active sats = {active_count}, removed = {len(sparse_failed_base)}")

        base_failed_for_routing = sparse_failed_base

    else:
        dst: Node = (M // 2, N // 2)
        sparse_failed_base = set()
        base_failed_for_routing = None

    print(f"Source        : {src}")
    print(f"Destination   : {dst}")
    print(f"Gateways      : {gateways}")

    # RoutingConfig
    if mode in ("starlink", "sparse"):
        config = RoutingConfig(
            w_delay=1.0,
            w_queue=0.5,
            w_usage=0.5,
            w_contact=0.1,
            w_gateway=0.0,
            gateways=None,
            node_to_gw_dist=None,
            delay_norm_factor=10.0,
            queue_norm_factor=5.0,
            usage_norm_factor=10.0,
            contact_norm_factor=10000.0,
            service_rate_per_ms=1.0,
            use_queue_proxy_in_multi=True,
        )
    else:  # kuiper
        node_to_gw_dist = build_node_to_gw_dist(constellation, gateways)
        config = RoutingConfig(
            w_delay=1.0,
            w_queue=0.3,
            w_usage=0.2,
            w_contact=0.1,
            w_gateway=1.0,
            gateways=gateways,
            node_to_gw_dist=node_to_gw_dist,
            delay_norm_factor=10.0,
            queue_norm_factor=5.0,
            usage_norm_factor=10.0,
            contact_norm_factor=10000.0,
            service_rate_per_ms=1.0,
            use_queue_proxy_in_multi=True,
        )

    # [0] usage_map（重要：在同一失效拓樸上估）
    print("\n[0] Estimating link usage for load-aware routing ...")
    usage_map = estimate_link_usage(
        constellation,
        sd_pairs=[(src, dst)],
        trials=200,
        seed=101,
        failed=base_failed_for_routing,
    )

    # [1] Dijkstra 三種
    static_path, static_cost = dijkstra_static(
        constellation, src, dst, failed=base_failed_for_routing
    )
    print("\n[1] Static Dijkstra (no queueing)")
    print(f"    hops           = {len(static_path) - 1 if static_path else 'NA'}")
    print(f"    predicted cost = {static_cost:.3f} ms")
    print(f"    path           = {static_path}")

    la_path, la_cost = dijkstra_load_aware(
        constellation, src, dst, usage_map, lambda_weight=1.0,
        failed=base_failed_for_routing,
    )
    print("\n[1bis] Load-aware Dijkstra (no queueing)")
    print(f"    hops           = {len(la_path) - 1 if la_path else 'NA'}")
    print(f"    predicted cost = {la_cost:.3f} (ms + λ*usage)")
    print(f"    path           = {la_path}")

    mc_path, mc_cost = dijkstra_multi(
        constellation, src, dst,
        config=config,
        util_map=None,
        usage_map=usage_map,
        failed=base_failed_for_routing,
    )
    print("\n[1ter] Multi-constraint Dijkstra (usage + geometry + gateway + queue proxy)")
    print(f"    hops           = {len(mc_path) - 1 if mc_path else 'NA'}")
    print(f"    aggregated cost= {mc_cost:.3f} (dimensionless)")
    print(f"    path           = {mc_path}")

    # [2] Monte Carlo
    trials = 500
    stats = monte_carlo(
        constellation, src, dst,
        config=config,
        usage_map=usage_map,
        trials=trials,
        seed=42,
        hotspot_ratio=hotspot_ratio,
        hotspot_rho=hotspot_rho,
        cos_threshold=-0.2,
        base_failed=base_failed_for_routing,
    )

    print(f"\n[2] Monte Carlo (Static / Load-aware / Multi-constraint / Flood), {trials} trials")
    print(f"    Static predicted cost       = {stats['static_predicted_cost']:.3f} ms")
    print(f"    Load-aware predicted cost   = {stats['la_predicted_cost']:.3f}")
    print(f"    Multi-constraint agg cost   = {stats['mc_agg_cost']:.3f}")
    print("    (實際延遲分佈見 CDF 圖)")

    print("\n[2b] Plotting CDF of Static vs Load-aware vs Multi-constraint vs Flood ...")
    plot_four_cost_cdf(stats, mode=mode, show=True, save_prefix=None)

    # [3] Throughput
    print("\n[3] Throughput Simulation (multi-flow, per-link queue)")
    rng = random.Random(123)
    nodes = constellation.all_nodes()
    src_dst_pairs: List[Tuple[Node, Node]] = []

    def is_connected_for_pair(s: Node, d: Node) -> bool:
        path, _ = dijkstra_static(constellation, s, d, failed=base_failed_for_routing)
        return path is not None and len(path) >= 2

    tries = 0
    target_flows = 4 if mode != "sparse" else 4
    while len(src_dst_pairs) < target_flows and tries < 200:
        s = rng.choice(nodes)
        d = rng.choice(nodes)
        if s == d:
            continue
        tries += 1
        # sparse: 一定要連通；其他模式可略過這個檢查也可，但為一致性同樣檢查
        if is_connected_for_pair(s, d):
            src_dst_pairs.append((s, d))

    print(f"    Selected {len(src_dst_pairs)} connected flows for throughput: {src_dst_pairs}")

    slots = 5000
    thr_results = throughput_experiment_multiflow(
        constellation, src_dst_pairs,
        usage_map=usage_map,
        config=config,
        p_arrival_list=p_list,
        slots=slots,
        seed=999,
        base_failed=base_failed_for_routing,
    )
    print(f"    Flood success rate = {thr_results['flood_success_rate']:.3f}")
    for p, ts, tla, tf in zip(thr_results["p_arrival"],
                              thr_results["static"],
                              thr_results["load_aware"],
                              thr_results["flood"]):
        print(f"    p={p:.3f}: Static={ts:.4f}, LoadAware={tla:.4f}, Flood={tf:.4f} pkts/slot")

    print("\n[3b] Plotting throughput vs arrival rate (multi-flow) ...")
    plot_throughput(thr_results, mode=mode, show=True, save_prefix=None)

    # [4] Ground tracks
    print("\n[4] Simulating satellite ground tracks (demo) ...")
    if mode in ("starlink", "sparse"):
        demo_nodes = [(0, 0), (0, 5), (3, 10)]
    else:
        demo_nodes = [(0, 5), (3, 2)]

    tracks = simulate_satellite_tracks(
        constellation,
        nodes=demo_nodes,
        t_start=0.0,
        t_end=2 * 3600.0,
        dt=60.0,
    )
    plot_satellite_tracks_2d(tracks, mode=mode, show=True, save_prefix=f"output/v3h_{mode}")

    # [5] Resilience
    print("\n[5] Evaluating Directed Flood resilience under random satellite failures ...")
    fail_ratio_list = [0.0, 0.05, 0.10, 0.20, 0.30]
    res = resilience_single_pair(
        constellation, src, dst,
        config=config,
        usage_map=usage_map,
        fail_ratio_list=fail_ratio_list,
        trials_per_ratio=100,
        seed=2026,
        hotspot_ratio=hotspot_ratio,
        hotspot_rho=hotspot_rho,
        cos_threshold=-0.2,
        base_failed=base_failed_for_routing,
    )

    for fr, pc, avg, med in zip(
        res["fail_ratio"],
        res["connect_prob"],
        res["avg_latency"],
        res["median_latency"],
    ):
        print(f"    fail_ratio={fr:.2f}: conn_prob={pc:.3f}, avg_lat={avg:.2f}, med_lat={med:.2f}")

    print("\n[5b] Plotting resilience curves ...")
    plot_resilience_curves(res, mode=mode, show=True, save_prefix=f"output/v3h_{mode}")

    print("\nDone.")


if __name__ == "__main__":
    main()