import numpy as np
from scipy.sparse.linalg import spsolve
from beam import neb_beam_matrices, newmark1step_mrhs
from noise import generate_noise_temporal

def kalman_filter(z, x_est, P, A, B, H, Q, R, u=0):
    # Prediction
    x_pred = A @ x_est + B * u
    P_pred = A @ P @ A.T + Q

    # Update
    K = P_pred @ H.T @ np.linalg.inv(H @ P_pred @ H.T + R)
    x_est = x_pred + K @ (z - H @ x_pred)
    P = (np.eye(len(x_est)) - K @ H) @ P_pred

    return x_est, P

def run_simulation(nstep=1000, delta=0.01, q=0.0001, r=0.0001, T=3, Kp=5e-1, Ki=1e-1, Kd=1e-2, mu=1e-1, beta=0.25, gamma=0.5, L=500, a=5, b=5, E=70000, rho=2700e-9, cy=1e-4, nx=10, nA=None, control_type='pid', meas_location='B'):
    if nA is None:
        nA = nx
    
    S = a * b
    Ix = a * b**3 / 12
    dx = L / nx
    time = np.arange(delta, (nstep+1)*delta, delta)
    time0 = np.concatenate([[0], time])
    
    ixe = np.arange(dx, L+dx, dx)
    
    is_matlab = True
    vseed1 = 0
    vseed2 = 1
    vseed3 = 2
    
    npt = nx + 1
    ndof = 2 * npt
    ndo1 = ndof - 2
    induA = 2 * nA - 1
    induB = 2 * nx - 1
    
    fpert = generate_noise_temporal(time0, 1, q, vseed1, is_matlab)
    mpert = generate_noise_temporal(time0, 1, r, vseed2, is_matlab)
    
    fA = np.zeros_like(time0)
    fA[len(fA)//10:] = 1
    
    ffull = np.zeros((ndof, nstep+1))
    Kfull, Cfull, Mfull = neb_beam_matrices(nx, dx, E, Ix, rho, S, cy)
    K = Kfull[2:, 2:]
    C = Cfull[2:, 2:]
    M = Mfull[2:, 2:]
    f = ffull[2:, :]
    f[induA, :] = fA + fpert
    
    u0 = np.zeros(ndo1)
    v0 = np.zeros(ndo1)
    
    u = np.zeros((ndo1, nstep+1))
    v = np.zeros((ndo1, nstep+1))
    a = np.zeros((ndo1, nstep+1))
    a0 = spsolve(M, f[:, 0] - C @ v0 - K @ u0)
    a[:, 0] = a0
    u[:, 0] = u0
    v[:, 0] = v0

    integ_e = 0
    prev_e = 0
    t0 = 1.0
    u_ref = (time0 >= t0).astype(float)
    e_set = [0]

    indu_meas = induA if meas_location == 'A' else induB

    # Initialize Kalman if needed
    if control_type in ['kalman', 'pid_kalman']:
        x_est = np.array([0.0, 0.0])
        P = np.eye(2) * 0.1
        A = np.array([[1, delta], [0, 1]])
        H = np.array([1, 0])
        Q = np.eye(2) * q
        R = r

    for step in range(1, nstep+1):
        if control_type == 'pid':
            u_meas = u[indu_meas, step-1] + mpert[step]
            e = u_ref[step] - u_meas
            integ_e += e * delta
            de = (e - prev_e) / delta
            Fpid = Kp * e + Ki * integ_e + Kd * de
            f_k = np.zeros(ndo1)
            f_k[induA] = Fpid + fpert[step]
            prev_e = e
            e_set.append(e)
        elif control_type == 'kalman':
            f_k = np.zeros(ndo1)
            f_k[induA] = fpert[step]  # open loop
            e_set.append(0)  # no error for kalman
        elif control_type == 'pid_kalman':
            z = u[indu_meas, step-1] + mpert[step]
            x_est, P = kalman_filter(z, x_est, P, A, np.zeros(2), H, Q, R)
            u_est = x_est[0]
            e = u_ref[step] - u_est
            integ_e += e * delta
            de = (e - prev_e) / delta
            Fpid = Kp * e + Ki * integ_e + Kd * de
            f_k = np.zeros(ndo1)
            f_k[induA] = Fpid + fpert[step]
            prev_e = e
            e_set.append(e)

        u[:, step], v[:, step], a[:, step] = newmark1step_mrhs(M, C, K, f_k, u[:, step-1], v[:, step-1], a[:, step-1], delta, beta, gamma)

    return time0, ixe, u, v, a, e_set, fA