"""Outer (heading-control) PID -- reference Python implementation.

Inputs:
    heading_setpoint_deg : desired compass heading (0..360)
    heading_measured_deg : current compass heading (0..360)
    rate_of_turn_dps     : current ROT in deg/s (positive = turning stbd)
    speed_kn             : current SOG in knots, used for gain scheduling
    allowed              : if False, output is zero and integrator bleeds

Output:
    rudder_setpoint_deg  : signed rudder angle in degrees, clamped to the
                            mechanical limit. Passed to the inner PID.

Algorithm:
    err            = shortest_arc(setpoint, measured)
    deadband applied: |err| < deadband_deg -> 0
    p_term         = kp(speed) * err
    integral      += ki(speed) * err * dt   (with anti-windup)
    derivative     = kd(speed) * (err - prev_err) / dt (LPF)
    rot_ff_term    = rot_ff_gain * (-rot)    # subtract ROT to anticipate
    raw_output     = p_term + integral + derivative + rot_ff_term
    rudder_setpoint = saturate(raw_output, max_rudder_deg)
    if saturated, back-calculate the integrator

The ROT feed-forward (brief sec. 6) is the term that distinguishes a
premium autopilot from a basic one: it subtracts the current rate of
turn from the rudder command, so as the vessel approaches the new
heading the rudder eases off before the heading is reached, preventing
overshoot due to yaw inertia.
"""

from __future__ import annotations

from dataclasses import dataclass, field

from arautopilot.studio.simulator.vessel_heading import heading_error_deg


@dataclass
class GainPoint:
    """One interpolation point for the SOG-indexed gain schedule."""

    speed_kn: float
    kp: float
    ki: float
    kd: float


@dataclass
class PidOuterConfig:
    # ----- Sampling -----------------------------------------------------
    freq_hz: float = 10.0

    # ----- Base gains (used if gain_schedule is empty) -----------------
    base_kp: float = 0.9
    base_ki: float = 0.02
    base_kd: float = 1.2

    # ----- SOG-indexed gain schedule (outer loop only) -----------------
    gain_schedule: list[GainPoint] = field(default_factory=list)

    # ----- Setpoint handling -------------------------------------------
    deadband_deg: float = 0.5
    integral_clamp: float = 35.0   # output-equivalent units (deg of rudder)

    # ----- Output saturation -------------------------------------------
    max_rudder_deg: float = 35.0
    rate_limit_dps: float = 30.0   # rate-limit on the produced rudder setpoint

    # ----- Anti-windup --------------------------------------------------
    aw_gain: float | None = None   # default: 1 / current kp

    # ----- Derivative low-pass -----------------------------------------
    d_lpf_tau_s: float = 0.10

    # ----- Rate-of-Turn feed-forward (brief sec. 6) --------------------
    rot_ff_gain: float = 1.5

    def dt(self) -> float:
        return 1.0 / self.freq_hz


@dataclass
class PidOuterState:
    integral: float = 0.0
    prev_error: float = 0.0
    prev_d_term: float = 0.0
    prev_rudder_setpoint: float = 0.0
    last_output_deg: float = 0.0


def interpolate_gains(
    schedule: list[GainPoint], speed_kn: float, base_kp: float, base_ki: float,
    base_kd: float,
) -> tuple[float, float, float]:
    """Return (kp, ki, kd) interpolated linearly at ``speed_kn``.

    Outside the range of the schedule the closest endpoint is held
    (no extrapolation). If the schedule is empty, the base gains are
    returned unchanged.
    """
    if not schedule:
        return base_kp, base_ki, base_kd
    points = sorted(schedule, key=lambda p: p.speed_kn)
    if speed_kn <= points[0].speed_kn:
        return points[0].kp, points[0].ki, points[0].kd
    if speed_kn >= points[-1].speed_kn:
        return points[-1].kp, points[-1].ki, points[-1].kd
    for i in range(len(points) - 1):
        a, b = points[i], points[i + 1]
        if a.speed_kn <= speed_kn <= b.speed_kn:
            span = b.speed_kn - a.speed_kn
            t = 0.0 if span == 0.0 else (speed_kn - a.speed_kn) / span
            return (
                a.kp + t * (b.kp - a.kp),
                a.ki + t * (b.ki - a.ki),
                a.kd + t * (b.kd - a.kd),
            )
    raise RuntimeError("Gain schedule interpolation failed unexpectedly")


