header_generator.py

from pathlib import Path
import re
import os
import json
import inspect
import time
import numpy as np
from crazyflie_model import CrazyFlieModel
from parameters import (
    AMP_X,
    AMP_Y,
    HORIZON_LENGTH,
    MODEL_FREQ,
    Q_DIAG,
    R_DIAG,
    RHO_EQ_PARAM,
    RHO_INEQ_PARAM,
    DELAY_STEPS,
    ADMM_USE_FLOAT,
    ADMM_ITERATIONS,
    ADMM_ENABLE_TRAJECTORY,
    ADMM_SOLVER_ARCH,
    ADMM_SOLVER_INPUT_WIDTH,
    MPC_LINEAR_DRAG_XY,
    MPC_LINEAR_DRAG_Z,
    STAR_INNER_RATIO,
    TRAJ_WARMSTART_PAD,
    TRAJ_TICK_DIV,
    TRAJ_SHAPE,
    U_ABS_MIN,
    U_ABS_MAX,
    XY_MIN,
    XY_MAX,
)

# Horizon length
N = HORIZON_LENGTH

# ADMM row-wise rho values (single source of truth for both KKT generation
# and emitted C/HLS constants).
rho_ineq = RHO_INEQ_PARAM
rho_eq = RHO_EQ_PARAM
SCRIPT_DIR = Path(__file__).resolve().parent
REPO_ROOT = SCRIPT_DIR.parent
TRAJ_DATA_HEADER_PATH = REPO_ROOT / "vitis_projects" / "ADMM" / "traj_data.h"
DATA_HEADER_PATH = REPO_ROOT / "vitis_projects" / "ADMM" / "data.h"
TEST_DATA_HEADER_PATH = REPO_ROOT / "vitis_projects" / "ADMM" / "test_data.h"
RTL_PARAMS_HEADER_PATH = REPO_ROOT / "vivado_project" / "vivado_project.srcs" / "sources_1" / "new" / "admm_autogen_params.vh"
RUNTIME_CONFIG_HEADER_PATH = REPO_ROOT / "vitis_projects" / "ADMM" / "admm_runtime_config.h"
ABQR_OVERRIDE_PATH = os.environ.get("ADMM_ABQR_OVERRIDE_PATH", "").strip()
GENERATOR_DIAGNOSTICS_PATH = os.environ.get("ADMM_GENERATOR_DIAGNOSTICS_PATH", "").strip()

# Initialize goal state
xg = np.zeros(13)
xg[3] = 1.0  # unit quaternion
# Create quadrotor instance
quad = CrazyFlieModel(freq=MODEL_FREQ)
ug = quad.hover_thrust
mass_kg = quad.mass

def _parse_matrix_block(text: str, name: str, rows: int, cols: int) -> np.ndarray:
    m = re.search(rf"\b{name}\s*<<\s*(.*?);", text, re.DOTALL)
    if m is None:
        raise ValueError(f"Missing block for {name} in override file")
    body = m.group(1)
    nums = re.findall(r"[-+]?\d*\.?\d+(?:[eE][-+]?\d+)?f?", body)
    vals = [float(tok[:-1] if tok.lower().endswith("f") else tok) for tok in nums]
    expected = rows * cols
    if len(vals) != expected:
        raise ValueError(f"{name} expected {expected} values, found {len(vals)}")
    return np.asarray(vals, dtype=np.float64).reshape(rows, cols)


def _load_abqr_override(path: Path) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    text = path.read_text()
    A_ovr = _parse_matrix_block(text, "A", 12, 12)
    B_ovr = _parse_matrix_block(text, "B", 12, 4)
    Q_ovr = _parse_matrix_block(text, "Q", 12, 12)
    R_ovr = _parse_matrix_block(text, "R", 4, 4)
    return A_ovr, B_ovr, Q_ovr, R_ovr


if ABQR_OVERRIDE_PATH:
    override_path = Path(ABQR_OVERRIDE_PATH)
    A, B, Q, R = _load_abqr_override(override_path)
    print(f"Using matrix override file: {override_path}")
else:
    # Get linearized system
    A, B = quad.get_linearized_dynamics(xg, ug)

    # Match dominant plant non-ideality (linear drag) in the MPC model.
    A[6, 6] -= (MPC_LINEAR_DRAG_XY / mass_kg) / MODEL_FREQ
    A[7, 7] -= (MPC_LINEAR_DRAG_XY / mass_kg) / MODEL_FREQ
    A[8, 8] -= (MPC_LINEAR_DRAG_Z / mass_kg) / MODEL_FREQ

    q_diag = Q_DIAG
    Q = np.diag(q_diag)
    R = np.diag(R_DIAG)


def load_traj_length_from_header(header_path: Path, horizon: int) -> int:
    text = header_path.read_text()
    m = re.search(r"#define\s+TRAJ_Q_PACKED_ROWS\s+(\d+)", text)
    if m is None:
        raise ValueError(f"Could not find TRAJ_Q_PACKED_ROWS in {header_path}")
    rows = int(m.group(1))
    traj_length = rows - horizon
    if traj_length <= 0:
        raise ValueError(f"Invalid TRAJ_LENGTH derived from {header_path}: {traj_length}")
    return traj_length


