import streamlit as st
from streamlit.runtime.uploaded_file_manager import UploadedFile
import pandas as pd
import numpy as np
from pose_format import Pose
from pose_format.pose_visualizer import PoseVisualizer
from pathlib import Path
from pyzstd import decompress
from PIL import Image
import cv2
import mediapipe as mp

mp_holistic = mp.solutions.holistic
FACEMESH_CONTOURS_POINTS = [
    str(p)
    for p in sorted(
        set([p for p_tup in list(mp_holistic.FACEMESH_CONTOURS) for p in p_tup])
    )
]


def pose_normalization_info(pose_header):
    if pose_header.components[0].name == "POSE_LANDMARKS":
        return pose_header.normalization_info(
            p1=("POSE_LANDMARKS", "RIGHT_SHOULDER"),
            p2=("POSE_LANDMARKS", "LEFT_SHOULDER"),
        )

    if pose_header.components[0].name == "BODY_135":
        return pose_header.normalization_info(
            p1=("BODY_135", "RShoulder"), p2=("BODY_135", "LShoulder")
        )

    if pose_header.components[0].name == "pose_keypoints_2d":
        return pose_header.normalization_info(
            p1=("pose_keypoints_2d", "RShoulder"), p2=("pose_keypoints_2d", "LShoulder")
        )


def pose_hide_legs(pose):
    if pose.header.components[0].name == "POSE_LANDMARKS":
        point_names = ["KNEE", "ANKLE", "HEEL", "FOOT_INDEX"]
        # pylint: disable=protected-access
        points = [
            pose.header._get_point_index("POSE_LANDMARKS", side + "_" + n)
            for n in point_names
            for side in ["LEFT", "RIGHT"]
        ]
        pose.body.confidence[:, :, points] = 0
        pose.body.data[:, :, points, :] = 0
        return pose
    else:
        raise ValueError("Unknown pose header schema for hiding legs")


# @st.cache_data(hash_funcs={UploadedFile: lambda p: str(p.name)})
def load_pose(uploaded_file: UploadedFile) -> Pose:

    # with input_path.open("rb") as f_in:
    if uploaded_file.name.endswith(".zst"):
        return Pose.read(decompress(uploaded_file.read()))
    else:
        return Pose.read(uploaded_file.read())


@st.cache_data(hash_funcs={Pose: lambda p: np.array(p.body.data)})
def get_pose_frames(pose: Pose, transparency: bool = False):
    v = PoseVisualizer(pose)
    frames = [frame_data for frame_data in v.draw()]

    if transparency:
        cv_code = v.cv2.COLOR_BGR2RGBA
    else:
        cv_code = v.cv2.COLOR_BGR2RGB
    images = [Image.fromarray(v.cv2.cvtColor(frame, cv_code)) for frame in frames]
    return frames, images


def get_pose_gif(pose: Pose, step: int = 1, fps: int = None):
    if fps is not None:
        pose.body.fps = fps
    v = PoseVisualizer(pose)
    frames = [frame_data for frame_data in v.draw()]
    frames = frames[::step]
    return v.save_gif(None, frames=frames)


st.write("# Pose-format explorer")
st.write(
    "`pose-format` is a toolkit/library for 'handling, manipulation, and visualization of poses'. See [The documentation](https://pose-format.readthedocs.io/en/latest/)"
)
st.write(
    "I made this app to help me visualize and understand the format, including different 'components' and 'points', and what they are named."
)
uploaded_file = st.file_uploader("Upload a .pose file", type=[".pose", ".pose.zst"])


if uploaded_file is not None:
    with st.spinner(f"Loading {uploaded_file.name}"):
        pose = load_pose(uploaded_file)
        frames, images = get_pose_frames(pose=pose)
    st.success("Done loading!")
    
    st.write("### File Info")
    with st.expander(f"Show full Pose-format header from {uploaded_file.name}"):
        st.write(pose.header)

    st.write(f"### Selection")
    component_selection = st.radio(
        "How to select components?", options=["manual", "signclip"]
    )

    if component_selection == "manual":
        chosen_component_names = st.multiselect(
            "Select components to visualize", options=[c.name for c in pose.header.components]
        )
        if chosen_component_names:
            pose = pose.get_components(chosen_component_names)

    elif component_selection == "signclip":
        st.write("Selected landmarks used for SignCLIP.")
        pose = pose.get_components(
            ["POSE_LANDMARKS", "FACE_LANDMARKS", "LEFT_HAND_LANDMARKS", "RIGHT_HAND_LANDMARKS"]
        )
        pose = pose_hide_legs(pose)

    # Filter button logic
    if st.button("Filter Components/Points"):
        st.write("### Filtered .pose file")
        with st.expander("Show header"):
            st.write(pose.header)
        with st.expander("Show body"):
            st.write(pose.body)
        
        pose_file_out = Path(uploaded_file.name).with_suffix(".pose")
        with pose_file_out.open("wb") as f:
            pose.write(f)

        with pose_file_out.open("rb") as f:
            st.download_button("Download Filtered Pose", f, file_name=pose_file_out.name)

    # Visualization button logic
    if st.button("Visualize"):
        st.write("### Visualization")
        step = st.select_slider("Step value to select every nth image", list(range(1, len(frames))), value=1)
        fps = st.slider("FPS for visualization", min_value=1.0, max_value=pose.body.fps, value=pose.body.fps)
        st.image(get_pose_gif(pose=pose, step=step, fps=fps))


            


            # st.write(pose.body.data.shape)

            # st.write(visualize_pose(pose=pose)) # bunch of ndarrays
            # st.write([Image.fromarray(v.cv2.cvtColor(frame, cv_code)) for frame in frames])

            # for i, image in enumerate(images[::n]):
            #     print(f"i={i}")
            #     st.image(image=image, width=width)