import streamlit as st
from brian2 import *
import numpy as np
import plotly.graph_objects as go

def run_simulation(num_neurons=12, sim_time=100*ms):
    # Установка времени симуляции и временного шага
    defaultclock.dt = 0.1*ms

    # Определение уравнений для нейронов
    eqs = '''
    dv/dt = (I - v)/tau : 1 (unless refractory)
    I : 1
    tau : second
    '''

    # Создание группы нейронов
    G = NeuronGroup(num_neurons, eqs, threshold='v>1', reset='v = 0', refractory=5*ms, method='euler')

    # Начальные условия
    G.I = np.linspace(1.5, 2.0, num_neurons)  # Разные токи для разных сегментов
    G.tau = 10*ms

    # Создание переменной для записи результатов
    M = StateMonitor(G, 'v', record=True)

    # Запуск симуляции
    run(sim_time)

    # Моделирование механики червя
    # Упрощенная модель: каждый сегмент движется в зависимости от активности нейронов

    # Параметры червя
    segment_length = 0.5  # Длина каждого сегмента
    num_segments = num_neurons
    worm_length = segment_length * num_segments

    # Начальные позиции сегментов (случайные)
    num_frames = len(M.t)
    positions = np.zeros((num_frames, num_segments, 2))
    initial_positions = np.random.rand(num_segments, 2) * 10  # Случайные начальные позиции в пространстве 10x10
    positions[0, :, :] = initial_positions

    # Обновление позиций сегментов на основе активности нейронов
    for frame in range(1, num_frames):
        for i in range(num_segments):
            if i == 0:
                # Головной сегмент движется случайным образом
                positions[frame, i, 0] = positions[frame-1, i, 0] + (np.random.rand() - 0.5) * 0.5
                positions[frame, i, 1] = positions[frame-1, i, 1] + (np.random.rand() - 0.5) * 0.5
            else:
                # Остальные сегменты следуют за предыдущими
                positions[frame, i, 0] = positions[frame-1, i-1, 0] + segment_length
                positions[frame, i, 1] = positions[frame-1, i-1, 1] + (M.v[i][frame] - 1) * 0.1

    return positions, M.t

def create_plotly_animation(positions, times):
    num_frames, num_segments, _ = positions.shape

    # Создание фигуры Plotly
    fig = go.Figure(
        layout=go.Layout(
            xaxis=dict(range=[-10, 20], autorange=False),
            yaxis=dict(range=[-10, 20], autorange=False),
            title="Движение червя",
            showlegend=False,
            width=800,
            height=600
        )
    )

    # Добавление кадров
    frames = []
    for frame in range(num_frames):
        x = positions[frame, :, 0]
        y = positions[frame, :, 1]
        frames.append(go.Frame(data=[go.Scatter(x=x, y=y, mode='lines+markers',
                                               marker=dict(size=10, color='blue'),
                                               line=dict(width=2, color='blue'))],
                               name=f'frame{frame}',
                               layout=go.Layout(title_text=f'Время: {times[frame]:.1f} ms')))

    fig.frames = frames

    # Добавление начального состояния
    initial_x = positions[0, :, 0]
    initial_y = positions[0, :, 1]
    fig.add_trace(go.Scatter(x=initial_x, y=initial_y, mode='lines+markers',
                             marker=dict(size=10, color='blue'),
                             line=dict(width=2, color='blue')))

    # Настройка анимации
    fig.update_layout(
        updatemenus=[
            dict(
                type="buttons",
                buttons=[
                    dict(label="▶️ Запустить",
                         method="animate",
                         args=[None, {"frame": {"duration": 100, "redraw": True},
                                      "fromcurrent": True, "transition": {"duration": 0}}]),
                    dict(label="⏹️ Остановить",
                         method="animate",
                         args=[[None], {"frame": {"duration": 0, "redraw": False},
                                        "mode": "immediate",
                                        "transition": {"duration": 0}}])
                ],
                showactive=False,
                x=0,
                y=1.05
            )
        ]
    )

    return fig

def main():
    st.title("Симуляция движения червя с использованием Streamlit и Brian2")
    st.write("""
    Это веб-приложение выполняет симуляцию нейронной активности и соответствующего движения червя.
    Используется **Plotly** для отображения интерактивной анимации.
    """)

    st.sidebar.header("Параметры Симуляции")
    num_neurons = st.sidebar.slider('Количество сегментов червя', min_value=5, max_value=20, value=12)
    sim_time_ms = st.sidebar.slider('Время симуляции (мс)', min_value=50, max_value=500, value=100, step=10)

    if st.button("Запустить симуляцию"):
        with st.spinner('Выполнение симуляции и создание анимации...'):
            try:
                positions, times = run_simulation(num_neurons=num_neurons, sim_time=sim_time_ms*ms)
                fig = create_plotly_animation(positions, times)
                st.plotly_chart(fig, use_container_width=True)
                st.success('Симуляция завершена!')
            except Exception as e:
                st.error(f"Произошла ошибка: {e}")

    st.write("""
    ### Описание
    - **Количество сегментов червя**: устанавливается пользователем (по умолчанию 12)
    - **Длина каждого сегмента**: 0.5 единиц
    - **Время симуляции**: устанавливается пользователем (по умолчанию 100 мс)
    - **Модель нейронов**: Простая модель с параметрами I и tau, влияющими на потенциал v.
    
    ### Как это работает
    1. **Симуляция нейронов**: Создается группа нейронов с разными токами и временем релаксации.
    2. **Механика червя**: Каждому сегменту соответствует нейрон. Активность нейронов влияет на движение сегментов.
    3. **Анимация**: Полученные позиции сегментов визуализируются в интерактивной анимации движения червя.
    """)

if __name__ == "__main__":
    main()