# 【功能描述】
# 本项目基于 SMPL 模型，通过随机生成或从库中加载参数，分析人体姿态与压力分布的关系。
# 主要功能包括：
# 1. 随机生成人体参数（包括 betas 和 pose）
# 2. 从预定义库中加载已保存的参数
# 3. 利用 SMPL 模型生成人体网格和压力分布
# 4. 可视化展示生成的人体模型和压力分布
# 5. 允许用户保存新生成的参数到库中【新增功能】

# streamlit run new_app.py --server.port 8501

import streamlit as st
import torch
import numpy as np
import plotly.graph_objects as go
import os
import matplotlib.pyplot as plt

from sample_utils import PoseSampler, sample_beta, sample_transl4pp
from generate_utils import PressureGenerator

# --- 1. 初始化 ---
st.set_page_config(page_title="SMPL2Pressure", layout="wide")

st.markdown("""
    <style>
    .block-container { padding-top: 3rem; padding-bottom: 0rem; }
    #MainMenu {visibility: hidden;}
    footer {visibility: hidden;}
    [data-testid="stMetric"] { margin-top: -1.1rem; margin-bottom: -1rem; }
    </style>
    """, unsafe_allow_html=True)

@st.cache_resource
def init_models(_device):
    p_sampler = PoseSampler(device=_device, dataset='pp')
    p_gen = PressureGenerator(device=_device)
    return p_sampler, p_gen

DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
pose_sampler, generator = init_models(DEVICE)

# 初始化 Session State
if 'betas' not in st.session_state:
    st.session_state.betas = torch.zeros((1, 10)).to(DEVICE)
    st.session_state.pose = torch.zeros((1, 72)).to(DEVICE)
    st.session_state.transl = sample_transl4pp(batch_size=1, device=DEVICE)
    st.session_state.run_trigger = 0
    st.session_state.selected_pose_idx = 0

# --- 2. 辅助函数 ---
@st.cache_data(show_spinner=False)
def compute_single_result(_beta, _pose, _transl, transfer=False, trigger=0):
    # 确保输入是 tensor 且在正确的设备上
    pmap = generator.generate(betas=_beta, transl=_transl, poses=_pose, transfer=transfer)
    pmap = pmap.flip(1)
    with torch.no_grad():
        output = generator.smpl_model(
            betas=_beta, 
            global_orient=_pose[:, :3], 
            body_pose=_pose[:, 3:], 
            transl=_transl
        )
        verts = output.vertices[0].cpu().numpy()
        faces = generator.smpl_model.faces

    if transfer:
        verts[:, 1] = 1.80 - verts[:, 1]
        verts[:, 2] = -verts[:, 2]
    
    return verts, faces, pmap.squeeze().cpu().numpy()

# --- 3. 侧边栏 ---
with st.sidebar:
    st.title("🎛️ Controls")
    mode = st.radio("Input Source", ["Random", "Library"], horizontal=True)
    st.divider()
    if mode == "Library":
        st.info("Library mode: Select a pose from the gallery on the right.")
    else:
        st.success("Random mode: Generate random parameters.")

# --- 4. 主界面标题与顶部控制栏 ---
head_c1, head_c2, head_c3 = st.columns([3, 1, 1])
with head_c1:
    st.subheader("🔬 Pressure Map Synthesis via SMPL Model")

# --- 5. 模式逻辑分流 ---

if mode == "Random":
    # --- RANDOM 模式控制按钮 ---
    with head_c2:
        if st.button("🔄 Generate Random", type="primary", use_container_width=True):
            st.session_state.betas = sample_beta(batch_size=1, device=DEVICE)
            st.session_state.pose = pose_sampler.sample(batch_size=1)
            st.session_state.transl = sample_transl4pp(batch_size=1, device=DEVICE)
            st.session_state.run_trigger += 100

    with head_c3:
        if st.button("Save to Library", type="secondary", use_container_width=True):
            try:
                save_dir = "static_stats"
                os.makedirs(save_dir, exist_ok=True)
                existing_files = [f for f in os.listdir(save_dir) if f.startswith("betas_")]
                indices = [int(f.split('_')[1].split('.')[0]) for f in existing_files if '_' in f]
                next_idx = max(indices) + 1 if indices else 1
                np.savez(f"{save_dir}/betas_{next_idx}.npz", betas=st.session_state.betas.cpu().numpy().squeeze(0))
                np.savez(f"{save_dir}/pose_{next_idx}.npz", pose=st.session_state.pose.cpu().numpy().squeeze(0))
                st.toast(f"Saved as index {next_idx}!", icon='✅')
            except Exception as e:
                st.error(f"Save failed: {e}")

    # --- RANDOM 模式渲染 (7:3 布局) ---
    verts, faces, pmap_np = compute_single_result(st.session_state.betas, st.session_state.pose, st.session_state.transl, trigger=st.session_state.run_trigger)
    
    view_c1, view_c2 = st.columns([7.3, 2.7])
    with view_c1:
        st.subheader("🌐 3D Mesh")
        fig_3d = go.Figure(data=[go.Mesh3d(
            x=verts[:, 0], y=verts[:, 1], z=verts[:, 2],
            i=faces[:, 0], j=faces[:, 1], k=faces[:, 2],
            color='LightBlue',
            opacity=1.0,
            flatshading=False,
            lighting=dict(ambient=0.5, diffuse=0.9, specular=0.5, roughness=0.6),
            lightposition=dict(x=100, y=200, z=150)
        )])
        
        fig_3d.update_layout(
            scene=dict(
                aspectmode='data',
                xaxis_visible=False, yaxis_visible=False, zaxis_visible=False,
                camera=dict(
                    eye=dict(x=0, y=0, z=2.0),     # Positioned high on Z-axis
                    up=dict(x=0, y=1, z=0),        # Y-axis points "up" on the screen
                    center=dict(x=0., y=0., z=0.)     # Looking at the origin
                )
            ),
            height=720, 
            margin=dict(l=0, r=0, b=0, t=0),
            paper_bgcolor="white",
        )
        st.plotly_chart(fig_3d, use_container_width=True)

    with view_c2:
        m_c1, m_c2 = st.columns(2)
        m_c1.metric("Peak Pressure", f"{pmap_np.max():.2f}")
        m_c2.metric("Contact Pixels", f"{(pmap_np > 0.5).sum()}")

        # with st.container():
        # aspect='equal' ensures no distortion
        fig_2d, ax = plt.subplots(figsize=(5, 8)) 
        im = ax.imshow(pmap_np, cmap='viridis', origin='lower', aspect='equal')
        plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
        ax.axis('off')
        # Use container width helps alignment
        st.pyplot(fig_2d, use_container_width=True)