def load_traj_loop_start_from_header(header_path: Path) -> int:
    text = header_path.read_text()
    m = re.search(r"#define\s+TRAJ_LOOP_START_IDX\s+(\d+)", text)
    if m is None:
        return 0
    return int(m.group(1))

# Control input constraints
u_max = np.array([U_ABS_MAX - ug[0]] * 4)
u_min = np.array([U_ABS_MIN - ug[0]] * 4)


def build_Aeq_interleaved(Anp, Bnp, N):
    """
    Build dense A_eq and b_eq for interleaved ordering:
    z = [ x0, u0, x1, u1, ..., xN ]

    A_eq rows:
    [x0 = x_init]         ← initial condition
    [dynamics constraints] x1 - A x0 - B u0 = 0
                            x2 - A x1 - B u1 = 0
                                ...
                            xN - A x(N-1) - B u(N-1) = 0
    """

    n = Anp.shape[0]
    m = Bnp.shape[1]

    num_rows = (N + 1) * n     # (1 initial + N dynamics)
    num_var  = (N + 1) * n + N * m

    def idx_x(k):   # k in 0..N
        return k * (n + m)

    def idx_u(k):   # k in 0..N-1
        return k * (n + m) + n

    A_eq = np.zeros((num_rows, num_var))
    b_eq = np.zeros(num_rows)

    # 1) Initial condition: x0 = x_init
    A_eq[0:n, idx_x(0):idx_x(0) + n] = np.eye(n)
    b_eq[0:n] = np.zeros(n) #x0.reshape(-1)

    # 2) Dynamics constraints
    for k in range(N):
        row_start = (k + 1) * n  # shifted down by n rows
        xk_col    = idx_x(k)
        xkp1_col  = idx_x(k + 1)
        uk_col    = idx_u(k)

        A_eq[row_start:row_start + n, xkp1_col:xkp1_col + n] += np.eye(n)
        A_eq[row_start:row_start + n, xk_col:xk_col + n]     -= Anp
        A_eq[row_start:row_start + n, uk_col:uk_col + m]     -= Bnp

    return A_eq, b_eq


def add_block(mat, row_start, col_start, block):
    rows, cols = block.shape
    mat[row_start:row_start + rows, col_start:col_start + cols] += block


def add_sparse_block(mat, row_start, col_start, block):
    rows, cols = block.shape
    for r in range(rows):
        for c in range(cols):
            val = block[r, c]
            if val != 0:
                mat[row_start + r, col_start + c] += val


def idx_x_interleaved(k, n, m):
    return k * (n + m)


def idx_u_interleaved(k, n, m):
    return k * (n + m) + n


def build_kkt_direct(Anp, Bnp, Pmat, horizon, rho_eq_val, rho_ineq_val):
    n = Anp.shape[0]
    m = Bnp.shape[1]
    KKT = Pmat.copy()

    aTa = Anp.T @ Anp
    aTb = Anp.T @ Bnp
    bTa = Bnp.T @ Anp
    bTb = Bnp.T @ Bnp
    eye_n = np.eye(n)

    # Initial condition row: x0 = x_init
    x0 = idx_x_interleaved(0, n, m)
    add_block(KKT, x0, x0, rho_eq_val * eye_n)

    # Dynamics rows: x_{k+1} - A x_k - B u_k = 0
    for k in range(horizon):
        xk = idx_x_interleaved(k, n, m)
        uk = idx_u_interleaved(k, n, m)
        xkp1 = idx_x_interleaved(k + 1, n, m)

        add_block(KKT, xk, xk, rho_eq_val * aTa)
        add_block(KKT, xk, uk, rho_eq_val * aTb)
        add_block(KKT, xk, xkp1, -rho_eq_val * Anp.T)

        add_block(KKT, uk, xk, rho_eq_val * bTa)
        add_block(KKT, uk, uk, rho_eq_val * bTb)
        add_block(KKT, uk, xkp1, -rho_eq_val * Bnp.T)

        add_block(KKT, xkp1, xk, -rho_eq_val * Anp)
        add_block(KKT, xkp1, uk, -rho_eq_val * Bnp)
        add_block(KKT, xkp1, xkp1, rho_eq_val * eye_n)

    # Input box constraints.
    for k in range(horizon):
        uk = idx_u_interleaved(k, n, m)
        KKT[uk:uk + m, uk:uk + m] += rho_ineq_val * np.eye(m)

    # x/y state box constraints for predicted states k=1..N.
    for k in range(1, horizon + 1):
        xk = idx_x_interleaved(k, n, m)
        KKT[xk + 0, xk + 0] += rho_ineq_val
        KKT[xk + 1, xk + 1] += rho_ineq_val

    return KKT


