import control
import numpy as np
import pandas as pd
from scipy.signal import butter, filtfilt
import matplotlib.pyplot as plt

# ============================================================
# Parameters
# ============================================================
n = 3  # number of masses in the chain

# <<< NEW: Wall-to-wall distance and equilibrium positions
L_wall = 1.0  # total distance between left and right wall [m]
spacing = L_wall / (n + 1)  # equal spacing at rest
p_eq = np.array([(i + 1) * spacing for i in range(n)])  # [0.25, 0.5, 0.75]

params = {
    'm': 1.0,  # m: mass,
    'k0': 50.0,  # k0: nominal spring coefficient,
    'c0': 2.0,  # c0: nominal damping coefficient,
    'alpha': 0.02,  # 2% stiffness loss per deg C, thermal stiffness sensitivity,
    'f_s': 0.4,  # Stiction, stiction amplitude.
    'lambda_wear': 5e-4  # Wear rate
}

# Reduced params: no thermal stiffness sensitivity, no stiction, no wear
params_reduced = {
    'm': 1.0,
    'k0': 50.0,
    'c0': 2.0,
    'alpha': 0.0,          # no thermal stiffness variation
    'f_s': 0.0,            # no stiction
    'lambda_wear': 0.0     # no wear
}

motor_params = {
    'R': 2.0,    # armature resistance
    'L': 0.5,    # armature inductance
    'Kt': 0.12,   # torque constant
    'Ke': 0.12,   # back-emf constant
    'N': 1.0,    # gear ratio
    'r': 0.02,   # pulley/pinion/effective radius
    'eta': 0.999,   # transmission efficiency
    'act_idx': 0   # actuator force applied to mass 0
}

Ts = 0.005


# ============================================================
# 1) LINEAR MODEL (chain + PMDC)
# ============================================================
# The linear model uses displacement coordinates internally.
# The A, B matrices are identical (they describe dynamics about equilibrium).
# We only change the C matrix and initial conditions to output absolute positions.

def build_system_matrices(n, m_val=1.0, k_val=50.0, c_val=1.0):
    m = np.ones(n) * m_val
    k = np.ones(n) * k_val
    c = np.ones(n) * c_val

    K_mat = np.zeros((n, n))
    C_mat = np.zeros((n, n))

    for i in range(n):
        K_mat[i, i] -= k[i]
        C_mat[i, i] -= c[i]
        if i > 0:
            K_mat[i, i-1] = k[i]
            C_mat[i, i-1] = c[i]
        if i < n - 1:
            K_mat[i, i] -= k[i+1]
            C_mat[i, i] -= c[i+1]
            K_mat[i, i+1] = k[i+1]
            C_mat[i, i+1] = c[i+1]

    K_mat[-1, -1] -= k[-1]
    C_mat[-1, -1] -= c[-1]

    Minv = np.diag(1.0 / m)

    A = np.block([
        [np.zeros((n, n)), np.eye(n)],
        [Minv @ K_mat,     Minv @ C_mat]
    ])

    B = np.zeros((2 * n, 1))
    B[n, 0] = 1.0 / m[0]

    C_out = np.zeros((2, 2 * n))
    C_out[0, 0] = 1.0
    C_out[1, n - 1] = 1.0

    return A, B, C_out


def build_augmented_linear_system(n, params, motor_params):
    A_chain, B_force, C_chain = build_system_matrices(
        n, m_val=params['m'], k_val=params['k0'], c_val=params['c0']
    )

    R   = motor_params['R']
    L   = motor_params['L']
    Kt  = motor_params['Kt']
    Ke  = motor_params['Ke']
    N   = motor_params.get('N', 1.0)
    r   = motor_params.get('r', 1.0)
    eta = motor_params.get('eta', 1.0)
    act_idx = motor_params.get('act_idx', 0)

    Kf = eta * N * Kt / r
    Kv = N * Ke / r

    nx_chain = 2 * n
    nx = nx_chain + 1

    A_aug = np.zeros((nx, nx))
    B_aug = np.zeros((nx, 1))

    A_aug[:nx_chain, :nx_chain] = A_chain
    A_aug[:nx_chain, -1] = B_force[:, 0] * Kf

    vel_idx = n + act_idx
    A_aug[-1, vel_idx] = -Kv / L
    A_aug[-1, -1]      = -R / L
    B_aug[-1, 0]        = 1.0 / L

    C_aug = np.zeros((4, nx))
    C_aug[0, 0]   = 1.0
    C_aug[1, n-1] = 1.0
    C_aug[2, -1]  = 1.0
    C_aug[3, -1]  = Kf

    D_aug = np.zeros((4, 1))

    return A_aug, B_aug, C_aug, D_aug, Kf, Kv


