import numpy as np
from scipy.optimize import linear_sum_assignment

def pairwise_mac(Phi: np.ndarray, Psi: np.ndarray) -> np.ndarray:
    """
    Pairwise MAC between mode shape sets Phi (d x m) and Psi (d x n).
    Columns are mode shapes. Works for real or complex.
    Returns an (m x n) array of MAC values in [0,1].
    """
    # Normalize columns
    Phi_n = Phi / np.linalg.norm(Phi, axis=0, keepdims=True)
    Psi_n = Psi / np.linalg.norm(Psi, axis=0, keepdims=True)
    G = Phi_n.conj().T @ Psi_n
    return np.abs(G) ** 2

def eig_to_freq_damp_from_s(lam: np.ndarray):
    """
    Continuous-time eigenvalues s = sigma + j*omega -> (f, zeta).
    f in Hz, zeta in [0,1].
    """
    sigma = lam.real
    omega = np.abs(lam.imag)
    f = omega / (2 * np.pi)
    zeta = -sigma / np.sqrt(sigma**2 + omega**2 + 1e-18)
    return f, zeta

def z_to_s(z: np.ndarray, dt: float) -> np.ndarray:
    """ Map discrete-time eigenvalues z to continuous-time s via log(z)/dt. """
    return np.log(z) / dt

def build_cost_matrix(
    f1, z1, Phi1, f2, z2, Phi2,
    w_f=1.0, w_z=0.5, w_mac=1.0, f_scale=None,
    gate_df=0.20, gate_z=0.05, gate_mac=0.7,
    cost_unassigned=2.0, big_M=1e6
):
    """
    Construct a padded square cost matrix for Hungarian assignment with gating and dummies.

    Parameters:
      - f1, z1: (m,) arrays for set A (reference)
      - Phi1: (d, m) mode shapes for set A
      - f2, z2: (n,) arrays for set B (to match)
      - Phi2: (d, n) mode shapes for set B
      - w_f, w_z, w_mac: weights
      - f_scale: normalization for frequency difference (default: max of both sets or 1.0)
      - gate_df: max allowed normalized freq diff (Hz normalized) for feasible match
      - gate_z: max allowed damping diff for feasible match
      - gate_mac: min required MAC for feasible match
      - cost_unassigned: penalty to leave a mode unmatched (tune to your problem)
      - big_M: prohibitive cost for infeasible pairs

    Returns:
      C_pad: (L, L) padded square cost matrix
      sizes: dict with m, n, L
    """
    f1 = np.asarray(f1).reshape(-1)
    z1 = np.asarray(z1).reshape(-1)
    f2 = np.asarray(f2).reshape(-1)
    z2 = np.asarray(z2).reshape(-1)

    m = f1.size
    n = f2.size

    if f_scale is None:
        f_scale = max(np.max(f1) if m else 1.0, np.max(f2) if n else 1.0, 1.0)

    # Pairwise components
    DF = np.abs(f1[:, None] - f2[None, :]) / f_scale           # (m,n)
    DZ = np.abs(z1[:, None] - z2[None, :])                      # (m,n)
    MAC = pairwise_mac(Phi1, Phi2) if (Phi1.size and Phi2.size) else np.zeros((m, n))
    one_minus_MAC = 1.0 - MAC

    # Base cost
    C = w_f * DF + w_z * DZ + w_mac * one_minus_MAC

    # Gating
    infeasible = (DF > gate_df) | (DZ > gate_z) | (MAC < gate_mac)
    C = np.where(infeasible, big_M, C)

    # Pad to square with dummy rows/cols so that unmatched is allowed
    L = max(m, n)
    if L == 0:
        return np.zeros((0, 0)), {"m": m, "n": n, "L": L}

    C_pad = np.full((L, L), cost_unassigned, dtype=float)

    # Place real costs into top-left block
    if m and n:
        C_pad[:m, :n] = C

    # For the padded parts (dummies), keep cost_unassigned (meaning: matching to dummy costs that penalty)

    return C_pad, {"m": m, "n": n, "L": L}