def build_kkt_sparse(Anp, Bnp, Pmat, horizon, rho_eq_val, rho_ineq_val):
    from scipy import sparse

    n = Anp.shape[0]
    m = Bnp.shape[1]
    KKT = Pmat.tolil()

    aTa = Anp.T @ Anp
    aTb = Anp.T @ Bnp
    bTa = Bnp.T @ Anp
    bTb = Bnp.T @ Bnp
    eye_n = np.eye(n)

    x0 = idx_x_interleaved(0, n, m)
    add_sparse_block(KKT, x0, x0, rho_eq_val * eye_n)

    for k in range(horizon):
        xk = idx_x_interleaved(k, n, m)
        uk = idx_u_interleaved(k, n, m)
        xkp1 = idx_x_interleaved(k + 1, n, m)

        add_sparse_block(KKT, xk, xk, rho_eq_val * aTa)
        add_sparse_block(KKT, xk, uk, rho_eq_val * aTb)
        add_sparse_block(KKT, xk, xkp1, -rho_eq_val * Anp.T)

        add_sparse_block(KKT, uk, xk, rho_eq_val * bTa)
        add_sparse_block(KKT, uk, uk, rho_eq_val * bTb)
        add_sparse_block(KKT, uk, xkp1, -rho_eq_val * Bnp.T)

        add_sparse_block(KKT, xkp1, xk, -rho_eq_val * Anp)
        add_sparse_block(KKT, xkp1, uk, -rho_eq_val * Bnp)
        add_sparse_block(KKT, xkp1, xkp1, rho_eq_val * eye_n)

    for k in range(horizon):
        uk = idx_u_interleaved(k, n, m)
        for i in range(m):
            KKT[uk + i, uk + i] += rho_ineq_val

    for k in range(1, horizon + 1):
        xk = idx_x_interleaved(k, n, m)
        KKT[xk + 0, xk + 0] += rho_ineq_val
        KKT[xk + 1, xk + 1] += rho_ineq_val

    return KKT.tocsc()


def build_full_constraint_matrix(Anp, Bnp, horizon, n, m, n_ineq, num_var):
    A_eq, _ = build_Aeq_interleaved(Anp, Bnp, horizon)
    A_ineq = np.zeros((n_ineq, num_var))
    row = 0

    for k in range(horizon):
        u_start = k * (n + m) + n
        u_end = u_start + m
        A_ineq[row : row + m, u_start:u_end] = np.eye(m)
        row += m

    for k in range(1, horizon + 1):
        xk_start = k * (n + m)
        A_ineq[row, xk_start + 0] = 1.0
        row += 1
        A_ineq[row, xk_start + 1] = 1.0
        row += 1

    if row != n_ineq:
        raise RuntimeError(f"Internal inequality row mismatch: row={row}, n_ineq={n_ineq}")

    return A_eq, A_ineq, np.vstack([A_eq, A_ineq])


def build_full_constraint_matrix_sparse(Anp, Bnp, horizon, n, m, n_ineq, num_var):
    from scipy import sparse

    num_eq_rows = (horizon + 1) * n
    num_constraints = num_eq_rows + n_ineq
    A_full = sparse.lil_matrix((num_constraints, num_var), dtype=np.float64)

    for i in range(n):
        A_full[i, idx_x_interleaved(0, n, m) + i] = 1.0

    for k in range(horizon):
        row_start = (k + 1) * n
        xk_col = idx_x_interleaved(k, n, m)
        xkp1_col = idx_x_interleaved(k + 1, n, m)
        uk_col = idx_u_interleaved(k, n, m)

        add_sparse_block(A_full, row_start, xkp1_col, np.eye(n))
        add_sparse_block(A_full, row_start, xk_col, -Anp)
        add_sparse_block(A_full, row_start, uk_col, -Bnp)

    row = num_eq_rows
    for k in range(horizon):
        u_start = idx_u_interleaved(k, n, m)
        for i in range(m):
            A_full[row + i, u_start + i] = 1.0
        row += m

    for k in range(1, horizon + 1):
        xk_start = idx_x_interleaved(k, n, m)
        A_full[row, xk_start + 0] = 1.0
        row += 1
        A_full[row, xk_start + 1] = 1.0
        row += 1

    if row != num_constraints:
        raise RuntimeError(f"Internal sparse inequality row mismatch: row={row}, n_constr={num_constraints}")

    A_full = A_full.tocsc()
    A_eq = A_full[:num_eq_rows, :]
    A_ineq = A_full[num_eq_rows:, :]
    return A_eq, A_ineq, A_full


# def compute_lqr_terminal_cost(
#     A: np.ndarray,
#     B: np.ndarray,
#     Q: np.ndarray,
#     R: np.ndarray,
#     *,
#     max_iter: int = 2000,
#     tol: float = 1e-10,
# ) -> np.ndarray:
#     """Solve the discrete-time algebraic Riccati equation via fixed-point iteration."""
#     P = Q.astype(np.float64).copy()
#     for _ in range(max_iter):
#         bt_p = B.T @ P
#         s = R + bt_p @ B
#         k = np.linalg.solve(s, bt_p @ A)
#         p_next = Q + A.T @ P @ A - A.T @ P @ B @ k
#         p_next = 0.5 * (p_next + p_next.T)
#         if np.linalg.norm(p_next - P, ord="fro") < tol:
#             return p_next
#         P = p_next
#     raise RuntimeError("LQR terminal-cost Riccati iteration did not converge")


# Number of states and inputs
n = A.shape[0]
m = B.shape[1]
A_stage = A.copy()
B_stage = B.copy()

# Total number of variables
num_x = n * (N + 1)
num_u = m * N
num_var = num_x + num_u

