import streamlit as st
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.backends.backend_agg import FigureCanvasAgg
from streamlit_drawable_canvas import st_canvas
import time
from PIL import Image
import io
from transformers import AutoModelForCausalLM, AutoTokenizer

# Constants
WIDTH, HEIGHT = 800, 400
AVATAR_WIDTH, AVATAR_HEIGHT = 300, 400

# Set up DialoGPT model
@st.cache_resource
def load_model():
    tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-small")
    model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-small")
    return tokenizer, model

tokenizer, model = load_model()

# Create sensation map for the avatar
def create_sensation_map(width, height):
    sensation_map = np.zeros((height, width, 3))  # RGB channels for pain, pleasure, and neutral
    for y in range(height):
        for x in range(width):
            # Base sensation
            base = np.sin(x/15) * np.cos(y/15) * 0.5 + np.random.normal(0, 0.1)
            
            # Pain regions (red channel)
            pain = np.exp(-((x-75)**2 + (y-100)**2) / 2000) + np.exp(-((x-225)**2 + (y-300)**2) / 2000)
            
            # Pleasure regions (green channel)
            pleasure = np.exp(-((x-150)**2 + (y-200)**2) / 2000) + np.exp(-((x-75)**2 + (y-300)**2) / 2000)
            
            # Neutral sensation (blue channel)
            neutral = 1 - (pain + pleasure)
            
            sensation_map[y, x] = [pain, pleasure, neutral]
    
    return sensation_map

avatar_sensation_map = create_sensation_map(AVATAR_WIDTH, AVATAR_HEIGHT)

# Streamlit app
st.title("Advanced Humanoid Touch Simulation")

# Create two columns
col1, col2 = st.columns(2)

# Avatar column
with col1:
    st.subheader("Humanoid Avatar")
    avatar_fig, avatar_ax = plt.subplots(figsize=(4, 6))
    avatar_ax.imshow(avatar_sensation_map)
    avatar_ax.axis('off')
    st.pyplot(avatar_fig)

# Touch interface column
with col2:
    st.subheader("Touch Interface")
    touch_fig, touch_ax = plt.subplots(figsize=(4, 6))
    touch_ax.add_patch(plt.Rectangle((0, 0), AVATAR_WIDTH, AVATAR_HEIGHT, fill=False))
    touch_ax.set_xlim(0, AVATAR_WIDTH)
    touch_ax.set_ylim(0, AVATAR_HEIGHT)
    touch_ax.axis('off')
    
    # Convert matplotlib figure to Image
    canvas = FigureCanvasAgg(touch_fig)
    canvas.draw()
    buf = io.BytesIO()
    plt.savefig(buf, format='png')
    buf.seek(0)
    img = Image.open(buf)

    # Use streamlit-drawable-canvas for interaction
    canvas_result = st_canvas(
        fill_color="rgba(255, 165, 0, 0.3)",
        stroke_width=3,
        stroke_color="#e00",
        background_color="#eee",
        background_image=img,
        update_streamlit=True,
        height=AVATAR_HEIGHT,
        width=AVATAR_WIDTH,
        drawing_mode="point",
        point_display_radius=5,
        key="canvas",
    )

def calculate_sensation(x, y, pressure, duration):
    sensation = avatar_sensation_map[int(y), int(x)]
    modified_sensation = sensation * pressure * (1 + np.log(duration + 1))
    return modified_sensation

def generate_description(x, y, pressure, duration, pain, pleasure, neutral):
    prompt = f"Human: Describe the sensation when touched at ({x:.1f}, {y:.1f}) with pressure {pressure:.2f} for {duration:.2f} seconds. Pain: {pain:.2f}, Pleasure: {pleasure:.2f}, Neutral: {neutral:.2f}.\nAvatar:"
    
    input_ids = tokenizer.encode(prompt, return_tensors="pt")
    output = model.generate(input_ids, max_length=150, num_return_sequences=1, no_repeat_ngram_size=2, top_k=50, top_p=0.95, temperature=0.7)
    
    return tokenizer.decode(output[0], skip_special_tokens=True).split("Avatar: ")[-1].strip()

# Initialize session state
if 'touch_start_time' not in st.session_state:
    st.session_state.touch_start_time = None
if 'last_touch_position' not in st.session_state:
    st.session_state.last_touch_position = None

# Handle touch events
if canvas_result.json_data is not None:
    objects = canvas_result.json_data["objects"]
    if len(objects) > 0:
        last_object = objects[-1]
        current_position = (last_object["left"], last_object["top"])
        
        if st.session_state.touch_start_time is None:
            st.session_state.touch_start_time = time.time()
            st.session_state.last_touch_position = current_position
        else:
            # Calculate pressure based on movement
            if st.session_state.last_touch_position is not None:
                dx = current_position[0] - st.session_state.last_touch_position[0]
                dy = current_position[1] - st.session_state.last_touch_position[1]
                distance = np.sqrt(dx**2 + dy**2)
                pressure = 1 + distance / 10  # Adjust this formula as needed
            else:
                pressure = 1.0

            duration = time.time() - st.session_state.touch_start_time
            x, y = current_position
            sensation = calculate_sensation(x, y, pressure, duration)
            pain, pleasure, neutral = sensation

            description = generate_description(x, y, pressure, duration, pain, pleasure, neutral)
            
            st.write(f"Touch at ({x:.1f}, {y:.1f}) with pressure {pressure:.2f} for {duration:.2f} seconds")
            st.write(f"Pain: {pain:.2f}, Pleasure: {pleasure:.2f}, Neutral: {neutral:.2f}")
            st.write("Avatar's response:")
            st.write(description)
            
            st.session_state.last_touch_position = current_position
    else:
        st.session_state.touch_start_time = None
        st.session_state.last_touch_position = None

st.write("Click and drag on the touch interface to simulate touching the avatar.")
st.write("The avatar's sensation map shows pain (red), pleasure (green), and neutral (blue) areas.")