package dev.lumareth.improvedAerodynamics.flight;

import org.joml.Matrix3d;
import org.joml.Matrix3dc;
import org.joml.Quaterniondc;
import org.joml.Vector3d;

/**
 * The virtual pilot of cursor aim. The cursor is a stick that springs back to the gun reticle:
 * <ul>
 *   <li>sideways it commands a bank angle and turn: the further out, the steeper the bank and the harder the pull
 *   that comes with it. Centered it levels the wings, so a centered cursor is straight flight;</li>
 *   <li>up and down it commands a pitch rate;</li>
 *   <li>close to the reticle it steers with yaw and pitch directly instead of banking, so aiming does not rock the
 *   wings.</li>
 * </ul>
 * The wanted rotation is turned into stick inputs by the measured effect of each input on the craft (torque per unit
 * of input, see {@link AeroState}), so it works for any craft however its surfaces are mixed or oriented: angular
 * acceleration wanted times inertia, solved for the three inputs with a little regularization for axes that have no
 * surfaces. Keys the pilot holds win over it, per axis.
 */
final class CursorPilot {
    private static final double FINE_START = Math.toRadians(4.0);
    private static final double FINE_END = Math.toRadians(16.0);
    private static final double MANUAL_OVERRIDE = 0.05;
    private static final double MAX_PULL_BANK = Math.toRadians(90.0);

    private CursorPilot() {
    }

    /**
     * @return the node with the pilot's wanted inputs in place of the keys, or the node itself if cursor aim is off
     */
    static NodeData apply(AeroState.NodeState state, NodeData node, Quaterniondc rotation, Matrix3dc inertia, double[] w,
                          double speed, boolean wingStalled) {
        if (!CursorAimConfig.ENABLED.get()) {
            return node;
        }
        CursorAimRegistry.Aim aim = CursorAimRegistry.get(node.key());
        if (aim == null || speed < CursorAimConfig.MIN_SPEED.get()) {
            return node;
        }

        double fullAngle = Math.toRadians(CursorAimConfig.FULL_DEFLECTION_DEGREES.get());
        double maxTurn = Math.toRadians(CursorAimConfig.MAX_TURN_RATE_DEGREES.get());
        double maxRoll = Math.toRadians(CursorAimConfig.MAX_ROLL_RATE_DEGREES.get());
        double maxBank = Math.toRadians(CursorAimConfig.MAX_BANK_DEGREES.get());

        double[] f = aim.forward();
        double[] r = aim.right();
        double[] u = aim.up();

        // Cursor as angles from the reticle: sideways (positive right) and up, both through any 270 degree cone
        double sideways = Math.atan2(aim.cr(), aim.cf());
        double vertical = Math.atan2(aim.cu(), Math.hypot(aim.cf(), aim.cr()));
        // Deadzone around the reticle: shrink the cursor offset by its radius, direction kept
        double offset = Math.hypot(sideways, vertical);
        double deadzone = Math.toRadians(CursorAimConfig.DEADZONE_DEGREES.get());
        if (offset > 1.0E-9) {
            double scale = Math.max(0.0, offset - deadzone) / offset;
            sideways *= scale;
            vertical *= scale;
        }
        double lateral = saturate(sideways / fullAngle);
        double blend = smoothstep(Math.abs(sideways), FINE_START, FINE_END);

        // Bank of the craft: positive when the right wing is down, through upside down
        Vector3d rightWorld = rotation.transform(new Vector3d(r[0], r[1], r[2]));
        Vector3d upWorld = rotation.transform(new Vector3d(u[0], u[1], u[2]));
        double bank = Math.atan2(-rightWorld.y, upWorld.y);

        double bankCommand = maxBank * lateral * blend;
        double rollRate = clamp(CursorAimConfig.BANK_GAIN.get() * (bankCommand - bank), -maxRoll, maxRoll);

        // Pitch: the stick, plus the pull that carries the turn while banked
        double pitchStick = saturate(vertical / fullAngle);
        double pitchRate = maxTurn * (pitchStick >= 0.0 ? pitchStick : pitchStick * CursorAimConfig.PUSH_FACTOR.get());
        if (Math.abs(bank) < MAX_PULL_BANK) {
            pitchRate += maxTurn * Math.abs(lateral) * blend * Math.sin(Math.abs(bank));
        }
        if (wingStalled && pitchRate > 0.0) {
            pitchRate *= CursorAimConfig.STALLED_PITCH_FACTOR.get();
        }
        // Close to the reticle the nose follows the cursor with the rudder
        double yawRate = maxTurn * lateral * (1.0 - blend);

        // Wanted angular velocity in ship axes: roll about forward, pitch about right, yaw (nose right) about -up
        Vector3d wanted = new Vector3d(
                rollRate * f[0] + pitchRate * r[0] - yawRate * u[0],
                rollRate * f[1] + pitchRate * r[1] - yawRate * u[1],
                rollRate * f[2] + pitchRate * r[2] - yawRate * u[2]);
        Vector3d error = wanted.sub(w[0], w[1], w[2]);
        Vector3d acceleration = error.mul(CursorAimConfig.RATE_GAIN.get());
        Vector3d torque = inertia.transform(acceleration, new Vector3d());

        double[] inputs = solve(state.pitchResponse, state.yawResponse, state.rollResponse, torque);
        if (inputs == null) {
            return node;
        }

        double pitch = clamp(inputs[0] - state.trimOutput, -1.0, 1.0);
        double yaw = clamp(inputs[1], -1.0, 1.0);
        double roll = clamp(inputs[2], -1.0, 1.0);
        return new NodeData(node.key(),
                Math.abs(node.pitch()) > MANUAL_OVERRIDE ? node.pitch() : pitch,
                Math.abs(node.yaw()) > MANUAL_OVERRIDE ? node.yaw() : yaw,
                Math.abs(node.roll()) > MANUAL_OVERRIDE ? node.roll() : roll,
                node.throttle(), node.autoTrim(), node.storedTrim());
    }

    /** Least squares for the three inputs whose responses add up to the torque, ridge-regularized so unused axes stay at zero. */
    private static double[] solve(Vector3d pitch, Vector3d yaw, Vector3d roll, Vector3d torque) {
        Matrix3d responses = new Matrix3d(pitch, yaw, roll);
        Matrix3d transposed = responses.transpose(new Matrix3d());
        Matrix3d normal = transposed.mul(responses, new Matrix3d());
        double trace = normal.m00() + normal.m11() + normal.m22();
        if (!(trace > 1.0E-9) || !Double.isFinite(trace)) {
            return null;
        }
        double ridge = trace / 3.0 * 1.0E-3;
        normal.m00(normal.m00() + ridge);
        normal.m11(normal.m11() + ridge);
        normal.m22(normal.m22() + ridge);
        Vector3d right = transposed.transform(torque, new Vector3d());
        Vector3d x = normal.invert().transform(right);
        if (!x.isFinite()) {
            return null;
        }
        return new double[]{x.x, x.y, x.z};
    }

    private static double saturate(double value) {
        return clamp(value, -1.0, 1.0);
    }

    private static double smoothstep(double value, double from, double to) {
        double t = clamp((value - from) / (to - from), 0.0, 1.0);
        return t * t * (3.0 - 2.0 * t);
    }

    private static double clamp(double value, double min, double max) {
        return Math.max(min, Math.min(max, value));
    }
}