# Build Hessian matrix P
# Q_terminal = compute_lqr_terminal_cost(A, B, Q, R)
blocks = []
for k in range(N):
    blocks.append(Q)     # Q_k
    blocks.append(R)     # R_k
blocks.append(Q)  # Final terminal Q_N from discrete LQR

try:
    from scipy import sparse
except ImportError as exc:
    raise RuntimeError(
        "scipy is required for sparse header generation. Activate ~/venv or install scipy."
    ) from exc

P_sparse = sparse.block_diag(blocks, format="csc", dtype=np.float64)

num_eq_rows = n * (N + 1)
num_input_ineq = m * N
num_state_box_ineq = 2 * N  # x and y for predicted states k=1..N
n_ineq = num_input_ineq + num_state_box_ineq
START_U_INEQ = 0
START_XY_INEQ = num_input_ineq

# Keep constraints aligned with trajectory frame for the rotated-square mode:
# midpoint of one side is treated as origin, so shift XY bounds equally.
xy_bound_shift = 0.0
xy_bound_halfspan = None
if TRAJ_SHAPE == "diamond1m_hold":
    xy_bound_shift = 0.5 / np.sqrt(2.0)  # side=1.0 => midpoint component
    # Match constraint box side to trajectory side (1.0 m).
    xy_bound_halfspan = 0.5
elif TRAJ_SHAPE == "star_hold":
    diagonal_angle = np.arctan2(AMP_X, AMP_Y)
    xy_bound_shift = AMP_X * STAR_INNER_RATIO * np.cos(diagonal_angle)

if xy_bound_halfspan is None:
    xy_min_eff = XY_MIN - xy_bound_shift
    xy_max_eff = XY_MAX - xy_bound_shift
else:
    xy_min_eff = -xy_bound_halfspan - xy_bound_shift
    xy_max_eff = xy_bound_halfspan - xy_bound_shift
l_u = np.hstack([np.tile(u_min, N), np.full(num_state_box_ineq, xy_min_eff)])
u_u = np.hstack([np.tile(u_max, N), np.full(num_state_box_ineq, xy_max_eff)])
b_eq = np.zeros(num_eq_rows)
l = np.hstack([b_eq, l_u])
u = np.hstack([b_eq, u_u])


num_constraints = num_eq_rows + n_ineq
rho_vect = rho_ineq * np.ones(num_constraints)
rho_vect[np.where(l == u)[0]] = rho_eq

generation_start = time.time()
KKT = build_kkt_sparse(A, B, P_sparse, N, rho_eq, rho_ineq)
A_eq, A_ineq, A_full = build_full_constraint_matrix_sparse(A, B, N, n, m, n_ineq, num_var)

try:
    from sksparse.cholmod import cholesky
except ImportError as exc:
    raise RuntimeError(
        "CHOLMOD sparse Cholesky requires scikit-sparse in the active Python environment. "
        "Run `source ~/venv/bin/activate` and install scikit-sparse/SuiteSparse before "
        "large-horizon benchmark generation."
    ) from exc

chol_parameters = inspect.signature(cholesky).parameters
if "lower" in chol_parameters and "order" in chol_parameters:
    factor_result = cholesky(KKT, lower=True, order="natural")
else:
    factor_result = cholesky(KKT, ordering_method="natural")
if isinstance(factor_result, tuple):
    L, chol_permutation = factor_result
else:
    L = factor_result.L()
    chol_permutation = factor_result.P()
L = L.tocsc()
if not np.array_equal(np.asarray(chol_permutation), np.arange(KKT.shape[0])):
    raise RuntimeError("CHOLMOD returned a non-identity permutation despite order='natural'")

def print_matrix(mat, name):
    print(f"{name}:")
    for row in mat:
        for val in row:
            str = f" {val:8.4f}" if val != 0 else "   .    "
            print(str, end="")
        print()
    print()

def get_max_bandwidth(L):
    n = L.shape[0]
    max_bandwidth = 0

    for i in range(n):
        row_nonzeros = np.nonzero(L[i, :])[0]
        if len(row_nonzeros) == 0:
            continue
        first_nonzero = row_nonzeros[0]
        last_nonzero = row_nonzeros[-1]
        bandwidth = last_nonzero - first_nonzero + 1
        if bandwidth > max_bandwidth:
            max_bandwidth = bandwidth

    return max_bandwidth

def convert_matrix_to_sparse_storage(M):
    n = M.shape[0]
    k = np.max([np.count_nonzero(M[i, :]) for i in range(n)])
    M_data = np.zeros((n, k))
    M_indexes = np.zeros((n, k), dtype=np.uint)
    
    for i in range(n):
        row_nonzeros = np.nonzero(M[i, :])[0]
        for j, col_idx in enumerate(row_nonzeros):
            M_data[i, j] = M[i, col_idx]
            M_indexes[i, j] = col_idx

    n_bits_idx = int(np.ceil(np.log2(np.max(M_indexes) + 1)))
    return M_data, M_indexes, n_bits_idx