def discretize_linear_zoh(A, B, C, D, Ts):
    sys_c = control.ss(A, B, C, D)
    sys_d = control.c2d(sys_c, Ts, method='zoh')
    return np.array(sys_d.A), np.array(sys_d.B), np.array(sys_d.C), np.array(sys_d.D)


A_aug, B_aug, C_aug, D_aug, Kf, Kv = build_augmented_linear_system(n, params, motor_params)
Ad, Bd, Cd, Dd = discretize_linear_zoh(A_aug, B_aug, C_aug, D_aug, Ts)


# ============================================================
# 2) FULL NONLINEAR TRUTH MODEL — ABSOLUTE COORDINATES
# ============================================================

def truth_dynamics_with_motor(t, state, V, n, params, motor_params, p_eq, L_wall):  # <<< CHANGED signature
    """
    Full nonlinear truth model in ABSOLUTE coordinates.
    p[i] is now the absolute position from the left wall.
    State (4n+1): [p(n), v(n), temp(n), wear(n), i_a]
    """
    p    = state[0:n]
    v    = state[n:2*n]
    temp = state[2*n:3*n]
    wear = state[3*n:4*n]
    i_a  = state[4*n]

    dp     = v.copy()
    dv     = np.zeros(n)
    d_temp = np.zeros(n)
    d_wear = np.zeros(n)

    m_val   = params['m']
    R       = motor_params['R']
    L       = motor_params['L']
    Kt      = motor_params['Kt']
    Ke      = motor_params['Ke']
    N       = motor_params.get('N', 1.0)
    r       = motor_params.get('r', 1.0)
    eta     = motor_params.get('eta', 1.0)
    act_idx = motor_params.get('act_idx', 0)

    Kf_val = eta * N * Kt / r
    Kv_val = N * Ke / r
    F_act  = Kf_val * i_a

    f_mech = np.zeros(n)

    # <<< CHANGED: Left wall (at x=0) <-> mass 0
    # Spring extension = p[0] - 0 - natural_length = p[0] - spacing
    k_left = params['k0'] * (1.0 - params['alpha'] * (temp[0] - 20.0))
    c_left = params['c0'] * wear[0]
    f_mech[0] += -k_left * (p[0] - p_eq[0]) - c_left * v[0]

    # <<< CHANGED: Internal links
    for j in range(n - 1):
        temp_link = 0.5 * (temp[j] + temp[j+1])
        wear_link = 0.5 * (wear[j] + wear[j+1])
        k_link = params['k0'] * (1.0 - params['alpha'] * (temp_link - 20.0))
        c_link = params['c0'] * wear_link
        # Spring extension = (p[j+1] - p[j]) - spacing
        rel_p = (p[j+1] - p[j]) - spacing
        rel_v = v[j+1] - v[j]
        f_on_j = k_link * rel_p + c_link * rel_v
        f_mech[j]   += f_on_j
        f_mech[j+1] -= f_on_j

    # <<< CHANGED: Right wall (at x=L_wall) <-> mass n-1
    # Spring extension = L_wall - p[n-1] - spacing  (compressed when mass moves right)
    k_right = params['k0'] * (1.0 - params['alpha'] * (temp[-1] - 20.0))
    c_right = params['c0'] * wear[-1]
    f_mech[-1] += -k_right * (p[-1] - p_eq[-1]) - c_right * v[-1]

    # Stiction (unchanged — depends on velocity only)
    eps = 0.01
    f_stiction = -params['f_s'] * np.exp(-(v / eps)**2) * np.sign(v)

    f_total = f_mech + f_stiction
    f_total[act_idx] += F_act
    dv = f_total / m_val

    for j in range(n):
        power_j = np.abs(f_mech[j] * v[j])
        d_temp[j] = 0.1 * power_j - 0.05 * (temp[j] - 20.0)
        d_wear[j] = -params['lambda_wear'] * power_j

    V_in = float(np.atleast_1d(V)[0])
    di_a = (-R * i_a - Kv_val * v[act_idx] + V_in) / L

    return np.concatenate([dp, dv, d_temp, d_wear, [di_a]])


# ============================================================
# 3) REDUCED NONLINEAR MODEL — ABSOLUTE COORDINATES
# ============================================================

