import gradio as gr
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from io import BytesIO
import base64
import time
from functools import lru_cache, partial
from robot_souple_helpers import *

@lru_cache(maxsize=10)
def cached_neb_matrices(nx, dx, E, Ix, rho, S, cy):
    return neb_beam_matrices(nx, dx, E, Ix, rho, S, cy)

def generate_animation(ixe, u_history, title="Animation de la déformation"):
    if u_history.shape[1] < 2:
        return "Animation non disponible (pas assez de frames)"
    fig, ax = plt.subplots()
    ax.set_xlim(0, max(ixe))
    ax.set_ylim(np.min(u_history) - 0.1, np.max(u_history) + 0.1)
    line, = ax.plot([], [], lw=2, color='blue')
    ax.set_title(title)
    ax.set_xlabel('Position x (m)')
    ax.set_ylabel('Déplacement u (m)')
    def init():
        line.set_data([], [])
        return line,
    def animate(i):
        line.set_data(ixe, u_history[:, i])
        return line,
    # Further reduce frames for faster generation
    max_frames = 10
    frames = min(u_history.shape[1], max_frames)
    step = max(1, u_history.shape[1] // frames)
    anim = FuncAnimation(fig, animate, init_func=init, frames=range(0, u_history.shape[1], step), interval=300, blit=True)
    buf = BytesIO()
    anim.save(buf, writer='pillow', fps=3)
    buf.seek(0)
    gif = base64.b64encode(buf.read()).decode('utf-8')
    plt.close(fig)
    return f'<img src="data:image/gif;base64,{gif}" alt="Animation">'

def run_simulation(nx, delta, q, r, Kp, Ki, Kd, mu, with_noise, with_pert, niter, nopt, stagEps, solicitation, compare_open, without_measure, mode="Boucle ouverte"):
    start_time = time.time()
    try:
        beta = 0.25
        gamma = 0.5
        L = 500
        a = 5
        b = 5
        E = 70000
        rho = 2700e-9
        cy = 1e-4
        nstep = min(200, int(2 / delta))
        nA = int(np.floor(nx / 2))
        S = a * b
        Ix = a * b**3 / 12
        dx = L / nx
        time0 = np.linspace(0, delta * nstep, nstep + 1)
        ixe = np.linspace(dx, L, nx)
        vseed1 = 0 if with_pert else 42
        vseed2 = 1 if with_noise else 43
        Kfull, Cfull, Mfull = cached_neb_matrices(nx, dx, E, Ix, rho, S, cy)
        ndo1 = Kfull.shape[0] - 2
        K = Kfull[2:, 2:]
        C = Cfull[2:, 2:]
        M = Mfull[2:, 2:]
        induA = 2 * nA - 2
        induB = 2 * nx - 2
        fpert = generateNoiseTemporal(time0, 1, q, vseed1) if with_pert else np.zeros(len(time0))
        mpert = generateNoiseTemporal(time0, 1, r, vseed2) if with_noise else np.zeros(len(time0))
        settling_time = ""

        if mode == "Boucle ouverte":
            if solicitation == "échelon":
                fA = Heaviside(time0 - 1)
            elif solicitation == "sinusoïdale":
                fA = np.sin(2 * np.pi * 0.1 * time0)
            else:
                fA = Heaviside(time0 - 1)
            f = np.zeros((ndo1, len(time0)))
            f[induA, :] = fA + fpert
            u0 = np.zeros(ndo1)
            v0 = np.zeros(ndo1)
            a0 = np.linalg.solve(M, f[:, 0] - C @ v0 - K @ u0)
            u, v, a = Newmark(M, C, K, f, u0, v0, a0, delta, beta, gamma)
            uB_final = u[induB, -1]
            threshold = 0.01 * abs(uB_final)
            stable_idx = np.where(np.abs(u[induB, :] - uB_final) < threshold)[0]
            settling_time = time0[stable_idx[0]] if len(stable_idx) > 0 else "Non stabilisé"
            fig, ax = plt.subplots()
            ax.plot(time0, u[induB, :], label='u_B (tip)', linewidth=2)
            ax.plot(time0, fA, label='Consigne f_A', linestyle='--')
            ax.legend()
            ax.set_title(f"Boucle ouverte ({solicitation}) : Déplacement du tip")
            ax.set_xlabel("Temps (s)")
            ax.set_ylabel("Déplacement (m)")
            animation = generate_animation(ixe, u[::2, :], f"Déformation en boucle ouverte ({solicitation})")
        elif mode == "Boucle fermée PID":
            u_ref = Heaviside(time0 - 1)
            u = np.zeros((ndo1, len(time0)))
            v = np.zeros((ndo1, len(time0)))
            a = np.zeros((ndo1, len(time0)))
            u[:, 0] = np.zeros(ndo1)
            v[:, 0] = np.zeros(ndo1)
            a[:, 0] = np.linalg.solve(M, -C @ v[:, 0] - K @ u[:, 0])
            integ_e = 0
            prev_e = 0
            errors = []
            for step in range(1, len(time0)):
                uB_meas = u[induB, step-1] + mpert[step-1]
                e = u_ref[step] - uB_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]
                u[:, step], v[:, step], a[:, step] = newmark1stepMRHS(M, C, K, f_k, u[:, step-1], v[:, step-1], a[:, step-1], delta, beta, gamma)
                prev_e = e
                errors.append(e)
            u_open = None
            if compare_open:
                fA_comp = Heaviside(time0 - 1)
                f_comp = np.zeros((ndo1, len(time0)))
                f_comp[induA, :] = fA_comp + fpert
                u_open, _, _ = Newmark(M, C, K, f_comp, np.zeros(ndo1), np.zeros(ndo1), np.linalg.solve(M, f_comp[:, 0]), delta, beta, gamma)
            fig, ax = plt.subplots()
            ax.plot(time0, u[induB, :], label='u_B PID', linewidth=2)
            if u_open is not None:
                ax.plot(time0, u_open[induB, :], label='u_B open loop', linestyle='--')
            ax.plot(time0, u_ref, label='Consigne', linestyle='-.')
            ax.legend()
            ax.set_title("Boucle fermée PID : Déplacement u_B")
            ax.set_xlabel("Temps (s)")
            ax.set_ylabel("Déplacement (m)")
            animation = generate_animation(ixe, u[::2, :], "Déformation en boucle fermée")
        elif mode == "Filtre de Kalman":
            n_state = 2 * ndo1
            Q = np.eye(n_state) * q
            R = r
            P = np.eye(n_state) * 1e-3
            x_est = np.zeros(n_state)
            A = Fconstruct(M, C, K, delta, beta, gamma)[:n_state, :n_state]
            B_mat = Bconstruct(M, C, K, nA, delta, beta, gamma)[:n_state, :ndo1]
            H = np.zeros((1, n_state))
            H[0, induB] = 1
            u_sim = np.zeros((ndo1, len(time0)))
            v_sim = np.zeros((ndo1, len(time0)))
            u_sim[:, 0] = np.zeros(ndo1)
            v_sim[:, 0] = np.zeros(ndo1)
            estimates = []
            P_history = [P.copy()]
            for step in range(1, len(time0)):
                f_k = np.zeros(ndo1)
                u_sim[:, step], v_sim[:, step], _ = newmark1stepMRHS(M, C, K, f_k, u_sim[:, step-1], v_sim[:, step-1], np.zeros(ndo1), delta, beta, gamma)
                z = None
                if not without_measure:
                    z = u_sim[induB, step] + mpert[step]
                x_pred = A @ x_est
                P_pred = A @ P @ A.T + Q
                if not without_measure and z is not None:
                    K_kal = P_pred @ H.T @ np.linalg.inv(H @ P_pred @ H.T + R)
                    x_est = x_pred + K_kal @ (z - H @ x_pred)
                    P = (np.eye(n_state) - K_kal @ H) @ P_pred
                else:
                    x_est = x_pred
                    P = P_pred
                estimates.append(x_est[induB])
                P_history.append(P.copy())
            fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))
            ax1.plot(time0[1:], estimates, label='Estimation Kalman', linewidth=2)
            ax1.plot(time0, u_sim[induB, :], label='Vraie u_B', linestyle='--')
            ax1.legend()
            ax1.set_title("Estimation de u_B")
            ax1.set_xlabel("Temps (s)")
            ax1.set_ylabel("Déplacement (m)")
            cov_uB = [P[induB, induB] for P in P_history]
            ax2.plot(time0, cov_uB, label='Covariance u_B', linewidth=2)
            ax2.legend()
            ax2.set_title("Évolution de la covariance")
            ax2.set_xlabel("Temps (s)")
            ax2.set_ylabel("Variance")
            animation = "Animation non disponible pour Kalman"
        else:
            return "Mode non implémenté", "", f"Temps: {time.time() - start_time:.2f}s"

        buf = BytesIO()
        fig.savefig(buf, format="png", dpi=100)
        buf.seek(0)
        plot_b64 = base64.b64encode(buf.read()).decode('utf-8')
        plt.close(fig)
        plot_html = f'<img src="data:image/png;base64,{plot_b64}" style="max-width:100%;">'
        elapsed = f"Temps de calcul: {time.time() - start_time:.2f}s"
        if mode == "Boucle ouverte":
            info_str = f"{elapsed}, Stabilité: {settling_time}"
        else:
            info_str = elapsed
        return plot_html, animation, info_str
    except Exception as e:
        return f"Erreur: {str(e)}", "", f"Temps: {time.time() - start_time:.2f}s"