def convert_sparse_matrix_to_sparse_storage(M):
    M = M.tocsr()
    rows, cols = M.shape
    counts = np.diff(M.indptr)
    k = int(counts.max()) if rows else 0
    M_data = np.zeros((rows, k))
    M_indexes = np.zeros((rows, k), dtype=np.uint)

    for i in range(rows):
        start = M.indptr[i]
        end = M.indptr[i + 1]
        row_cols = M.indices[start:end]
        row_vals = M.data[start:end]
        for j, (col_idx, val) in enumerate(zip(row_cols, row_vals)):
            M_data[i, j] = val
            M_indexes[i, j] = col_idx

    max_idx = int(M_indexes.max()) if M_indexes.size else max(cols - 1, 0)
    n_bits_idx = max(1, int(np.ceil(np.log2(max_idx + 1))))
    return M_data, M_indexes, n_bits_idx

def convert_chol_to_banded_storage(L):
    n = L.shape[0]
    max_bandwidth = get_max_bandwidth(L)
    L_row = np.zeros((n, max_bandwidth))
    
    for i in range(n):
        # Start column index of the band for this row
        start_col = max(0, i - max_bandwidth + 1)
        
        # Extract elements from start_col to diagonal (inclusive)
        row_values = [L[i, j] for j in range(start_col, i + 1)]
        
        # Pad with zeros on the left if needed
        while len(row_values) < max_bandwidth:
            row_values.insert(0, 0.0)
        
        row_values[-1] = 1.0 / L[i, i]  # Store reciprocal of diagonal

        L_row[i, :] = row_values
    
    return L_row


def get_sparse_lower_bandwidth(L):
    L = L.tocsr()
    max_bandwidth = 0
    for i in range(L.shape[0]):
        start = L.indptr[i]
        end = L.indptr[i + 1]
        cols = L.indices[start:end]
        cols = cols[cols <= i]
        if cols.size == 0:
            continue
        bandwidth = i - int(cols.min()) + 1
        max_bandwidth = max(max_bandwidth, bandwidth)
    return max_bandwidth


def convert_sparse_chol_to_banded_storage(L, max_bandwidth):
    L = L.tocsr()
    n = L.shape[0]
    L_row = np.zeros((n, max_bandwidth))

    for i in range(n):
        start_col = max(0, i - max_bandwidth + 1)
        values = dict(zip(L.indices[L.indptr[i]:L.indptr[i + 1]], L.data[L.indptr[i]:L.indptr[i + 1]]))
        row_values = [values.get(j, 0.0) for j in range(start_col, i + 1)]
        while len(row_values) < max_bandwidth:
            row_values.insert(0, 0.0)
        diag = values.get(i)
        if diag is None or diag == 0:
            raise RuntimeError(f"Missing or zero Cholesky diagonal at row {i}")
        row_values[-1] = 1.0 / diag
        L_row[i, :] = row_values

    return L_row

def convert_chol_transposed_to_banded_storage(L):
    n = L.shape[0]
    max_bandwidth = get_max_bandwidth(L)
    L_row = np.zeros((n, max_bandwidth))
    
    for i in range(n):
        end_col = min(n, i + max_bandwidth)        

        # print("i:",i," end_col:",end_col)

        row_values = [L[i, j] for j in range(i, end_col)]
        
        # Pad with zeros on the left if needed
        while len(row_values) < max_bandwidth:
            row_values.append(0.0)
        
        row_values[0] = 1.0 / L[i, i]  # Store reciprocal of diagonal

        L_row[i, :] = row_values
    
    return L_row


def convert_sparse_chol_transposed_to_banded_storage(L, max_bandwidth):
    U = L.T.tocsr()
    n = U.shape[0]
    L_row = np.zeros((n, max_bandwidth))

    for i in range(n):
        end_col = min(n, i + max_bandwidth)
        values = dict(zip(U.indices[U.indptr[i]:U.indptr[i + 1]], U.data[U.indptr[i]:U.indptr[i + 1]]))
        row_values = [values.get(j, 0.0) for j in range(i, end_col)]
        while len(row_values) < max_bandwidth:
            row_values.append(0.0)
        diag = values.get(i)
        if diag is None or diag == 0:
            raise RuntimeError(f"Missing or zero transposed Cholesky diagonal at row {i}")
        row_values[0] = 1.0 / diag
        L_row[i, :] = row_values

    return L_row


def generate_constants_header(constants_dict):
    lines = []
    lines.append("// Constants definitions\n")
    for name, value in constants_dict.items():
        lines.append(f"#define {name} {value}\n")
    lines.append("// end constants definitions\n")
    return "".join(lines)


