import streamlit as st
import plotly.graph_objects as go
import numpy as np
from sklearn.decomposition import PCA
import time
from sandbox import UniversalSymbolicSandbox

# --- 页面配置 ---
st.set_page_config(layout="wide", page_title="LAWS-ECHO Sandbox", page_icon="🌌")

# 自定义 CSS 打造极客感
st.markdown("""
<style>
    .reportview-container { background: #0e1117; }
    .stMetric value { color: #00ff00; }
</style>
""", unsafe_allow_html=True)

st.title("🌌 LAWS-ECHO: White-Box Geometric Reasoning Sandbox")
st.markdown("基于偏态高斯实体与角算子的可微物理推演引擎。观察不可见潜空间的算子盆地（Operator Basins）涌现。")

# --- 侧边栏：控制面板 ---
with st.sidebar:
    st.header("⚙️ 引擎控制台")
    formula_input = st.text_input("输入物理公式", value="F = m * a", help="支持 SymPy 解析，如 E = 1/2 * m * v**2")
    epochs = st.slider("推演 Epochs", min_value=50, max_value=500, value=200, step=10)
    start_btn = st.button("🚀 启动几何推演", type="primary")
    
    st.divider()
    st.markdown("### 实时监控指标")
    metric_loss = st.empty()
    metric_angle = st.empty()
    status_text = st.empty()

# --- 主界面：潜空间可视化 ---
main_plot = st.empty()

def render_pca(mu_dict, alpha_dict, epoch, loss):
    entities = list(mu_dict.keys())
    mu_matrix = np.array([mu_dict[k] for k in entities])
    
    # 使用 PCA 将 16 维降至 2 维进行可视化
    pca = PCA(n_components=2)
    try:
        mu_2d = pca.fit_transform(mu_matrix)
    except:
        mu_2d = np.zeros((len(entities), 2))
        
    fig = go.Figure()

    # 实体节点 (mu)
    colors = ['#00d2ff', '#ff0055', '#00ffaa', '#ffcc00', '#cc00ff']
    fig.add_trace(go.Scatter(
        x=mu_2d[:, 0], y=mu_2d[:, 1],
        mode='markers+text',
        marker=dict(size=18, color=colors[:len(entities)], line=dict(width=2, color='white')),
        text=entities,
        textposition="top center",
        textfont=dict(size=16, color="white"),
        name="Entity Center (μ)"
    ))

    # 各向异性趋势向量 (alpha)
    for i, entity in enumerate(entities):
        # 将 alpha 投影到 PCA 空间并按比例缩放以优化显示
        alpha_2d = pca.transform((mu_dict[entity] + alpha_dict[entity] * 2.0).reshape(1, -1))[0]
        fig.add_trace(go.Scatter(
            x=[mu_2d[i, 0], alpha_2d[0]],
            y=[mu_2d[i, 1], alpha_2d[1]],
            mode='lines',
            line=dict(width=3, color='rgba(255, 255, 255, 0.4)', dash='dot'),
            showlegend=False
        ))

    fig.update_layout(
        title=f"Latent Space Evolution | Epoch: {epoch}/{epochs}",
        plot_bgcolor='#0e1117',
        paper_bgcolor='#0e1117',
        font=dict(color='white'),
        xaxis=dict(showgrid=True, gridcolor='#333333', zeroline=True, zerolinecolor='#666666'),
        yaxis=dict(showgrid=True, gridcolor='#333333', zeroline=True, zerolinecolor='#666666'),
        height=650,
        margin=dict(l=10, r=10, t=40, b=10),
        showlegend=False
    )
    return fig

# --- 核心调度逻辑 ---
if start_btn:
    try:
        status_text.info("正在初始化符号抽象语法树 (AST)...")
        sandbox = UniversalSymbolicSandbox(formula_input)
        status_text.success("沙盒环境与几何实体分配就绪！")
        
        # 预热渲染一张初始图
        mu_init, alpha_init = sandbox.get_geometric_state()
        main_plot.plotly_chart(render_pca(mu_init, alpha_init, 0, 0.0), use_container_width=True)
        time.sleep(1.0)
        
        # 训练循环
        for ep in range(1, epochs + 1):
            loss = sandbox.step()
            
            # 为了保证前端不卡顿，每 3 个 Epoch 渲染一帧
            if ep % 3 == 0 or ep == epochs:
                mu_state, alpha_state = sandbox.get_geometric_state()
                current_angle = sandbox.compute_angles()
                
                # 更新图表
                fig = render_pca(mu_state, alpha_state, ep, loss)
                main_plot.plotly_chart(fig, use_container_width=True)
                
                # 更新指标看板
                with metric_loss.container():
                    st.metric(label="MSE Loss", value=f"{loss:.6f}")
                with metric_angle.container():
                    # 判断算子落在哪个盆地
                    basin = "MUL Basin (≈178°)" if current_angle > 140 else "ADD Basin (≈90°)" if current_angle > 70 else "Searching..."
                    st.metric(label="Primary Operator Angle", value=f"{current_angle:.1f}°", delta=basin)
                    
        status_text.success(f"✅ 演化完成！系统已通过纯标量监督，在 16 维潜空间内自动构建了 {formula_input} 的几何拓扑结构。")
        
    except Exception as e:
        status_text.error(f"解析错误: 请确保公式格式正确且变量字母标准。详情: {e}")
else:
    # 默认空状态展示
    st.info("👈 请在左侧面板输入物理公式并点击启动。")