"""Upstream and downstream chain building."""
from collections import defaultdict
import numpy as np

from .stats import MIN_CHAIN_OBS, MAX_UPSTREAM, MAX_DOWNSTREAM, est_pax


def build_upstream(pairs):
    """Upstream feeder routes per time window."""
    print("Building upstream chains per time window...")

    pairs["curr_key"] = pairs["curr_origin"] + "-" + pairs["curr_dest"]
    pairs["prev_key"] = pairs["prev_origin"] + "-" + pairs["prev_dest"]

    result = defaultdict(lambda: defaultdict(list))  # route -> time -> [links]

    grouped = pairs.groupby(["curr_key", "curr_time", "prev_key"])
    for (curr_key, tw, prev_key), g in grouped:
        n = len(g)
        if n < MIN_CHAIN_OBS:
            continue

        prev_d = g["prev_arr_delay"].dropna().values

        # avg inherited delay -- how much of the upstream delay bleeds through
        avg_inh = float(np.nanmean(prev_d)) if len(prev_d) > 0 else 0

        # risk classification
        if avg_inh < 10:
            risk = "low"
        elif avg_inh < 25:
            risk = "moderate"
        else:
            risk = "high"

        parts = prev_key.split("-")
        result[curr_key][tw].append({
            "origin": parts[0],
            "dest": parts[1],
            "flight_count": n,
            "avg_delay_inherited": round(avg_inh, 1),
            "risk": risk,
        })

    # sort and trim to top N per time window
    for route in result:
        for tw in result[route]:
            result[route][tw].sort(key=lambda x: x["flight_count"], reverse=True)
            result[route][tw] = result[route][tw][:MAX_UPSTREAM]

    cnt = sum(1 for r in result if any(result[r][tw] for tw in result[r]))
    print(f"  {cnt} routes have upstream data")
    return result


def build_downstream(pairs):
    """Downstream continuation routes per time window."""
    print("Building downstream chains per time window...")

    if "curr_key" not in pairs.columns:
        pairs["curr_key"] = pairs["curr_origin"] + "-" + pairs["curr_dest"]
    if "prev_key" not in pairs.columns:
        pairs["prev_key"] = pairs["prev_origin"] + "-" + pairs["prev_dest"]

    result = defaultdict(lambda: defaultdict(list))

    grouped = pairs.groupby(["prev_key", "prev_time", "curr_key"])
    for (prev_key, tw, curr_key), g in grouped:
        n = len(g)
        if n < MIN_CHAIN_OBS:
            continue

        avg_prop = float(g["curr_dep_delay"].mean()) if len(g) > 0 else 0
        avg_dist = float(g["curr_distance"].mean()) if g["curr_distance"].notna().any() else 800

        parts = curr_key.split("-")
        result[prev_key][tw].append({
            "origin": parts[0],
            "dest": parts[1],
            "flight_count": n,
            "avg_delay_propagated": round(avg_prop, 1),
            "est_pax_impact": est_pax(n, avg_dist),
        })

    for route in result:
        for tw in result[route]:
            result[route][tw].sort(key=lambda x: x["flight_count"], reverse=True)
            result[route][tw] = result[route][tw][:MAX_DOWNSTREAM]

    cnt = sum(1 for r in result if any(result[r][tw] for tw in result[r]))
    print(f"  {cnt} routes have downstream data")
    return result