def generate_verilog_params_header(constants_dict):
    lines = []
    lines.append("// Auto-generated by scripts/header_generator.py. Do not edit.\n")
    lines.append("`ifndef ADMM_AUTOGEN_PARAMS_VH\n")
    lines.append("`define ADMM_AUTOGEN_PARAMS_VH\n\n")
    lines.append(f"`define ADMM_HORIZON_LENGTH {constants_dict['HORIZON_LENGTH']}\n")
    lines.append(f"`define ADMM_N_STATE {constants_dict['STATE_SIZE']}\n")
    lines.append(f"`define ADMM_STAGE_SIZE {constants_dict['STAGE_SIZE']}\n")
    lines.append(f"`define ADMM_N_VAR {constants_dict['N_VAR']}\n")
    lines.append(f"`define ADMM_N_CONSTR {constants_dict['N_CONSTR']}\n")
    lines.append(f"`define ADMM_START_INEQ {constants_dict['START_INEQ']}\n")
    lines.append(f"`define ADMM_DELAY_STEPS {constants_dict['DELAY_STEPS']}\n")
    lines.append(f"`define ADMM_ENABLE_TRAJECTORY {constants_dict['ADMM_ENABLE_TRAJECTORY']}\n")
    lines.append(f"`define ADMM_SOLVER_INPUT_WIDTH {constants_dict['ADMM_SOLVER_INPUT_WIDTH']}\n")
    lines.append(f"`define ADMM_SOLVER_ARCH_STAGED_A {constants_dict['ADMM_SOLVER_ARCH_STAGED_A']}\n")
    lines.append(f"`define ADMM_SOLVER_ARCH_FULL_SPARSE {constants_dict['ADMM_SOLVER_ARCH_FULL_SPARSE']}\n")
    lines.append("`endif // ADMM_AUTOGEN_PARAMS_VH\n")
    return "".join(lines)


def generate_runtime_config_header() -> str:
    lines = []
    lines.append("// Auto-generated by scripts/header_generator.py. Do not edit.\n")
    lines.append("#ifndef ADMM_RUNTIME_CONFIG_H\n")
    lines.append("#define ADMM_RUNTIME_CONFIG_H\n\n")
    lines.append(f"#define ADMM_USE_FLOAT {1 if ADMM_USE_FLOAT else 0}\n")
    lines.append(f"#define ADMM_ITERATIONS {int(ADMM_ITERATIONS)}\n")
    lines.append(f"#define ADMM_ENABLE_TRAJECTORY {1 if ADMM_ENABLE_TRAJECTORY else 0}\n")
    lines.append(f"#define ADMM_SOLVER_ARCH_STAGED_A {1 if ADMM_SOLVER_ARCH == 'staged_a' else 0}\n")
    lines.append(f"#define ADMM_SOLVER_ARCH_FULL_SPARSE {1 if ADMM_SOLVER_ARCH == 'full_sparse' else 0}\n")
    lines.append(f"#define ADMM_SOLVER_ARCH_NAME \"{ADMM_SOLVER_ARCH}\"\n")
    lines.append(f"#define ADMM_SOLVER_INPUT_WIDTH {int(ADMM_SOLVER_INPUT_WIDTH)}\n")
    lines.append("\n#endif // ADMM_RUNTIME_CONFIG_H\n")
    return "".join(lines)

def generate_matrix_header(M, name, type="fp_t"):
    size_0, size_1 = M.shape
    lines = []
    lines.append(f"// Matrix {name} of size {size_0} x {size_1}\n")
    lines.append(f"#define {name.upper()}_ROWS {size_0}\n")
    lines.append(f"#define {name.upper()}_COLS {size_1}\n")
    lines.append(f"const {type} {name}[{name.upper()}_ROWS][{name.upper()}_COLS] = {{\n")
    for i in range(size_0):
        lines.append("   { " + ", ".join(f"({type}){M[i,j]:.8f}" for j in range(size_1)) + " },\n")
    lines.append(f"}};\n// end Matrix {name}\n")
    return "".join(lines)

def generate_vector_header(v, name, type="fp_t"):
    size = v.shape[0]
    lines = []
    lines.append(f"// Vector {name} of size {size}\n")
    lines.append(f"#define {name.upper()}_SIZE {size}\n")
    lines.append(f"const {type} {name}[{name.upper()}_SIZE] = {{\n")
    for i in range(size):
        lines.append(f"   ({type}){v[i]:.8f},\n")
    lines.append(f"}};\n// end Vector {name}\n")
    return "".join(lines)


def build_sparse_rows(M: np.ndarray, tol: float = 1e-12):
    rows, cols = M.shape
    counts = np.count_nonzero(np.abs(M) > tol, axis=1).astype(np.int32)
    max_nnz = int(counts.max()) if rows else 0
    idx = np.zeros((rows, max_nnz), dtype=np.int32)
    vals = np.zeros((rows, max_nnz), dtype=np.float64)
    for r in range(rows):
        nz_cols = np.flatnonzero(np.abs(M[r, :]) > tol)
        for k, c in enumerate(nz_cols):
            idx[r, k] = int(c)
            vals[r, k] = float(M[r, c])
    return counts, idx, vals, max_nnz


def build_sparse_cols(M: np.ndarray, tol: float = 1e-12):
    return build_sparse_rows(M.T, tol)

def generate_full_header(data, filename="data.h", guard="DATA_H"):
    with open(filename, "w") as f:
        f.write(f"#ifndef {guard}\n")
        f.write(f"#define {guard}\n\n")

        f.write('#include "data_types.h"\n\n')
        
        for d in data:
            f.write(d)
            f.write("\n// ===================== \n\n")
        
        f.write(f"#endif // {guard}\n")

def _checked_pow2_shift(name: str, value: int) -> int:
    if value <= 0:
        raise ValueError(f"{name} must be > 0, got {value}")
    shift = int(np.log2(value))
    if (1 << shift) != value:
        raise ValueError(f"{name} must be a power of two for shift-based rho ops, got {value}")
    return shift


