"""Phase coordination helpers for TNFR dynamics."""

from __future__ import annotations

import math
from collections import deque
from collections.abc import Mapping, MutableMapping, Sequence
from concurrent.futures import ProcessPoolExecutor
from typing import TYPE_CHECKING, Any, TypeVar, cast

from ..alias import get_theta_attr, set_theta
from ..constants import (
    DEFAULTS,
    METRIC_DEFAULTS,
    STATE_DISSONANT,
    STATE_STABLE,
    STATE_TRANSITION,
    normalise_state_token,
)
from ..glyph_history import append_metric
from ..helpers.numeric import angle_diff
from ..metrics.common import ensure_neighbors_map
from ..metrics.trig import neighbor_phase_mean_list
from ..metrics.trig_cache import get_trig_cache
from ..observers import DEFAULT_GLYPH_LOAD_SPAN, glyph_load, kuramoto_order
from ..types import NodeId, Phase, TNFRGraph
from ..utils import get_numpy
from .._compat import TypeAlias

if TYPE_CHECKING:  # pragma: no cover - typing imports only
    try:
        import numpy as np_typing
        import numpy.typing as npt
    except ImportError:  # pragma: no cover - optional typing dependency
        FloatArray: TypeAlias = Any
    else:
        FloatArray: TypeAlias = npt.NDArray[np_typing.float_]
else:  # pragma: no cover - runtime without numpy typing
    FloatArray: TypeAlias = Any

_DequeT = TypeVar("_DequeT")

ChunkArgs = tuple[
    Sequence[NodeId],
    Mapping[NodeId, Phase],
    Mapping[NodeId, float],
    Mapping[NodeId, float],
    Mapping[NodeId, Sequence[NodeId]],
    float,
    float,
    float,
]

__all__ = ("coordinate_global_local_phase",)


def _ensure_hist_deque(
    hist: MutableMapping[str, Any], key: str, maxlen: int
) -> deque[_DequeT]:
    """Ensure history entry ``key`` is a deque with ``maxlen``."""

    dq = hist.setdefault(key, deque(maxlen=maxlen))
    if not isinstance(dq, deque):
        dq = deque(dq, maxlen=maxlen)
        hist[key] = dq
    return cast("deque[_DequeT]", dq)


def _read_adaptive_params(
    g: Mapping[str, Any],
) -> tuple[Mapping[str, Any], float, float]:
    """Obtain configuration and current values for phase adaptation."""

    cfg = g.get("PHASE_ADAPT", DEFAULTS.get("PHASE_ADAPT", {}))
    kG = float(g.get("PHASE_K_GLOBAL", DEFAULTS["PHASE_K_GLOBAL"]))
    kL = float(g.get("PHASE_K_LOCAL", DEFAULTS["PHASE_K_LOCAL"]))
    return cast(Mapping[str, Any], cfg), kG, kL


def _compute_state(G: TNFRGraph, cfg: Mapping[str, Any]) -> tuple[str, float, float]:
    """Return the canonical network state and supporting metrics."""

    R = kuramoto_order(G)
    dist = glyph_load(G, window=DEFAULT_GLYPH_LOAD_SPAN)
    disr = float(dist.get("_disruptors", 0.0)) if dist else 0.0

    R_hi = float(cfg.get("R_hi", 0.90))
    R_lo = float(cfg.get("R_lo", 0.60))
    disr_hi = float(cfg.get("disr_hi", 0.50))
    disr_lo = float(cfg.get("disr_lo", 0.25))
    if (R >= R_hi) and (disr <= disr_lo):
        state = STATE_STABLE
    elif (R <= R_lo) or (disr >= disr_hi):
        state = STATE_DISSONANT
    else:
        state = STATE_TRANSITION
    return state, float(R), disr