class PidOuter:
    def __init__(self, config: PidOuterConfig | None = None) -> None:
        self.config: PidOuterConfig = config or PidOuterConfig()
        self.state: PidOuterState = PidOuterState()
        self._dt: float = self.config.dt()

    def reset(self, *, heading_deg: float = 0.0) -> None:
        self.state = PidOuterState()

    def update_config(self, config: PidOuterConfig) -> None:
        self.config = config
        self._dt = config.dt()

    def step(
        self,
        *,
        heading_setpoint_deg: float,
        heading_measured_deg: float,
        rate_of_turn_dps: float,
        speed_kn: float,
        allowed: bool = True,
    ) -> float:
        cfg = self.config
        st = self.state
        dt = self._dt

        # Error with shortest-arc semantics.
        raw_err = heading_error_deg(heading_setpoint_deg, heading_measured_deg)
        if abs(raw_err) <= cfg.deadband_deg:
            error = 0.0
        else:
            sign = 1.0 if raw_err > 0.0 else -1.0
            error = raw_err - sign * cfg.deadband_deg

        if not allowed:
            st.integral *= 0.95
            st.prev_error = error
            st.last_output_deg = 0.0
            return 0.0

        # Gain scheduling: pull gains for the current speed.
        kp, ki, kd = interpolate_gains(
            cfg.gain_schedule, speed_kn, cfg.base_kp, cfg.base_ki, cfg.base_kd
        )

        # PID terms.
        p_term = kp * error
        st.integral += ki * error * dt
        st.integral = _clamp(st.integral, -cfg.integral_clamp, cfg.integral_clamp)

        d_raw = kd * (error - st.prev_error) / dt if dt > 0.0 else 0.0
        alpha = _lpf_alpha(cfg.d_lpf_tau_s, dt)
        d_term = (1.0 - alpha) * st.prev_d_term + alpha * d_raw
        st.prev_d_term = d_term

        # Rate-of-Turn feed-forward: subtract the current ROT to anticipate
        # the vessel coasting through the setpoint. Premium autopilot
        # behaviour per brief sec. 6.
        rot_ff_term = -cfg.rot_ff_gain * rate_of_turn_dps

        raw_rudder = p_term + st.integral + d_term + rot_ff_term

        # Saturate to mechanical rudder limit + back-calculation anti-windup.
        rudder = _clamp(raw_rudder, -cfg.max_rudder_deg, cfg.max_rudder_deg)
        if raw_rudder != rudder:
            aw = cfg.aw_gain if cfg.aw_gain is not None else (
                1.0 / kp if kp != 0.0 else 0.0
            )
            st.integral -= aw * (raw_rudder - rudder) * dt
            st.integral = _clamp(st.integral, -cfg.integral_clamp, cfg.integral_clamp)

        # Rate-limit the rudder setpoint so the inner loop has a smooth
        # reference. (The inner loop also has a rate limit; this one keeps
        # the setpoint that the outer loop hands off bounded in delta.)
        max_delta = cfg.rate_limit_dps * dt
        delta = rudder - st.prev_rudder_setpoint
        if delta > max_delta:
            rudder = st.prev_rudder_setpoint + max_delta
        elif delta < -max_delta:
            rudder = st.prev_rudder_setpoint - max_delta
        st.prev_rudder_setpoint = rudder

        st.prev_error = error
        st.last_output_deg = rudder
        return rudder


def _clamp(x: float, lo: float, hi: float) -> float:
    if x < lo:
        return lo
    if x > hi:
        return hi
    return x


def _lpf_alpha(tau_s: float, dt: float) -> float:
    if tau_s <= 0.0:
        return 1.0
    return dt / (tau_s + dt)