def ADMM_iteration(l, u, iter):
    x = np.zeros(P.shape[0])
    z = np.zeros(A_full.shape[0])
    y = np.zeros(A_full.shape[0])
    q_vec = np.zeros(P.shape[0])
    state_rows = n
    ineq_start = A_eq.shape[0]
    u_ineq_start = A_eq.shape[0] + START_U_INEQ
    xy_ineq_start = A_eq.shape[0] + START_XY_INEQ

    for _ in range(iter):
        x = np.linalg.solve(KKT, A_full.T @ (rho_vect * (z - y)) - q_vec)
        Ax = A_full @ x

        z_new = np.zeros_like(z)
        for idx in range(A_full.shape[0]):
            if idx < state_rows:
                z_new[idx] = l[idx]
            elif idx < ineq_start:
                z_new[idx] = 0.0
            else:
                zi = Ax[idx] + y[idx]
                if idx >= xy_ineq_start:
                    z_new[idx] = np.clip(zi, l[idx], u[idx])
                elif idx >= u_ineq_start:
                    z_new[idx] = np.clip(zi, l[idx], u[idx])
                else:
                    z_new[idx] = np.clip(zi, l[idx], u[idx])

        y = y + (Ax - z_new)
        z = z_new
    
    return x, z, y

def testOSQP(l,u, iter):
    import osqp
    from scipy import sparse

    P_csc = sparse.csc_matrix(P)
    A_csc = sparse.csc_matrix(A_full)

    prob = osqp.OSQP()
    prob.setup(P_csc, np.zeros(P.shape[0]), A_csc, l, u, verbose=False, rho=rho_eq, adaptive_rho = False, max_iter=iter)
    res = prob.solve()

    return res.x

chol_bandwidth = get_sparse_lower_bandwidth(L)
L_banded = convert_sparse_chol_to_banded_storage(L, chol_bandwidth)
LT_banded = convert_sparse_chol_transposed_to_banded_storage(L, chol_bandwidth)
generation_elapsed_s = time.time() - generation_start

generator_diagnostics = {
    "horizon": int(N),
    "solver_arch": ADMM_SOLVER_ARCH,
    "admm_iterations": int(ADMM_ITERATIONS),
    "n_var": int(num_var),
    "n_constr": int(num_constraints),
    "n_ineq": int(n_ineq),
    "kkt_shape": [int(KKT.shape[0]), int(KKT.shape[1])],
    "kkt_nnz": int(KKT.nnz),
    "a_full_shape": [int(A_full.shape[0]), int(A_full.shape[1])],
    "a_full_nnz": int(A_full.nnz),
    "chol_l_nnz": int(L.nnz),
    "chol_bandwidth": int(chol_bandwidth),
    "generation_elapsed_s": float(generation_elapsed_s),
    "factorization": "cholmod_natural",
}
print("HEADER_GENERATOR_DIAGNOSTICS " + json.dumps(generator_diagnostics, sort_keys=True))
if GENERATOR_DIAGNOSTICS_PATH:
    diagnostics_path = Path(GENERATOR_DIAGNOSTICS_PATH)
    diagnostics_path.parent.mkdir(parents=True, exist_ok=True)
    diagnostics_path.write_text(json.dumps(generator_diagnostics, indent=2, sort_keys=True) + "\n")
constants = {}
constants["HORIZON_LENGTH"] = N
constants["STATE_SIZE"] = n
constants["INPUT_SIZE"] = m
constants["STAGE_SIZE"] = n + m
constants["N_VAR"] = num_var
constants["N_CONSTR"] = num_constraints
constants["N_INEQ"] = n_ineq
constants["START_INEQ"] = num_eq_rows
constants["START_U_INEQ"] = num_eq_rows + START_U_INEQ
constants["START_XY_INEQ"] = num_eq_rows + START_XY_INEQ
constants["DELAY_STEPS"] = DELAY_STEPS
rho_shift_ineq = _checked_pow2_shift("rho_ineq", int(rho_ineq))
rho_shift_eq = _checked_pow2_shift("rho_eq", int(rho_eq))
constants["RHO_SHIFT_INEQ"] = rho_shift_ineq
constants["RHO_SHIFT_EQ"] = rho_shift_eq
# Backward-compat alias for existing code/tests that still use RHO_SHIFT.
constants["RHO_SHIFT"] = rho_shift_eq
constants["U_MIN"] = u_min[0]
constants["U_MAX"] = u_max[0]
constants["XY_MIN"] = xy_min_eff
constants["XY_MAX"] = xy_max_eff
constants["U_HOVER"] = ug[0]
constants["ADMM_ENABLE_TRAJECTORY"] = 1 if ADMM_ENABLE_TRAJECTORY else 0
constants["ADMM_SOLVER_INPUT_WIDTH"] = ADMM_SOLVER_INPUT_WIDTH
constants["ADMM_SOLVER_ARCH_STAGED_A"] = 1 if ADMM_SOLVER_ARCH == "staged_a" else 0
constants["ADMM_SOLVER_ARCH_FULL_SPARSE"] = 1 if ADMM_SOLVER_ARCH == "full_sparse" else 0
if ADMM_ENABLE_TRAJECTORY:
    constants["TRAJ_LENGTH"] = load_traj_length_from_header(TRAJ_DATA_HEADER_PATH, horizon=N)
    constants["TRAJ_TICK_DIV"] = TRAJ_TICK_DIV
    constants["TRAJ_WARMSTART_PAD"] = TRAJ_WARMSTART_PAD
    constants["TRAJ_LOOP_START_IDX"] = load_traj_loop_start_from_header(TRAJ_DATA_HEADER_PATH)