def _smooth_adjust_k(
    kG: float, kL: float, state: str, cfg: Mapping[str, Any]
) -> tuple[float, float]:
    """Smoothly update kG/kL toward targets according to state."""

    kG_min = float(cfg.get("kG_min", 0.01))
    kG_max = float(cfg.get("kG_max", 0.20))
    kL_min = float(cfg.get("kL_min", 0.05))
    kL_max = float(cfg.get("kL_max", 0.25))

    state = normalise_state_token(state)

    if state == STATE_DISSONANT:
        kG_t = kG_max
        kL_t = 0.5 * (
            kL_min + kL_max
        )  # local medio para no perder plasticidad
    elif state == STATE_STABLE:
        kG_t = kG_min
        kL_t = kL_min
    else:
        kG_t = 0.5 * (kG_min + kG_max)
        kL_t = 0.5 * (kL_min + kL_max)

    up = float(cfg.get("up", 0.10))
    down = float(cfg.get("down", 0.07))

    def _step(curr: float, target: float, mn: float, mx: float) -> float:
        gain = up if target > curr else down
        nxt = curr + gain * (target - curr)
        return max(mn, min(mx, nxt))

    return _step(kG, kG_t, kG_min, kG_max), _step(kL, kL_t, kL_min, kL_max)


def _phase_adjust_chunk(args: ChunkArgs) -> list[tuple[NodeId, Phase]]:
    """Return coordinated phase updates for the provided chunk."""

    (
        nodes,
        theta_map,
        cos_map,
        sin_map,
        neighbors_map,
        thG,
        kG,
        kL,
    ) = args
    updates: list[tuple[NodeId, Phase]] = []
    for node in nodes:
        th = float(theta_map.get(node, 0.0))
        neigh = neighbors_map.get(node, ())
        if neigh:
            thL = neighbor_phase_mean_list(
                neigh,
                cos_map,
                sin_map,
                np=None,
                fallback=th,
            )
        else:
            thL = th
        dG = angle_diff(thG, th)
        dL = angle_diff(thL, th)
        updates.append((node, cast(Phase, th + kG * dG + kL * dL)))
    return updates


