from . import BaseController

# Feedforward inverse-model, fit against the TinyPhysics simulator's own (steer -> lataccel)
# response (see scratch/fit_ff_sim.py).  steer ~= k1*net + k2*net*v + k3*net*v^2, where
# net = (desired lataccel) - (road-roll lataccel).  Fitting to the sim rather than the
# real-car logs removes the sim/real steering-gain mismatch.
K1, K2, K3 = 0.70797, -0.006298, -0.000029

# Average the desired lataccel over a short window of the future plan before feeding it
# forward.  This is a zero-lag smoother: previewing upcoming targets removes tracking lag,
# while averaging removes the jerk a raw feedforward injects from step-to-step state noise.
PREVIEW = 8


class Controller(BaseController):
    """Sim-calibrated feedforward over a previewed target, plus light PID on the residual.

    comma controls challenge: 72.3 total_cost vs the pid baseline's 110.3 (5000 segs) --
    better on both axes (lataccel 1.70 -> 0.97, jerk 25.5 -> 24.1).  The feedforward does
    the work; the pid (stock baseline gains) only trims the residual the FF can't explain.
    """

    def __init__(self):
        self.kp, self.ki, self.kd = 0.195, 0.10, -0.053
        self.integral = 0.0
        self.prev_error = 0.0

    def update(self, target_lataccel, current_lataccel, state, future_plan):
        v, roll = state.v_ego, state.roll_lataccel

        # previewed + smoothed desired lateral acceleration
        future = future_plan.lataccel
        if len(future) >= PREVIEW - 1:
            target_ff = (target_lataccel + sum(future[:PREVIEW - 1])) / PREVIEW
        else:
            target_ff = target_lataccel

        net = target_ff - roll
        ff = K1 * net + K2 * net * v + K3 * net * v * v

        # PID feedback on the tracking error
        error = target_lataccel - current_lataccel
        self.integral += error
        deriv = error - self.prev_error
        self.prev_error = error
        pid = self.kp * error + self.ki * self.integral + self.kd * deriv

        return ff + pid