// =============================================================================
// kalman_2x2.sv : 2-state Kalman filter for one axis (roll OR pitch)
// -----------------------------------------------------------------------------
// State X = [angle ; gyro_rate]
// A = [1 dt ; 0 1]   H = I   Q = diag(Q00,Q11)   R = diag(R00,R11)
//
// I/O is Q7.9 (16-bit signed); internal storage is Q7.25 (32-bit signed).
// Q7.9 -> Q7.25 : sign-extend, shift left 16
// Q7.25 -> Q7.9 : arithmetic right shift 16, with saturation to 16 bits.
// =============================================================================
`default_nettype none

module kalman_2x2 #(
    // Process model -- override at instantiation if you want to retune.
    parameter signed [31:0] DT_Q725    = 32'sd33554,     // 0.001 s
    parameter signed [31:0] Q00_Q725   = 32'sd33554,     // 0.001
    parameter signed [31:0] Q11_Q725   = 32'sd100663,    // 0.003
    parameter signed [31:0] R00_Q725   = 32'sd1006633,   // 0.03
    parameter signed [31:0] R11_Q725   = 32'sd1006633,   // 0.03
    parameter signed [31:0] P0_00_Q725 = 32'sd33554432,  // 1.0
    parameter signed [31:0] P0_11_Q725 = 32'sd33554432   // 1.0
) (
    input  wire               clk,
    input  wire               rst_n,
    input  wire               din_valid,
    input  wire signed [15:0] ang_meas_q9,    // Q7.9 measured angle
    input  wire signed [15:0] gyro_meas_q9,   // Q7.9 measured gyro
    output reg  signed [15:0] est_angle_q9,
    output reg                dout_valid,
    output wire               busy
);

    // ---- Q7.9 <-> Q7.25 conversion ---------------------------------------
    // Q7.9 -> Q7.25: sign-extend to 32, then shift left 16 (= 25-9)
    function automatic signed [31:0] q9_to_q25 (input signed [15:0] x);
        q9_to_q25 = {{16{x[15]}}, x} <<< 16;
    endfunction

    // Q7.25 -> Q7.9: arithmetic right shift by 16, saturate to int16 range.
    function automatic signed [15:0] q25_to_q9 (input signed [31:0] x);
        reg signed [31:0] sh;
        begin
            sh = x >>> 16;
            if (sh > 32'sh00007FFF)
                q25_to_q9 = 16'sh7FFF;
            else if (sh < 32'shFFFF8000)
                q25_to_q9 = 16'sh8000;
            else
                q25_to_q9 = sh[15:0];
        end
    endfunction

    // ---- State storage ----------------------------------------------------
    reg signed [31:0] X0, X1;                       // X_prev
    reg signed [31:0] P00, P01, P10, P11;           // P_prev

    reg signed [31:0] Xkp0, Xkp1;                   // predicted state
    reg signed [31:0] Pkp00, Pkp01, Pkp10, Pkp11;
    reg signed [31:0] S00, S01, S10, S11;           // innovation cov
    reg signed [31:0] Sinv00, Sinv01, Sinv10, Sinv11;
    reg signed [31:0] K00, K01, K10, K11;           // Kalman gain
    reg signed [31:0] Y0, Y1;                       // measurement
    reg signed [31:0] innov0, innov1;               // residual
    reg signed [31:0] Knew0, Knew1;                 // K * innov
    reg signed [31:0] IK00, IK01, IK10, IK11;       // I - K
    reg signed [31:0] Pnew00, Pnew01, Pnew10, Pnew11;

    // ---- Shared 2x2 matrix multiplier -------------------------------------
    reg  signed [31:0] mm_a00, mm_a01, mm_a10, mm_a11;
    reg  signed [31:0] mm_b00, mm_b01, mm_b10, mm_b11;
    wire signed [31:0] mm_p00, mm_p01, mm_p10, mm_p11;

    mat2x2_mult u_mm (
        .m00(mm_a00), .m01(mm_a01), .m10(mm_a10), .m11(mm_a11),
        .n00(mm_b00), .n01(mm_b01), .n10(mm_b10), .n11(mm_b11),
        .p00(mm_p00), .p01(mm_p01), .p10(mm_p10), .p11(mm_p11)
    );

    // ---- Shared 2x2 add/sub -----------------------------------------------
    reg                ms_sub;
    reg  signed [31:0] ms_a00, ms_a01, ms_a10, ms_a11;
    reg  signed [31:0] ms_b00, ms_b01, ms_b10, ms_b11;
    wire signed [31:0] ms_p00, ms_p01, ms_p10, ms_p11;

    mat2x2_addsub u_ms (
        .sub(ms_sub),
        .m00(ms_a00), .m01(ms_a01), .m10(ms_a10), .m11(ms_a11),
        .n00(ms_b00), .n01(ms_b01), .n10(ms_b10), .n11(ms_b11),
        .p00(ms_p00), .p01(ms_p01), .p10(ms_p10), .p11(ms_p11)
    );

    // ---- 2x2 inverse engine -----------------------------------------------
    reg                inv_start;
    reg  signed [31:0] inv_in00, inv_in01, inv_in10, inv_in11;
    wire signed [31:0] inv_out00, inv_out01, inv_out10, inv_out11;
    wire               inv_done;

    mat2x2_inv u_inv (
        .clk(clk), .rst_n(rst_n),
        .start(inv_start),
        .m00(inv_in00), .m01(inv_in01), .m10(inv_in10), .m11(inv_in11),
        .inv00(inv_out00), .inv01(inv_out01), .inv10(inv_out10), .inv11(inv_out11),
        .done(inv_done)
    );

    // ---- A and A^T from DT ------------------------------------------------
    localparam signed [31:0] ONE_Q725 = 32'sd33554432;
    wire signed [31:0] A00  = ONE_Q725;
    wire signed [31:0] A01  = DT_Q725;
    wire signed [31:0] A10  = 32'sd0;
    wire signed [31:0] A11  = ONE_Q725;
    wire signed [31:0] AT00 = ONE_Q725;
    wire signed [31:0] AT01 = 32'sd0;
    wire signed [31:0] AT10 = DT_Q725;
    wire signed [31:0] AT11 = ONE_Q725;

    // ---- FSM --------------------------------------------------------------
    localparam [4:0] S_IDLE      = 5'd0;
    localparam [4:0] S_LATCH     = 5'd1;
    localparam [4:0] S_PRED_X    = 5'd2;
    localparam [4:0] S_PRED_P1   = 5'd3;
    localparam [4:0] S_PRED_P2   = 5'd4;
    localparam [4:0] S_PRED_P3   = 5'd5;
    localparam [4:0] S_INNOV_S   = 5'd6;
    localparam [4:0] S_INV_START = 5'd7;
    localparam [4:0] S_INV_WAIT  = 5'd8;
    localparam [4:0] S_GAIN      = 5'd9;
    localparam [4:0] S_RESID     = 5'd10;
    localparam [4:0] S_KINNOV    = 5'd11;
    localparam [4:0] S_UPD_X     = 5'd12;
    localparam [4:0] S_IK        = 5'd13;
    localparam [4:0] S_UPD_P     = 5'd14;
    localparam [4:0] S_COMMIT    = 5'd15;
    localparam [4:0] S_DONE      = 5'd16;

    reg [4:0] state, nstate;
    assign busy = (state != S_IDLE);

    always @(posedge clk or negedge rst_n) begin
        if (!rst_n) state <= S_IDLE;
        else        state <= nstate;
    end

    always @* begin
        nstate = state;
        case (state)
            S_IDLE      : if (din_valid)  nstate = S_LATCH;
            S_LATCH     :                 nstate = S_PRED_X;
            S_PRED_X    :                 nstate = S_PRED_P1;
            S_PRED_P1   :                 nstate = S_PRED_P2;
            S_PRED_P2   :                 nstate = S_PRED_P3;
            S_PRED_P3   :                 nstate = S_INNOV_S;
            S_INNOV_S   :                 nstate = S_INV_START;
            S_INV_START :                 nstate = S_INV_WAIT;
            S_INV_WAIT  : if (inv_done)   nstate = S_GAIN;
            S_GAIN      :                 nstate = S_RESID;
            S_RESID     :                 nstate = S_KINNOV;
            S_KINNOV    :                 nstate = S_UPD_X;
            S_UPD_X     :                 nstate = S_IK;
            S_IK        :                 nstate = S_UPD_P;
            S_UPD_P     :                 nstate = S_COMMIT;
            S_COMMIT    :                 nstate = S_DONE;
            S_DONE      :                 nstate = S_IDLE;
            default     :                 nstate = S_IDLE;
        endcase
    end

    // ---- Operand muxing for shared blocks ---------------------------------
    always @* begin
        mm_a00 = 32'd0; mm_a01 = 32'd0; mm_a10 = 32'd0; mm_a11 = 32'd0;
        mm_b00 = 32'd0; mm_b01 = 32'd0; mm_b10 = 32'd0; mm_b11 = 32'd0;
        ms_sub = 1'b0;
        ms_a00 = 32'd0; ms_a01 = 32'd0; ms_a10 = 32'd0; ms_a11 = 32'd0;
        ms_b00 = 32'd0; ms_b01 = 32'd0; ms_b10 = 32'd0; ms_b11 = 32'd0;

        case (state)
            S_PRED_P1: begin
                // tmp = A * P_prev
                mm_a00 = A00; mm_a01 = A01; mm_a10 = A10; mm_a11 = A11;
                mm_b00 = P00; mm_b01 = P01; mm_b10 = P10; mm_b11 = P11;
            end
            S_PRED_P2: begin
                // Pkp = tmp * A^T
                mm_a00 = Pkp00; mm_a01 = Pkp01; mm_a10 = Pkp10; mm_a11 = Pkp11;
                mm_b00 = AT00;  mm_b01 = AT01;  mm_b10 = AT10;  mm_b11 = AT11;
            end
            S_PRED_P3: begin
                // Pkp = Pkp + Q
                ms_sub = 1'b0;
                ms_a00 = Pkp00; ms_a01 = Pkp01; ms_a10 = Pkp10; ms_a11 = Pkp11;
                ms_b00 = Q00_Q725; ms_b01 = 32'd0; ms_b10 = 32'd0; ms_b11 = Q11_Q725;
            end
            S_INNOV_S: begin
                // S = Pkp + R
                ms_sub = 1'b0;
                ms_a00 = Pkp00; ms_a01 = Pkp01; ms_a10 = Pkp10; ms_a11 = Pkp11;
                ms_b00 = R00_Q725; ms_b01 = 32'd0; ms_b10 = 32'd0; ms_b11 = R11_Q725;
            end
            S_GAIN: begin
                // K = Pkp * Sinv
                mm_a00 = Pkp00;  mm_a01 = Pkp01;  mm_a10 = Pkp10;  mm_a11 = Pkp11;
                mm_b00 = Sinv00; mm_b01 = Sinv01; mm_b10 = Sinv10; mm_b11 = Sinv11;
            end
            S_RESID: begin
                // innov = Y - X_kp  (treat as col vector in col 0)
                ms_sub = 1'b1;
                ms_a00 = Y0;   ms_a01 = 32'd0; ms_a10 = Y1;   ms_a11 = 32'd0;
                ms_b00 = Xkp0; ms_b01 = 32'd0; ms_b10 = Xkp1; ms_b11 = 32'd0;
            end
            S_KINNOV: begin
                // Knew = K * innov
                mm_a00 = K00;     mm_a01 = K01;    mm_a10 = K10;    mm_a11 = K11;
                mm_b00 = innov0;  mm_b01 = 32'd0;  mm_b10 = innov1; mm_b11 = 32'd0;
            end
            S_UPD_X: begin
                // X_new col 0 = X_kp + Knew
                ms_sub = 1'b0;
                ms_a00 = Xkp0;  ms_a01 = 32'd0; ms_a10 = Xkp1;  ms_a11 = 32'd0;
                ms_b00 = Knew0; ms_b01 = 32'd0; ms_b10 = Knew1; ms_b11 = 32'd0;
            end
            S_IK: begin
                // IK = I - K
                ms_sub = 1'b1;
                ms_a00 = ONE_Q725; ms_a01 = 32'd0;    ms_a10 = 32'd0;    ms_a11 = ONE_Q725;
                ms_b00 = K00;      ms_b01 = K01;      ms_b10 = K10;      ms_b11 = K11;
            end
            S_UPD_P: begin
                // P_new = IK * Pkp
                mm_a00 = IK00;  mm_a01 = IK01;  mm_a10 = IK10;  mm_a11 = IK11;
                mm_b00 = Pkp00; mm_b01 = Pkp01; mm_b10 = Pkp10; mm_b11 = Pkp11;
            end
            default: ;
        endcase
    end

    // ---- Predicted angle helper: X0 + DT * X1 ----------------------------
    wire signed [31:0] dt_x1;
    wire signed [31:0] pred_x0;
    q_mult   u_predmul (.a(DT_Q725), .b(X1),    .p(dt_x1));
    q_addsub u_predadd (.sub(1'b0),  .a(X0),    .b(dt_x1), .y(pred_x0));

    // ---- Datapath sequential block ---------------------------------------
    always @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            X0 <= 32'd0;  X1 <= 32'd0;
            P00 <= P0_00_Q725; P01 <= 32'd0;
            P10 <= 32'd0;      P11 <= P0_11_Q725;

            Xkp0 <= 32'd0; Xkp1 <= 32'd0;
            Pkp00 <= 32'd0; Pkp01 <= 32'd0; Pkp10 <= 32'd0; Pkp11 <= 32'd0;
            S00 <= 32'd0; S01 <= 32'd0; S10 <= 32'd0; S11 <= 32'd0;
            Sinv00 <= 32'd0; Sinv01 <= 32'd0; Sinv10 <= 32'd0; Sinv11 <= 32'd0;
            K00 <= 32'd0; K01 <= 32'd0; K10 <= 32'd0; K11 <= 32'd0;
            Y0 <= 32'd0; Y1 <= 32'd0;
            innov0 <= 32'd0; innov1 <= 32'd0;
            Knew0 <= 32'd0; Knew1 <= 32'd0;
            IK00 <= 32'd0; IK01 <= 32'd0; IK10 <= 32'd0; IK11 <= 32'd0;
            Pnew00 <= 32'd0; Pnew01 <= 32'd0; Pnew10 <= 32'd0; Pnew11 <= 32'd0;

            est_angle_q9 <= 16'd0;
            dout_valid   <= 1'b0;
            inv_start    <= 1'b0;
            inv_in00 <= 32'd0; inv_in01 <= 32'd0; inv_in10 <= 32'd0; inv_in11 <= 32'd0;
        end
        else begin
            dout_valid <= 1'b0;
            inv_start  <= 1'b0;

            case (state)
                S_LATCH: begin
                    Y0 <= q9_to_q25(ang_meas_q9);
                    Y1 <= q9_to_q25(gyro_meas_q9);
                end

                S_PRED_X: begin
                    // X_kp = A * X_prev  (Xkp0 = X0 + DT*X1, Xkp1 = X1)
                    Xkp0 <= pred_x0;
                    Xkp1 <= X1;
                end

                S_PRED_P1: begin
                    Pkp00 <= mm_p00; Pkp01 <= mm_p01;
                    Pkp10 <= mm_p10; Pkp11 <= mm_p11;
                end

                S_PRED_P2: begin
                    Pkp00 <= mm_p00; Pkp01 <= mm_p01;
                    Pkp10 <= mm_p10; Pkp11 <= mm_p11;
                end

                S_PRED_P3: begin
                    Pkp00 <= ms_p00; Pkp01 <= ms_p01;
                    Pkp10 <= ms_p10; Pkp11 <= ms_p11;
                end

                S_INNOV_S: begin
                    S00 <= ms_p00; S01 <= ms_p01;
                    S10 <= ms_p10; S11 <= ms_p11;
                end

                S_INV_START: begin
                    inv_in00 <= S00; inv_in01 <= S01;
                    inv_in10 <= S10; inv_in11 <= S11;
                    inv_start <= 1'b1;
                end

                S_INV_WAIT: begin
                    if (inv_done) begin
                        Sinv00 <= inv_out00; Sinv01 <= inv_out01;
                        Sinv10 <= inv_out10; Sinv11 <= inv_out11;
                    end
                end

                S_GAIN: begin
                    K00 <= mm_p00; K01 <= mm_p01;
                    K10 <= mm_p10; K11 <= mm_p11;
                end

                S_RESID: begin
                    innov0 <= ms_p00;
                    innov1 <= ms_p10;
                end

                S_KINNOV: begin
                    Knew0 <= mm_p00;
                    Knew1 <= mm_p10;
                end

                S_UPD_X: begin
                    // Reuse Xkp regs to hold the new X temporarily.
                    Xkp0 <= ms_p00;
                    Xkp1 <= ms_p10;
                end

                S_IK: begin
                    IK00 <= ms_p00; IK01 <= ms_p01;
                    IK10 <= ms_p10; IK11 <= ms_p11;
                end

                S_UPD_P: begin
                    Pnew00 <= mm_p00; Pnew01 <= mm_p01;
                    Pnew10 <= mm_p10; Pnew11 <= mm_p11;
                end

                S_COMMIT: begin
                    X0  <= Xkp0;
                    X1  <= Xkp1;
                    P00 <= Pnew00; P01 <= Pnew01;
                    P10 <= Pnew10; P11 <= Pnew11;
                    est_angle_q9 <= q25_to_q9(Xkp0);
                    dout_valid   <= 1'b1;
                end

                default: ;
            endcase
        end
    end

endmodule

`default_nettype wire