def truth_dynamics_reduced(t, state, V, n, params_red, motor_params, p_eq, L_wall):  # <<< CHANGED signature
    """
    Reduced nonlinear model in ABSOLUTE coordinates.
    State (3n+1): [p(n), v(n), temp(n), i_a]
    """
    p    = state[0:n]
    v    = state[n:2*n]
    temp = state[2*n:3*n]
    i_a  = state[3*n]

    dp = v.copy()
    dv = np.zeros(n)
    d_temp = np.zeros(n)

    m_val   = params_red['m']
    k0      = params_red['k0']
    c0      = params_red['c0']

    R       = motor_params['R']
    L       = motor_params['L']
    Kt      = motor_params['Kt']
    Ke      = motor_params['Ke']
    N       = motor_params.get('N', 1.0)
    r       = motor_params.get('r', 1.0)
    eta     = motor_params.get('eta', 1.0)
    act_idx = motor_params.get('act_idx', 0)

    Kf_val = eta * N * Kt / r
    Kv_val = N * Ke / r
    F_act  = Kf_val * i_a

    f_mech = np.zeros(n)

    # <<< CHANGED: Left wall <-> mass 0
    f_mech[0] += -k0 * (p[0] - p_eq[0]) - c0 * v[0]

    # <<< CHANGED: Internal links
    for j in range(n - 1):
        rel_p = (p[j+1] - p[j]) - spacing
        rel_v = v[j+1] - v[j]
        f_on_j = k0 * rel_p + c0 * rel_v
        f_mech[j]   += f_on_j
        f_mech[j+1] -= f_on_j

    # <<< CHANGED: Right wall <-> mass n-1
    f_mech[-1] += -k0 * (p[-1] - p_eq[-1]) - c0 * v[-1]

    f_total = f_mech.copy()
    f_total[act_idx] += F_act
    dv = f_total / m_val

    for j in range(n):
        power_j = np.abs(f_mech[j] * v[j])
        d_temp[j] = 0.1 * power_j - 0.05 * (temp[j] - 20.0)

    V_in = float(np.atleast_1d(V)[0])
    di_a = (-R * i_a - Kv_val * v[act_idx] + V_in) / L

    return np.concatenate([dp, dv, d_temp, [di_a]])


# ============================================================
# RK4 integrator (unchanged)
# ============================================================

def rk4_step(f, t, x, u, dt, *args):
    k1 = f(t,        x,             u, *args)
    k2 = f(t + dt/2, x + dt/2 * k1, u, *args)
    k3 = f(t + dt/2, x + dt/2 * k2, u, *args)
    k4 = f(t + dt,   x + dt   * k3, u, *args)
    return x + (dt / 6.0) * (k1 + 2*k2 + 2*k3 + k4)


# ============================================================
# Output function (unchanged — reads state[0] and state[n-1])
# ============================================================

def compute_output(state, n, motor_params, model_type='full'):
    Kf_val = (motor_params['eta'] * motor_params['N']
              * motor_params['Kt'] / motor_params['r'])
    if model_type == 'full':
        i_a = state[4*n]
    elif model_type == 'reduced':
        i_a = state[3*n]
    elif model_type == 'linear':
        i_a = state[2*n]
    else:
        raise ValueError(f"Unknown model_type: {model_type}")

    return np.array([
        state[0],        # p1 — now absolute position from left wall
        state[n - 1],    # p3 — now absolute position from left wall
        i_a,
        Kf_val * i_a
    ])


# ============================================================
# DATA GENERATION
# ============================================================

np.random.seed(42)

T_sim   = 10.0
N_steps = int(T_sim / Ts)
t_vec   = np.arange(N_steps + 1) * Ts

raw_noise = np.random.randn(N_steps + 1) * 12.0
b_filt, a_filt = butter(2, 0.05)
V_input = filtfilt(b_filt, a_filt, raw_noise)


# ---- A) Linear model ----
# Linear model internally uses displacement coords.
# We add p_eq offset when saving states/outputs.

nx_lin = Ad.shape[0]
ny_lin = Cd.shape[0]
x_lin  = np.zeros(nx_lin)

states_lin      = np.zeros((N_steps, nx_lin))
next_states_lin = np.zeros((N_steps, nx_lin))
inputs_lin      = np.zeros((N_steps, 1))
outputs_lin     = np.zeros((N_steps, ny_lin))

