package dev.lumareth.improvedAerodynamics.flight;

import dev.lumareth.improvedAerodynamics.Config;
import it.unimi.dsi.fastutil.longs.Long2BooleanOpenHashMap;
import it.unimi.dsi.fastutil.longs.Long2ObjectOpenHashMap;
import org.joml.Matrix3dc;
import org.joml.Quaterniondc;
import org.joml.Vector3d;
import org.joml.Vector3dc;
import org.valkyrienskies.core.api.ships.PhysShip;
import org.valkyrienskies.core.api.ships.properties.ShipTransform;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * Aerodynamics of one ship's Warium surfaces, evaluated every physics substep. The server thread registers the
 * surfaces whose procedures ran (and so passed their conditions) during a tick, plus the control nodes they listen
 * to; at the end of the tick that set is handed to the physics thread, which computes the forces from the ship's
 * motion in each substep until the next tick arrives.
 * <p>
 * The physics side also smooths pilot input, runs auto trim, and integrates damping implicitly so stiff damping
 * cannot overshoot.
 */
public final class AeroState implements PhysicsBridge.Publishable {
    // Warium's seat moves inputs by 0.2 per tick when pressing a key, smoothing uses that for releasing too
    private static final double INPUT_RATE_PER_TICK = 0.2;
    private static final double VELOCITY_EPSILON = 0.05;
    private static final double ANGULAR_EPSILON = 0.005;

    private final SubstepClock clock;

    // Server thread
    private Long2ObjectOpenHashMap<AeroSurface> buildingSurfaces = new Long2ObjectOpenHashMap<>();
    private Long2ObjectOpenHashMap<NodeData> buildingNodes = new Long2ObjectOpenHashMap<>();

    private volatile Frame published;

    // Physics thread, except that the server thread reads trim values from nodeStates to save them
    private Frame current;
    private boolean[] stalled = new boolean[0];
    private boolean wingStalled;
    private final Map<Long, NodeState> nodeStates = new ConcurrentHashMap<>();
    private final Wrench wrench = new Wrench();
    private final Wrench probe = new Wrench();

    private record Frame(long tick, AeroSurface[] surfaces, Long2ObjectOpenHashMap<NodeData> nodes) {
    }

    /** Smoothed inputs and trim of a control node. */
    static final class NodeState {
        double pitch;
        double yaw;
        double roll;
        double trimIntegral;
        double trimOutput;
        // Torque (ship axes) one unit of pitch input on this node's pitch surfaces makes, measured every substep
        final Vector3d pitchResponse = new Vector3d();
        final Vector3d pitchResponseBuilding = new Vector3d();
        volatile double savedTrim;

        NodeState(NodeData data) {
            pitch = data.pitch();
            yaw = data.yaw();
            roll = data.roll();
            trimIntegral = data.autoTrim() ? data.storedTrim() : 0.0;
            trimOutput = trimIntegral;
            savedTrim = trimIntegral;
        }
    }

    public AeroState(SubstepClock clock) {
        this.clock = clock;
    }

    // ---- Server thread ----

    public boolean has(long surfaceKey) {
        return buildingSurfaces.containsKey(surfaceKey);
    }

    public boolean hasNode(long nodeKey) {
        return buildingNodes.containsKey(nodeKey);
    }

    public void add(AeroSurface surface) {
        if (buildingSurfaces.isEmpty()) {
            PhysicsBridge.touch(this);
        }
        buildingSurfaces.put(surface.key(), surface);
    }

    public void addNode(NodeData node) {
        buildingNodes.put(node.key(), node);
    }

    /** Trim auto trim currently holds on a node, NaN if the physics thread has not seen the node. */
    public double trim(long nodeKey) {
        NodeState state = nodeStates.get(nodeKey);
        return state != null ? state.savedTrim : Double.NaN;
    }

    @Override
    public void publish(long tick) {
        published = new Frame(tick, buildingSurfaces.values().toArray(new AeroSurface[0]), buildingNodes);
        buildingSurfaces = new Long2ObjectOpenHashMap<>();
        buildingNodes = new Long2ObjectOpenHashMap<>();
    }

    // ---- Physics thread ----