def generate_schema():
    fig, ax = plt.subplots(figsize=(10, 6))
    ax.text(0.1, 0.8, 'i', fontsize=14, ha='center')
    ax.arrow(0.15, 0.8, 0.1, 0, head_width=0.03, head_length=0.03, fc='k', ec='k')
    ax.text(0.3, 0.8, 'Moteur', fontsize=14, ha='center', bbox=dict(boxstyle="round,pad=0.3", facecolor="lightblue"))
    ax.arrow(0.4, 0.8, 0.1, 0, head_width=0.03, head_length=0.03, fc='k', ec='k')
    ax.text(0.55, 0.8, 'F_m', fontsize=14, ha='center')
    ax.arrow(0.6, 0.8, 0.1, 0, head_width=0.03, head_length=0.03, fc='k', ec='k')
    ax.text(0.75, 0.8, 'Transmission', fontsize=14, ha='center', bbox=dict(boxstyle="round,pad=0.3", facecolor="lightgreen"))
    ax.arrow(0.85, 0.8, 0.1, 0, head_width=0.03, head_length=0.03, fc='k', ec='k')
    ax.text(0.95, 0.8, 'F_t', fontsize=14, ha='center')
    ax.arrow(0.15, 0.5, 0.1, 0, head_width=0.03, head_length=0.03, fc='k', ec='k')
    ax.text(0.3, 0.5, 'C_m', fontsize=14, ha='center')
    ax.arrow(0.4, 0.5, 0.1, 0, head_width=0.03, head_length=0.03, fc='k', ec='k')
    ax.text(0.55, 0.5, 'θ', fontsize=14, ha='center')
    ax.arrow(0.6, 0.5, 0.1, 0, head_width=0.03, head_length=0.03, fc='k', ec='k')
    ax.text(0.75, 0.5, 'f_A', fontsize=14, ha='center')
    ax.arrow(0.85, 0.5, 0.1, 0, head_width=0.03, head_length=0.03, fc='k', ec='k')
    ax.text(0.95, 0.5, 'u_A', fontsize=14, ha='center')
    ax.arrow(0.5, 0.3, 0, -0.1, head_width=0.03, head_length=0.03, fc='k', ec='k')
    ax.text(0.5, 0.2, 'Poutre Flexible', fontsize=14, ha='center', bbox=dict(boxstyle="round,pad=0.3", facecolor="lightcoral"))
    ax.arrow(0.5, 0.1, 0, -0.05, head_width=0.03, head_length=0.03, fc='k', ec='k')
    ax.text(0.5, 0.05, 'u_B', fontsize=14, ha='center')
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    ax.axis('off')
    buf = BytesIO()
    fig.savefig(buf, format="png", dpi=100)
    buf.seek(0)
    img_b64 = base64.b64encode(buf.read()).decode('utf-8')
    plt.close(fig)
    return f'<img src="data:image/png;base64,{img_b64}" style="max-width:100%;">'

