import streamlit as st
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from scipy import linalg
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import time
from simulation import run_simulation

def show_robot_tp():
    """Interface pour le TP du robot souple"""
    
    st.header("🤖 TP - Pilotage d'un Robot Souple")
    
    # Description du système
    st.markdown("""
    ### Système étudié : Poutre encastrée-libre
    
    **Caractéristiques** :
    - Longueur : 500 mm
    - Section : 5×5 mm
    - Matériau : Aluminium (E=70 GPa)
    - Point d'action : A (effort appliqué)
    - Point de mesure : B (déplacement désiré)
    
    **Objectif** : Commander la trajectoire de B en utilisant uniquement la mesure en A
    """)
    
    # Paramètres de simulation
    col1, col2, col3 = st.columns(3)
    
    with col1:
        st.markdown("### Paramètres géométriques")
        nstep = st.slider("Nombre de pas", 100, 2000, 1000)
        delta = st.slider("Pas de temps", 0.001, 0.1, 0.01, 0.001)
        L = st.slider("Longueur (mm)", 100, 1000, 500)
        a = st.slider("Largeur section", 1, 10, 5)
        b = st.slider("Hauteur section", 1, 10, 5)
        nx = st.slider("Nombre d'éléments", 5, 20, 10)
    
    with col2:
        st.markdown("### Paramètres de contrôle")
        control_type = st.selectbox("Type de contrôle",
                                   ["pid", "kalman", "pid_kalman"],
                                   format_func=lambda x: {"pid": "PID seul", "kalman": "Kalman seul", "pid_kalman": "PID + Kalman"}[x])
        
        meas_location = st.selectbox("Position de mesure", ["A", "B"],
                                     format_func=lambda x: f"Point {x}")

        Kp = 0.5
        Ki = 0.1
        Kd = 0.01
        if control_type in ['pid', 'pid_kalman']:
            Kp = st.slider("Kp", 0.0, 1.0, 0.5, 0.01)
            Ki = st.slider("Ki", 0.0, 1.0, 0.1, 0.01)
            Kd = st.slider("Kd", 0.0, 1.0, 0.01, 0.001)
    
    with col3:
        st.markdown("### Bruit")
        q = st.slider("Bruit force (q)", 1e-6, 1e-3, 1e-4, 1e-6, format="%.2e")
        r = st.slider("Bruit mesure (r)", 1e-6, 1e-3, 1e-4, 1e-6, format="%.2e")
        
        E = st.slider("Module Young (MPa)", 10000, 100000, 70000)
        rho = st.slider("Densité", 1e-9, 1e-8, 2.7e-9, 1e-10, format="%.2e")
        cy = st.slider("Amortissement", 1e-5, 1e-3, 1e-4, 1e-5, format="%.2e")
    
    # Simulation
    if st.button("🚀 Lancer la simulation", type="primary"):
        time0, ixe, u, v, a_sim, e_set, fA = run_simulation(nstep, delta, q, r, T=3, Kp=Kp if 'Kp' in locals() else 0.5, 
                                                          Ki=Ki if 'Ki' in locals() else 0.1, Kd=Kd if 'Kd' in locals() else 0.01, 
                                                          mu=0.01, beta=0.25, gamma=0.5, L=L, a=a, b=b, E=E, rho=rho, cy=cy, nx=nx, 
                                                          control_type=control_type, meas_location=meas_location)

        # Visualisation
        fig = make_subplots(rows=2, cols=1,
                           subplot_titles=("Configuration finale", "Déplacement pointe vs temps"),
                           vertical_spacing=0.15)
        
        # Configuration finale
        fig.add_trace(go.Scatter(x=ixe, y=u[::2, -1], name='Déplacement u',
                                line=dict(color='blue')),
                     row=1, col=1)
        fig.add_trace(go.Scatter(x=ixe, y=u[1::2, -1], name='Rotation θ',
                                line=dict(color='red')),
                     row=1, col=1)
        
        # Déplacement pointe
        fig.add_trace(go.Scatter(x=time0, y=u[-1, :], name='Pointe',
                                line=dict(color='blue')),
                     row=2, col=1)
        fig.add_trace(go.Scatter(x=time0, y=fA, name='Référence',
                                line=dict(color='red', dash='dash')),
                     row=2, col=1)
        
        fig.update_layout(height=600, showlegend=True)
        st.plotly_chart(fig, use_container_width=True)
        
        if control_type in ['pid', 'pid_kalman']:
            # Erreur
            fig_error = go.Figure()
            fig_error.add_trace(go.Scatter(x=time0, y=e_set, name='Erreur',
                                          line=dict(color='green')))
            fig_error.update_layout(title="Erreur de suivi", xaxis_title="Temps (s)", yaxis_title="Erreur")
            st.plotly_chart(fig_error, use_container_width=True)