    /** Applies this substep's aerodynamic forces. Call after {@link SubstepClock#step()}. */
    public void apply(PhysShip ship, boolean autoScale) {
        Frame frame = published;
        if (frame == null || frame.tick < PhysicsBridge.publishedTick() || clock.stale()) {
            return;
        }
        if (frame != current) {
            switchFrame(frame);
        }

        ShipTransform transform = ship.getTransform();
        Quaterniondc rotation = transform.getShipToWorldRotation();
        Vector3dc centerOfMass = transform.getPositionInShip();
        Vector3dc worldVelocity = ship.getVelocity();
        Vector3d velocity = rotation.transformInverse(worldVelocity, new Vector3d());
        Vector3d omega = rotation.transformInverse(ship.getOmega(), new Vector3d());
        double[] v = {velocity.x, velocity.y, velocity.z};
        double[] w = {omega.x, omega.y, omega.z};
        double dt = clock.substepSeconds();
        boolean improved = Config.IMPROVED_FLIGHT_MODEL.get();
        Matrix3dc inertia = ship.getMomentOfInertia();

        updateNodes(frame, worldVelocity, rotation, inertia, w, dt, improved);
        evaluate(frame, centerOfMass, v, w, improved, true, wrench);

        if (Config.IMPLICIT_DAMPING.get()) {
            // Implicit Euler for the damping part: scale each axis by 1 / (1 + dt * damping / inertia), using the
            // diagonal of the force's derivative with respect to velocity and angular velocity
            double[] diagonal = {inertia.m00(), inertia.m11(), inertia.m22()};
            for (int axis = 0; axis < 3; axis++) {
                v[axis] += VELOCITY_EPSILON;
                evaluate(frame, centerOfMass, v, w, improved, false, probe);
                v[axis] -= VELOCITY_EPSILON;
                double damping = -(probe.force[axis] - wrench.force[axis]) / VELOCITY_EPSILON;
                if (damping > 0.0) {
                    wrench.force[axis] /= 1.0 + dt * damping / ship.getMass();
                }

                w[axis] += ANGULAR_EPSILON;
                evaluate(frame, centerOfMass, v, w, improved, false, probe);
                w[axis] -= ANGULAR_EPSILON;
                double angularDamping = -(probe.totalTorque(axis) - wrench.totalTorque(axis)) / ANGULAR_EPSILON;
                if (angularDamping > 0.0 && diagonal[axis] > 0.0) {
                    double factor = 1.0 / (1.0 + dt * angularDamping / diagonal[axis]);
                    wrench.forceTorque[axis] *= factor;
                    wrench.torque[axis] *= factor;
                }
            }
        }

        // Same frames and scaling as Warium applies its forces in
        double scale = transform.getShipToWorldScaling().x();
        double forceScale = autoScale ? scale * scale * scale : 1.0;
        Vector3d force = new Vector3d(wrench.force[0], wrench.force[1], wrench.force[2]).mul(forceScale);
        Vector3d forceTorque = new Vector3d(wrench.forceTorque[0], wrench.forceTorque[1], wrench.forceTorque[2]).mul(forceScale);
        if (force.isFinite() && forceTorque.isFinite()) {
            ship.applyInvariantForce(rotation.transform(force));
            ship.applyInvariantTorque(rotation.transform(forceTorque));
        }
        Vector3d torque = new Vector3d(wrench.torque[0], wrench.torque[1], wrench.torque[2]);
        if (torque.lengthSquared() > 0.0 && torque.isFinite()) {
            ship.applyRotDependentTorque(torque);
        }
    }

    private void switchFrame(Frame frame) {
        // Carry stall state over by block, the surface order changes from tick to tick
        Long2BooleanOpenHashMap previous = new Long2BooleanOpenHashMap();
        if (current != null) {
            for (int i = 0; i < current.surfaces.length; i++) {
                if (stalled[i]) {
                    previous.put(current.surfaces[i].key(), true);
                }
            }
        }
        stalled = new boolean[frame.surfaces.length];
        for (int i = 0; i < stalled.length; i++) {
            stalled[i] = previous.get(frame.surfaces[i].key());
        }
        nodeStates.keySet().removeIf(key -> !frame.nodes.containsKey((long) key));
        for (NodeData node : frame.nodes.values()) {
            nodeStates.computeIfAbsent(node.key(), key -> new NodeState(node));
        }
        current = frame;
    }