def compute_ft_gt(rendement, cm, theta_dot, fa_dot, ua_dot):
    if rendement == 1:
        ft = cm * theta_dot
        gt = fa_dot * ua_dot
        equal = "Oui" if np.isclose(ft, gt) else "Non"
        return f"F_t = {ft:.3f}, G_t = {gt:.3f}, Égal: {equal}"
    else:
        return r"Pour rendement unitaire, F_t = C_m * \dot{\theta}, G_t = f_A * \dot{u_A}"

def quiz_answer(choice):
    if choice == "Estimer u_B sans mesure directe":
        return "Correct ! Le jumeau permet d'estimer états non mesurés."
    else:
        return "Incorrect. Réessayez."

desc_prep = r"""
### Travail Préparatoire (Section 4.1 des rapports)
Étude du système : Schéma bloc en boucle ouverte. Relation F_t et G_t : Sous hypothèse de conservation d'énergie (rendement unitaire), F_t = G_t (d'après conservation C_m \dot{\theta} = f_A \dot{u_A}).  
Intérêt du jumeau : Estimer u_B sans mesure directe pour contrôle en boucle fermée.  
"""
desc_open = """
### Boucle Ouverte (Section 4.2.1)
Simulation sans rétroaction. Observez le déplacement u_B avec/sans perturbation. La position de B suit A avec retard (forme sinusoïdale ou échelon).  
"""
desc_pid = """
### Boucle Fermée PID
Contrôle avec PID pour suivre la consigne. Observez l'erreur de suivi.
"""
desc_kalman = """
### Filtre de Kalman
Estimation d'état avec filtre Kalman pour filtrer le bruit sur la mesure u_B.
"""