else:
    constants["TRAJ_LENGTH"] = 0
    constants["TRAJ_TICK_DIV"] = 1
    constants["TRAJ_WARMSTART_PAD"] = 0
    constants["TRAJ_LOOP_START_IDX"] = 0

data = []
data.append(generate_constants_header(constants))
data.append(generate_matrix_header(L_banded, "L_banded"))
data.append(generate_matrix_header(LT_banded, "LT_banded"))

if ADMM_SOLVER_ARCH == "staged_a":
    A_stage_row_counts, A_stage_row_cols, A_stage_row_vals, A_stage_nnz_max = build_sparse_rows(A_stage)
    A_stage_col_counts, A_stage_col_rows, A_stage_col_vals, A_stage_t_nnz_max = build_sparse_cols(A_stage)
    constants["A_STAGE_NNZ_MAX"] = A_stage_nnz_max
    constants["A_STAGE_T_NNZ_MAX"] = A_stage_t_nnz_max

    data[0] = generate_constants_header(constants)
    data.append(generate_vector_header(A_stage_row_counts, "A_stage_row_counts", type="int"))
    data.append(generate_matrix_header(A_stage_row_cols, "A_stage_row_cols", type="int"))
    data.append(generate_matrix_header(A_stage_row_vals, "A_stage_row_vals"))
    data.append(generate_vector_header(A_stage_col_counts, "A_stage_col_counts", type="int"))
    data.append(generate_matrix_header(A_stage_col_rows, "A_stage_col_rows", type="int"))
    data.append(generate_matrix_header(A_stage_col_vals, "A_stage_col_vals"))
    data.append(generate_matrix_header(B_stage, "B_stage"))
elif ADMM_SOLVER_ARCH == "full_sparse":
    A_sparse_data, A_sparse_indexes, A_n_bits_idx = convert_sparse_matrix_to_sparse_storage(A_full)
    AT_sparse_data, AT_sparse_indexes, AT_n_bits_idx = convert_sparse_matrix_to_sparse_storage(A_full.T)
    constants["A_SPARSE_INDEX_BITS"] = A_n_bits_idx
    constants["AT_SPARSE_INDEX_BITS"] = AT_n_bits_idx

    data[0] = generate_constants_header(constants)
    data.append(generate_matrix_header(A_sparse_data, "A_sparse_data"))
    data.append(generate_matrix_header(A_sparse_indexes, "A_sparse_indexes", type=f"ap_uint<{A_n_bits_idx}>"))
    data.append(generate_matrix_header(AT_sparse_data, "AT_sparse_data"))
    data.append(generate_matrix_header(AT_sparse_indexes, "AT_sparse_indexes", type=f"ap_uint<{AT_n_bits_idx}>"))
else:
    raise RuntimeError(f"Unsupported ADMM_SOLVER_ARCH: {ADMM_SOLVER_ARCH}")

generate_full_header(data, filename=str(DATA_HEADER_PATH))
RTL_PARAMS_HEADER_PATH.write_text(generate_verilog_params_header(constants))
RUNTIME_CONFIG_HEADER_PATH.write_text(generate_runtime_config_header())

def generate_csim_reference():
    """Write the cold-start, finite-iteration reference for ADMM_test.cpp.

    The testbench supplies this initial state and no trajectory-start command,
    so q is zero. This double-precision reference checks implementation agreement;
    it is not an optimal solution or a convergence certificate.
    """
    from scipy.linalg import cho_factor, cho_solve

    lower, upper = l.copy(), u.copy()
    lower[:12] = upper[:12] = np.array([.1, .1, -.1] + [0.] * 9)
    factor = cho_factor(KKT.toarray())
    x = np.zeros(P_sparse.shape[0])
    z = np.zeros(A_full.shape[0])
    dual = np.zeros_like(z)
    for _ in range(ADMM_ITERATIONS):
        x = cho_solve(factor, A_full.T @ (rho_vect * (z - dual)))
        ax = A_full @ x
        z = np.clip(ax + dual, lower, upper)
        dual += ax - z
    if not np.all(np.isfinite(x)):
        raise ValueError("Nonfinite C simulation reference")
    generate_full_header(
        [generate_vector_header(x, "ADMM_x_after_hw_iters", type="double")],
        filename=str(TEST_DATA_HEADER_PATH), guard="TEST_DATA_H",
    )
    print(f"Generated cold-start C simulation reference: H={N}, iterations={ADMM_ITERATIONS}")


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Generate solver headers for the configured MPC problem.")
    parser.add_argument("--csim-reference", action="store_true",
                        help="also generate the finite-iteration reference for ADMM_test.cpp")
    args = parser.parse_args()
    if args.csim_reference:
        generate_csim_reference()