    /**
     * Sums the forces of all surfaces into out.
     *
     * @param commit whether this is the real evaluation of the substep (updates stall state and trim response)
     *               rather than a probe for the damping derivatives
     */
    private void evaluate(Frame frame, Vector3dc centerOfMass, double[] v, double[] w, boolean improved, boolean commit, Wrench out) {
        out.clear();
        double weight = clock.weight();
        double liftConstant = ImprovedAero.liftConstant(Config.DESIGN_SPEED.get());
        double stallAngleScale = Config.STALL_ANGLE_SCALE.get();
        double[] r = new double[3];
        if (commit) {
            wingStalled = false;
            for (NodeState state : nodeStates.values()) {
                state.pitchResponseBuilding.zero();
            }
        }

        for (int i = 0; i < frame.surfaces.length; i++) {
            AeroSurface s = frame.surfaces[i];
            r[0] = s.x() - centerOfMass.x();
            r[1] = s.y() - centerOfMass.y();
            r[2] = s.z() - centerOfMass.z();
            NodeData node = s.node() != AeroSurface.NO_NODE ? frame.nodes.get(s.node()) : null;
            NodeState state = node != null ? nodeStates.get(node.key()) : null;
            double input = s.type().controllable ? input(s, node, state, 0.0) : 0.0;

            if (improved) {
                boolean nowStalled = ImprovedAero.evaluate(s, v, w, r, input, weight, liftConstant, stallAngleScale, stalled[i], out);
                if (commit) {
                    stalled[i] = nowStalled;
                    wingStalled |= nowStalled && s.type() == SurfaceType.WING;
                }
            } else {
                WariumAero.evaluate(s, v, w, r, input, weight, out);
            }

            // How much a little more pitch input on this node would change the torque, for auto trim
            if (commit && node != null && node.autoTrim() && s.type().controllable && s.listensToPitch()) {
                double delta = 0.05;
                probe.clear();
                double base = 0.0;
                if (improved) {
                    ImprovedAero.evaluate(s, v, w, r, input, weight, liftConstant, stallAngleScale, stalled[i], probe);
                } else {
                    WariumAero.evaluate(s, v, w, r, input, weight, probe);
                }
                double[] before = {probe.totalTorque(0), probe.totalTorque(1), probe.totalTorque(2)};
                probe.clear();
                double nudged = input(s, node, state, delta);
                if (improved) {
                    ImprovedAero.evaluate(s, v, w, r, nudged, weight, liftConstant, stallAngleScale, stalled[i], probe);
                } else {
                    WariumAero.evaluate(s, v, w, r, nudged, weight, probe);
                }
                state.pitchResponseBuilding.add(
                        (probe.totalTorque(0) - before[0]) / delta,
                        (probe.totalTorque(1) - before[1]) / delta,
                        (probe.totalTorque(2) - before[2]) / delta);
            }
        }

        if (commit) {
            for (NodeState state : nodeStates.values()) {
                state.pitchResponse.set(state.pitchResponseBuilding);
            }
        }
    }

    /** A control surface's input, combined from its two keys and clamped like Warium does. */
    private static double input(AeroSurface s, NodeData node, NodeState state, double extraPitch) {
        if (node == null || state == null) {
            return 0.0;
        }
        double value = axisValue(s.axis1(), node, state, extraPitch) * s.multiplier1()
                + axisValue(s.axis2(), node, state, extraPitch) * s.multiplier2();
        return Math.max(-1.0, Math.min(1.0, value));
    }

    private static double axisValue(ControlAxis axis, NodeData node, NodeState state, double extraPitch) {
        return switch (axis) {
            case PITCH -> state.pitch + state.trimOutput + extraPitch;
            case YAW -> state.yaw;
            case ROLL -> state.roll;
            case THROTTLE -> node.throttle();
            case NONE -> 0.0;
        };
    }

    private void updateNodes(Frame frame, Vector3dc worldVelocity, Quaterniondc rotation, Matrix3dc inertia,
                             double[] w, double dt, boolean improved) {
        double step = Config.CONTROL_SMOOTHING.get() ? INPUT_RATE_PER_TICK * clock.weight() : Double.POSITIVE_INFINITY;
        for (NodeData node : frame.nodes.values()) {
            NodeState state = nodeStates.get(node.key());
            if (state == null) {
                continue;
            }
            state.pitch = moveTowards(state.pitch, node.pitch(), step);
            state.yaw = moveTowards(state.yaw, node.yaw(), step);
            state.roll = moveTowards(state.roll, node.roll(), step);
            AutoTrim.update(state, node, worldVelocity, rotation, inertia, w, dt, improved && wingStalled);
        }
    }

    static double moveTowards(double value, double target, double maxStep) {
        if (Math.abs(target - value) <= maxStep) {
            return target;
        }
        return value + Math.copySign(maxStep, target - value);
    }
}