with gr.Blocks() as demo:
    gr.Markdown("# Jumeaux Numériques : Simulation de Robot Souple\nDémo interactive basée sur le cours de Renaud Ferrier - Centrale Lyon ENISE.")
    
    with gr.Tabs():
        with gr.Tab("Préparatoire"):
            gr.Markdown(desc_prep)
            gr.Markdown("### Schéma Bloc du Système")
            schema_img = gr.HTML(value=generate_schema())
            gr.Markdown("### Conservation d'Énergie : F_t = G_t")
            with gr.Row():
                rendement_slider = gr.Slider(0.5, 1, 1, label="Rendement (η)", info="Hypothèse rendement unitaire pour conservation")
                cm_input = gr.Number(1e-3, label="C_m (constante moteur)")
                theta_dot_input = gr.Number(10, label="\\dot{\\theta} (vitesse angulaire)")
                fa_dot_input = gr.Number(100, label="f_A (force en A)")
                ua_dot_input = gr.Number(0.1, label="\\dot{u_A} (vitesse en A)")
            compute_btn = gr.Button("Calculer F_t et G_t")
            ft_gt_output = gr.Textbox(label="Résultat")
            compute_btn.click(compute_ft_gt, inputs=[rendement_slider, cm_input, theta_dot_input, fa_dot_input, ua_dot_input], outputs=ft_gt_output)
            gr.Markdown("### Intérêt du Jumeau Numérique")
            quiz_radio = gr.Radio(["Estimer u_B sans mesure directe", "Simuler rapidement", "Réduire coûts"], label="Quel est l'intérêt principal du jumeau ?")
            quiz_btn = gr.Button("Vérifier")
            quiz_output = gr.Textbox()
            quiz_btn.click(quiz_answer, inputs=quiz_radio, outputs=quiz_output)
            gr.Markdown("### Probabilités Gaussiennes pour Bruits\nDensité : $ p(x) = \\frac{1}{\\sqrt{2\\pi \\sigma^2}} \\exp\\left(-\\frac{(x-\\mu)^2}{2\\sigma^2}\\right) $\nUtilisée pour covariances Q, R dans Kalman.")

        with gr.Tab("Boucle ouverte"):
            gr.Markdown(desc_open)
            solicitation_dropdown = gr.Dropdown(["échelon", "sinusoïdale"], value="échelon", label="Type de sollicitation", info="Échelon à t=1s ou sinusoïdale")
            with gr.Row():
                nx_slider = gr.Slider(5, 20, 10, 1, label="nx (éléments)", info="Précision poutre")
                delta_slider = gr.Slider(0.001, 0.05, 0.01, label="δt (pas temps)")
                q_slider = gr.Slider(0, 0.001, 0.0001, label="q (variance perturbation)")
            with_pert = gr.Checkbox(True, label="Perturbation ?")
            compare_open_open = gr.Checkbox(False, visible=False)
            without_measure_open = gr.Checkbox(False, visible=False)
            mode_open = gr.Textbox("Boucle ouverte", visible=False)
            run_open = gr.Button("Simuler")
            plot_open = gr.HTML()
            anim_open = gr.HTML()
            time_open = gr.Textbox(label="Info")
            run_open.click(run_simulation, inputs=[nx_slider, delta_slider, q_slider, gr.Number(value=0, visible=False), gr.Number(value=0, visible=False), gr.Number(value=0, visible=False), gr.Number(value=0, visible=False), gr.Number(value=0, visible=False), gr.Checkbox(value=False, visible=False), with_pert, gr.Number(value=0, visible=False), gr.Number(value=0, visible=False), gr.Number(value=0, visible=False), solicitation_dropdown, compare_open_open, without_measure_open, mode_open], outputs=[plot_open, anim_open, time_open])
        
        with gr.Tab("Boucle fermée PID"):
            gr.Markdown(desc_pid)
            with gr.Row():
                nx_pid = gr.Slider(5, 20, 10, 1, label="nx")
                delta_pid = gr.Slider(0.001, 0.05, 0.01, label="δt")
                r_pid = gr.Slider(0, 0.001, 0.0001, label="r (bruit mesure)")
                Kp_slider = gr.Slider(0, 10, 0.5, label="Kp")
                Ki_slider = gr.Slider(0, 1, 0.1, label="Ki")
                Kd_slider = gr.Slider(0, 0.1, 0.01, label="Kd")
            with_noise_pid = gr.Checkbox(True, label="Bruit mesure ?")
            compare_open_pid = gr.Checkbox(False, label="Comparer avec boucle ouverte")
            solicitation_pid = gr.Textbox("échelon", visible=False)
            without_measure_pid = gr.Checkbox(False, visible=False)
            mode_pid = gr.Textbox("Boucle fermée PID", visible=False)
            run_pid = gr.Button("Simuler")
            plot_pid = gr.HTML()
            anim_pid = gr.HTML()
            time_pid = gr.Textbox()
            run_pid.click(run_simulation, inputs=[nx_pid, delta_pid, gr.Number(value=0, visible=False), r_pid, Kp_slider, Ki_slider, Kd_slider, gr.Number(value=0, visible=False), with_noise_pid, gr.Checkbox(value=False, visible=False), gr.Number(value=0, visible=False), gr.Number(value=0, visible=False), gr.Number(value=0, visible=False), solicitation_pid, compare_open_pid, without_measure_pid, mode_pid], outputs=[plot_pid, anim_pid, time_pid])
        
        with gr.Tab("Filtre de Kalman"):
            gr.Markdown(desc_kalman)
            with gr.Row():
                nx_kal = gr.Slider(5, 15, 10, 1, label="nx")
                delta_kal = gr.Slider(0.001, 0.05, 0.01, label="δt")
                q_kal = gr.Slider(0, 0.001, 0.0001, label="q (bruit processus)")
                r_kal = gr.Slider(0, 0.001, 0.0001, label="r (bruit mesure)")
            without_measure_kal = gr.Checkbox(False, label="Sans mesure u_B (jumeau)")
            solicitation_kal = gr.Textbox("échelon", visible=False)
            compare_open_kal = gr.Checkbox(False, visible=False)
            mode_kal = gr.Textbox("Filtre de Kalman", visible=False)
            run_kal = gr.Button("Simuler")
            plot_kal = gr.HTML()
            anim_kal = gr.HTML()
            time_kal = gr.Textbox()
            run_kal.click(run_simulation, inputs=[nx_kal, delta_kal, q_kal, r_kal, gr.Number(value=0, visible=False), gr.Number(value=0, visible=False), gr.Number(value=0, visible=False), gr.Number(value=0, visible=False), gr.Checkbox(value=False, visible=False), gr.Checkbox(value=False, visible=False), gr.Number(value=0, visible=False), gr.Number(value=0, visible=False), gr.Number(value=0, visible=False), solicitation_kal, compare_open_kal, without_measure_kal, mode_kal], outputs=[plot_kal, anim_kal, time_kal])

demo.launch()