else:
    # --- LIBRARY 模式 ---
    st.write("### 1️⃣ Step: Select Human Pose")
    pose_cols = st.columns(8)
    for i in range(8):
        img_path = f"src/static/fig/pose_{i}.png"
        with pose_cols[i]:
            if os.path.exists(img_path):
                st.image(img_path, use_container_width=True)
            # 点击按钮更新全局索引
            if st.button(f"Pose {i}", key=f"btn_p_{i}", use_container_width=True, type="primary" if st.session_state.selected_pose_idx == i else "secondary"):
                st.session_state.selected_pose_idx = i
                st.session_state.run_trigger += 1
                st.rerun()
    
    # import pdb; pdb.set_trace()

    # 加载选中的数据
    try:
        # p_path = f"src/static/poses/POSE_{st.session_state.selected_pose_idx}.npz"
        p_path = f"src/static_stats/pose_{st.session_state.selected_pose_idx}.npz"
        p_data = np.load(p_path)['pose']
        current_pose = torch.from_numpy(p_data).float().unsqueeze(0).to(DEVICE)
        
        # 加载 5 个 Betas
        beta_tensors = []
        for i in range(5):
            b_path = f"src/static/betas/MALE_BETA_{i}.npz"
            if os.path.exists(b_path):
                b_data = np.load(b_path)['beta']
                beta_tensors.append(torch.from_numpy(b_data).float().unsqueeze(0).to(DEVICE))
            else:
                test_beta = sample_beta(batch_size=1, device=DEVICE)
                beta_tensors.append(test_beta)

    except Exception as e:
        # import pdb; pdb.set_trace()
        st.error(f"Error loading files: {e}. Check if static/poses/ and static/betas/ exist.")
        st.stop()

    # 统一的平移量
    lib_transl = sample_transl4pp(batch_size=1, device=DEVICE)

    st.divider()
    st.write(f"### 2️⃣ Step: Multi-Body Shape Comparison (Current: Pose {st.session_state.selected_pose_idx})")

    # 并排展示
    cols = st.columns(5)
    for i in range(5):
        with cols[i]:
            # print(f"{beta_tensors[i].shape} {beta_tensors[i].cpu().numpy()}")
            # import pdb; pdb.set_trace()
            v, f, p = compute_single_result(beta_tensors[i], current_pose, lib_transl, transfer=False, trigger=st.session_state.run_trigger + i*1000)
            st.markdown(f"<center><b>Body Type {i}</b></center>", unsafe_allow_html=True)
            
            if not "stats" in p_path:
                # print(v.shape, f.shape, p.shape)
                v[:, 1] = 1.80 - v[:, 1]
                v[:, 2] = -v[:, 2]
                p = torch.from_numpy(p).float().flip(0).cpu().numpy()

            # 3D Mesh
            fig_3d = go.Figure(data=[go.Mesh3d(x=v[:,0], y=v[:,1], z=v[:,2], i=f[:,0], j=f[:,1], k=f[:,2], color='LightBlue')])
            fig_3d.update_layout(scene=dict(aspectmode='data', xaxis_visible=False, yaxis_visible=False, zaxis_visible=False,
                                           camera=dict(eye=dict(x=0, y=0, z=2.2), up=dict(x=0, y=1, z=0))),
                                 height=350, margin=dict(l=0,r=0,b=0,t=0))
            # fig_3d.update_layout(scene=dict(aspectmode='manual', xaxis=dict(range=[-0.2, 1.2], autorange=False), 
            #                                 yaxis=dict(range=[0, 1.95], autorange=False),
            #                                 zaxis=dict(range=[0, 0.5], autorange=True),
            #                                 xaxis_visible=False, yaxis_visible=False, zaxis_visible=False,
            #                                camera=dict(eye=dict(x=0, y=0, z=2.2), up=dict(x=0, y=1, z=0))),
            #                      height=350, margin=dict(l=0,r=0,b=0,t=0))
            st.plotly_chart(fig_3d, use_container_width=True, key=f"mesh_lib_{i}")
            
            # Pressure Map
            fig_2d, ax = plt.subplots()
            ax.imshow(p, cmap='viridis', origin='lower')
            ax.axis('off')
            st.pyplot(fig_2d, use_container_width=True)
            plt.close(fig_2d)
            # st.metric("Peak", f"{p.max():.2f}")