def hungarian_match(
    f1, z1, Phi1, f2, z2, Phi2,
    w_f=1.0, w_z=0.5, w_mac=1.0, f_scale=None,
    gate_df=0.20, gate_z=0.05, gate_mac=0.7,
    cost_unassigned=2.0, big_M=1e6,
    accept_cost_threshold=None
):
    """
    Run Hungarian assignment and return matches and unmatched lists.

    Returns:
      matches: list of (i, j, cost) with i in [0,m) and j in [0,n)
      unmatched_A: list of indices in A not matched to any real B (or rejected by accept threshold)
      unmatched_B: list of indices in B not matched to any real A (or rejected by accept threshold)
      C_pad: the square padded cost matrix actually optimized
    """
    C_pad, sizes = build_cost_matrix(
        f1, z1, Phi1, f2, z2, Phi2,
        w_f, w_z, w_mac, f_scale,
        gate_df, gate_z, gate_mac,
        cost_unassigned, big_M
    )
    m, n, L = sizes["m"], sizes["n"], sizes["L"]
    if L == 0:
        return [], [], [], C_pad

    row_ind, col_ind = linear_sum_assignment(C_pad)

    matches = []
    unmatched_A = []
    unmatched_B = []

    # Interpret the assignment on the unpadded indices
    for r, c in zip(row_ind, col_ind):
        # Case 1: real-to-real
        if r < m and c < n:
            cost = C_pad[r, c]
            if accept_cost_threshold is not None and cost > accept_cost_threshold:
                unmatched_A.append(r)
                unmatched_B.append(c)
            elif cost >= big_M:
                # Infeasible pair forced by padding; treat as unmatched
                unmatched_A.append(r)
                unmatched_B.append(c)
            else:
                matches.append((r, c, float(cost)))
        # Case 2: A matched to a dummy column => A unmatched
        elif r < m and c >= n:
            unmatched_A.append(r)
        # Case 3: dummy row matched to real B => B unmatched
        elif r >= m and c < n:
            unmatched_B.append(c)
        # Case 4: dummy-dummy (should not affect)
        else:
            pass

    # Deduplicate unmatched indices in case of multiple reasons
    unmatched_A = sorted(set([i for i in unmatched_A if 0 <= i < m]))
    unmatched_B = sorted(set([j for j in unmatched_B if 0 <= j < n]))

    return matches, unmatched_A, unmatched_B, C_pad

def demo():
    # Synthetic example: 3 reference modes vs 4 estimated modes
    rng = np.random.default_rng(42)
    d = 3  # DOF of mode shapes

    # Reference modes (A)
    fA = np.array([1.60, 3.00, 7.50])  # Hz
    zA = np.array([0.010, 0.020, 0.005])  # damping ratios
    PhiA = np.array([
        [1.0, 0.1, 0.2],
        [0.2, 1.0, 0.1],
        [0.1, 0.2, 1.0],
    ], dtype=float)
    PhiA = PhiA / np.linalg.norm(PhiA, axis=0, keepdims=True)

    # Estimated modes (B): 3 near matches + 1 spurious
    fB = np.array([1.62, 7.48, 2.95, 10.0])
    zB = np.array([0.011, 0.006, 0.022, 0.010])
    # Perturb PhiA columns to create similar shapes (plus one unrelated)
    Psi_sim = PhiA.copy()
    Psi_sim += 0.03 * rng.standard_normal(Psi_sim.shape)
    Psi_sim = Psi_sim / np.linalg.norm(Psi_sim, axis=0, keepdims=True)
    Psi_spur = rng.standard_normal((d, 1))
    Psi_spur = Psi_spur / np.linalg.norm(Psi_spur, axis=0, keepdims=True)
    PhiB = np.concatenate([Psi_sim[:, [0, 2, 1]], Psi_spur], axis=1)  # reorder to make matching non-trivial

    matches, unA, unB, C_pad = hungarian_match(
        fA, zA, PhiA, fB, zB, PhiB,
        w_f=1.0, w_z=0.5, w_mac=1.0,
        f_scale=max(fA.max(), fB.max()),
        gate_df=0.25, gate_z=0.05, gate_mac=0.6,
        cost_unassigned=1.2,   # tune this: lower => more likely to leave spurious unmatched
        big_M=1e6,
        accept_cost_threshold=0.9
    )

    print("Cost matrix (padded):\n", np.round(C_pad, 3))
    print("\nMatches (A_i -> B_j, cost):")
    for i, j, c in matches:
        print(f"  A[{i}] -> B[{j}], cost={c:.3f}, fA={fA[i]:.2f}, fB={fB[j]:.2f}, zA={zA[i]:.3f}, zB={zB[j]:.3f}")

    print("\nUnmatched in A:", unA)
    print("Unmatched in B:", unB)

if __name__ == "__main__":
    demo()