# <<< NEW: offset vector for linear state [p1,p2,p3, v1,v2,v3, i_a]
lin_offset = np.zeros(nx_lin)
lin_offset[0:n] = p_eq  # add equilibrium positions to position states

for k in range(N_steps):
    u_k = np.array([V_input[k]])
    y_k = Cd @ x_lin + Dd @ u_k
    x_next = Ad @ x_lin + Bd @ u_k

    states_lin[k]      = x_lin + lin_offset       # <<< CHANGED: save absolute
    inputs_lin[k]      = u_k
    # <<< CHANGED: offset outputs y_p1 and y_p3
    y_abs = y_k.copy()
    y_abs[0] += p_eq[0]      # y_p1
    y_abs[1] += p_eq[n-1]    # y_p3
    outputs_lin[k]     = y_abs
    next_states_lin[k] = x_next + lin_offset       # <<< CHANGED: save absolute
    x_lin = x_next

state_names_lin      = [f'p{i+1}' for i in range(n)] + [f'v{i+1}' for i in range(n)] + ['i_a']
next_state_names_lin = [f'next_{s}' for s in state_names_lin]
output_names         = ['y_p1', 'y_p3', 'y_ia', 'y_Fact']

df_lin = pd.DataFrame(
    np.hstack([states_lin, inputs_lin, next_states_lin, outputs_lin]),
    columns=state_names_lin + ['V'] + next_state_names_lin + output_names
)
df_lin.insert(0, 'time', t_vec[:N_steps])
df_lin.to_csv('out/linear_model_data.csv', index=False)
print(f"Linear model data saved: {df_lin.shape[0]} rows, {df_lin.shape[1]} cols")


# ---- B) Full nonlinear truth model ----

nx_full = 4 * n + 1
x_full  = np.zeros(nx_full)
x_full[0:n]    = p_eq        # <<< CHANGED: initial positions = equilibrium (absolute)
x_full[2*n:3*n] = 20.0
x_full[3*n:4*n] = 1.0

states_full      = np.zeros((N_steps, nx_full))
next_states_full = np.zeros((N_steps, nx_full))
inputs_full      = np.zeros((N_steps, 1))
outputs_full     = np.zeros((N_steps, 4))

for k in range(N_steps):
    u_k = V_input[k]
    y_k = compute_output(x_full, n, motor_params, model_type='full')
    x_next = rk4_step(
        truth_dynamics_with_motor, t_vec[k], x_full, u_k, Ts,
        n, params, motor_params, p_eq, L_wall    # <<< CHANGED: pass p_eq, L_wall
    )
    states_full[k]      = x_full
    inputs_full[k]      = u_k
    outputs_full[k]     = y_k
    next_states_full[k] = x_next
    x_full = x_next

state_names_full = (
    [f'p{i+1}' for i in range(n)]
    + [f'v{i+1}' for i in range(n)]
    + [f'temp{i+1}' for i in range(n)]
    + [f'wear{i+1}' for i in range(n)]
    + ['i_a']
)
next_state_names_full = [f'next_{s}' for s in state_names_full]

df_full = pd.DataFrame(
    np.hstack([states_full, inputs_full, next_states_full, outputs_full]),
    columns=state_names_full + ['V'] + next_state_names_full + output_names
)
df_full.insert(0, 'time', t_vec[:N_steps])
df_full.to_csv('out/nonlinear_full_truth_data.csv', index=False)
print(f"Full nonlinear truth data saved: {df_full.shape[0]} rows, {df_full.shape[1]} cols")


# ---- C) Reduced nonlinear model ----

nx_red = 3 * n + 1
x_red  = np.zeros(nx_red)
x_red[0:n]    = p_eq         # <<< CHANGED: initial positions = equilibrium (absolute)
x_red[2*n:3*n] = 20.0

states_red      = np.zeros((N_steps, nx_red))
next_states_red = np.zeros((N_steps, nx_red))
inputs_red      = np.zeros((N_steps, 1))
outputs_red     = np.zeros((N_steps, 4))

for k in range(N_steps):
    u_k = V_input[k]
    y_k = compute_output(x_red, n, motor_params, model_type='reduced')
    x_next = rk4_step(
        truth_dynamics_reduced, t_vec[k], x_red, u_k, Ts,
        n, params_reduced, motor_params, p_eq, L_wall    # <<< CHANGED: pass p_eq, L_wall
    )
    states_red[k]      = x_red
    inputs_red[k]      = u_k
    outputs_red[k]     = y_k
    next_states_red[k] = x_next
    x_red = x_next