def coordinate_global_local_phase(
    G: TNFRGraph,
    global_force: float | None = None,
    local_force: float | None = None,
    *,
    n_jobs: int | None = None,
) -> None:
    """Coordinate phase using a blend of global and neighbour coupling."""

    g = cast(dict[str, Any], G.graph)
    hist = cast(dict[str, Any], g.setdefault("history", {}))
    maxlen = int(
        g.get("PHASE_HISTORY_MAXLEN", METRIC_DEFAULTS["PHASE_HISTORY_MAXLEN"])
    )
    hist_state = cast(deque[str], _ensure_hist_deque(hist, "phase_state", maxlen))
    if hist_state:
        normalised_states = [normalise_state_token(item) for item in hist_state]
        if normalised_states != list(hist_state):
            hist_state.clear()
            hist_state.extend(normalised_states)
    hist_R = cast(deque[float], _ensure_hist_deque(hist, "phase_R", maxlen))
    hist_disr = cast(deque[float], _ensure_hist_deque(hist, "phase_disr", maxlen))

    if (global_force is not None) or (local_force is not None):
        kG = float(
            global_force
            if global_force is not None
            else g.get("PHASE_K_GLOBAL", DEFAULTS["PHASE_K_GLOBAL"])
        )
        kL = float(
            local_force
            if local_force is not None
            else g.get("PHASE_K_LOCAL", DEFAULTS["PHASE_K_LOCAL"])
        )
    else:
        cfg, kG, kL = _read_adaptive_params(g)

        if bool(cfg.get("enabled", False)):
            state, R, disr = _compute_state(G, cfg)
            kG, kL = _smooth_adjust_k(kG, kL, state, cfg)

            hist_state.append(state)
            hist_R.append(float(R))
            hist_disr.append(float(disr))

    g["PHASE_K_GLOBAL"] = kG
    g["PHASE_K_LOCAL"] = kL
    append_metric(hist, "phase_kG", float(kG))
    append_metric(hist, "phase_kL", float(kL))

    jobs: int | None
    try:
        jobs = None if n_jobs is None else int(n_jobs)
    except (TypeError, ValueError):
        jobs = None
    if jobs is not None and jobs <= 1:
        jobs = None

    np = get_numpy()
    if np is not None:
        jobs = None

    nodes: list[NodeId] = [cast(NodeId, node) for node in G.nodes()]
    num_nodes = len(nodes)
    if not num_nodes:
        return

    trig = get_trig_cache(G, np=np)
    theta_map = cast(dict[NodeId, Phase], trig.theta)
    cos_map = cast(dict[NodeId, float], trig.cos)
    sin_map = cast(dict[NodeId, float], trig.sin)

    neighbors_proxy = ensure_neighbors_map(G)
    neighbors_map: dict[NodeId, tuple[NodeId, ...]] = {}
    for n in nodes:
        try:
            neighbors_map[n] = tuple(cast(Sequence[NodeId], neighbors_proxy[n]))
        except KeyError:
            neighbors_map[n] = ()

    def _theta_value(node: NodeId) -> float:
        cached = theta_map.get(node)
        if cached is not None:
            return float(cached)
        attr_val = get_theta_attr(G.nodes[node], 0.0)
        return float(attr_val if attr_val is not None else 0.0)

    theta_vals = [_theta_value(n) for n in nodes]
    cos_vals = [
        float(cos_map.get(n, math.cos(theta_vals[idx])))
        for idx, n in enumerate(nodes)
    ]
    sin_vals = [
        float(sin_map.get(n, math.sin(theta_vals[idx])))
        for idx, n in enumerate(nodes)
    ]

    if np is not None:
        theta_arr = cast(FloatArray, np.fromiter(theta_vals, dtype=float))
        cos_arr = cast(FloatArray, np.fromiter(cos_vals, dtype=float))
        sin_arr = cast(FloatArray, np.fromiter(sin_vals, dtype=float))
        if cos_arr.size:
            mean_cos = float(np.mean(cos_arr))
            mean_sin = float(np.mean(sin_arr))
            thG = float(np.arctan2(mean_sin, mean_cos))
        else:
            thG = 0.0
        neighbor_means = [
            neighbor_phase_mean_list(
                neighbors_map.get(n, ()),
                cos_map,
                sin_map,
                np=np,
                fallback=theta_vals[idx],
            )
            for idx, n in enumerate(nodes)
        ]
        neighbor_arr = cast(FloatArray, np.fromiter(neighbor_means, dtype=float))
        theta_updates = theta_arr + kG * (thG - theta_arr) + kL * (
            neighbor_arr - theta_arr
        )
        for idx, node in enumerate(nodes):
            set_theta(G, node, float(theta_updates[int(idx)]))
        return

    mean_cos = math.fsum(cos_vals) / num_nodes
    mean_sin = math.fsum(sin_vals) / num_nodes
    thG = math.atan2(mean_sin, mean_cos)

    if jobs is None:
        for node in nodes:
            th = float(theta_map.get(node, 0.0))
            neigh = neighbors_map.get(node, ())
            if neigh:
                thL = neighbor_phase_mean_list(
                    neigh,
                    cos_map,
                    sin_map,
                    np=None,
                    fallback=th,
                )
            else:
                thL = th
            dG = angle_diff(thG, th)
            dL = angle_diff(thL, th)
            set_theta(G, node, float(th + kG * dG + kL * dL))
        return

    chunk_size = max(1, math.ceil(len(nodes) / jobs))
    chunks = [
        nodes[idx : idx + chunk_size]
        for idx in range(0, len(nodes), chunk_size)
    ]
    args: list[ChunkArgs] = [
        (
            chunk,
            theta_map,
            cos_map,
            sin_map,
            neighbors_map,
            thG,
            kG,
            kL,
        )
        for chunk in chunks
    ]
    results: dict[NodeId, Phase] = {}
    with ProcessPoolExecutor(max_workers=jobs) as executor:
        for res in executor.map(_phase_adjust_chunk, args):
            for node, value in res:
                results[node] = value
    for node in nodes:
        new_theta = results.get(node)
        base_theta = theta_map.get(node, 0.0)
        set_theta(G, node, float(new_theta if new_theta is not None else base_theta))