state_names_red = (
    [f'p{i+1}' for i in range(n)]
    + [f'v{i+1}' for i in range(n)]
    + [f'temp{i+1}' for i in range(n)]
    + ['i_a']
)
next_state_names_red = [f'next_{s}' for s in state_names_red]

df_red = pd.DataFrame(
    np.hstack([states_red, inputs_red, next_states_red, outputs_red]),
    columns=state_names_red + ['V'] + next_state_names_red + output_names
)
df_red.insert(0, 'time', t_vec[:N_steps])
df_red.to_csv('out/nonlinear_reduced_model_data.csv', index=False)
print(f"Reduced nonlinear model data saved: {df_red.shape[0]} rows, {df_red.shape[1]} cols")


# ============================================================
# COMPARISON PLOT (unchanged)
# ============================================================

fig, axs = plt.subplots(6, 1, figsize=(13, 16), sharex=True)

axs[0].plot(t_vec[:N_steps], V_input[:N_steps], 'k-', alpha=0.6)
axs[0].set_ylabel('Voltage [V]')
axs[0].set_title('Motor input voltage')
axs[0].grid(True)

axs[1].plot(t_vec[:N_steps], states_lin[:, 0],  label='Linear')
axs[1].plot(t_vec[:N_steps], states_full[:, 0],  '--', label='Full NL')
axs[1].plot(t_vec[:N_steps], states_red[:, 0],  ':', lw=2, label='Reduced NL')
axs[1].set_ylabel('Position [m]')
axs[1].set_title('Mass 1 absolute position (p1)')
axs[1].legend()
axs[1].grid(True)

axs[2].plot(t_vec[:N_steps], states_lin[:, n-1],  label='Linear')
axs[2].plot(t_vec[:N_steps], states_full[:, n-1],  '--', label='Full NL')
axs[2].plot(t_vec[:N_steps], states_red[:, n-1],  ':', lw=2, label='Reduced NL')
axs[2].set_ylabel('Position [m]')
axs[2].set_title('Mass 3 absolute position (p3)')
axs[2].legend()
axs[2].grid(True)

axs[3].plot(t_vec[:N_steps], states_lin[:, -1],  label='Linear')
axs[3].plot(t_vec[:N_steps], states_full[:, -1],  '--', label='Full NL')
axs[3].plot(t_vec[:N_steps], states_red[:, -1],  ':', lw=2, label='Reduced NL')
axs[3].set_ylabel('Current [A]')
axs[3].set_title('Motor armature current (i_a)')
axs[3].legend()
axs[3].grid(True)

axs[4].plot(t_vec[:N_steps], states_full[:, 2*n],   label='Full NL temp1')
axs[4].plot(t_vec[:N_steps], states_full[:, 2*n+1], label='Full NL temp2')
axs[4].plot(t_vec[:N_steps], states_full[:, 2*n+2], label='Full NL temp3')
axs[4].plot(t_vec[:N_steps], states_red[:, 2*n],   '--', label='Reduced NL temp1')
axs[4].plot(t_vec[:N_steps], states_red[:, 2*n+1], '--', label='Reduced NL temp2')
axs[4].plot(t_vec[:N_steps], states_red[:, 2*n+2], '--', label='Reduced NL temp3')
axs[4].set_ylabel('Temperature [°C]')
axs[4].set_title('Spring temperatures')
axs[4].legend(fontsize=8, ncol=2)
axs[4].grid(True)

axs[5].plot(t_vec[:N_steps], states_full[:, 3*n],   label='wear1')
axs[5].plot(t_vec[:N_steps], states_full[:, 3*n+1], label='wear2')
axs[5].plot(t_vec[:N_steps], states_full[:, 3*n+2], label='wear3')
axs[5].axhline(1.0, color='gray', ls='--', label='Reduced NL (fixed=1.0)')
axs[5].set_ylabel('Wear factor')
axs[5].set_xlabel('Time [s]')
axs[5].set_title('Wear states (full NL only)')
axs[5].legend()
axs[5].grid(True)

plt.tight_layout()
plt.savefig('out/three_model_comparison.png', dpi=150)
plt.show()

print("\nDone! Files generated:")
print("  - out/linear_model_data.csv")
print("  - out/nonlinear_full_truth_data.csv")
print("  - out/nonlinear_reduced_model_data.csv")
print("  - out/three_model_comparison.png")
