import streamlit as st
import numpy as np
from typing import List, Tuple
from transformers import pipeline
import matplotlib.pyplot as plt
import time

# Constants
WIDTH, HEIGHT = 600, 600
ELASTICITY = 0.3
DAMPING = 0.7

# Create touch points
num_points = 20
touch_points: List[Tuple[float, float]] = [(x, y) for x in range(50, WIDTH-50, int((WIDTH-100)/(num_points-1))) for y in range(50, HEIGHT-50, int((HEIGHT-100)/(num_points-1)))]
original_points = touch_points.copy()
velocities: List[Tuple[float, float]] = [(0.0, 0.0)] * len(touch_points)
is_affected: List[bool] = [False] * len(touch_points)

# Create pain and pleasure points
pain_points = [(100, 100), (500, 500)]
pleasure_points = [(300, 300), (200, 400)]

# Set up the Hugging Face pipeline
@st.cache_resource
def load_model():
    return pipeline('text-generation', model='gpt2')

text_generator = load_model()

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

# Create a Streamlit container for the touch simulation
touch_container = st.container()

def update_points():
    global touch_points, velocities, is_affected
    
    # Apply spring force
    for i, (x, y) in enumerate(touch_points):
        force_x = (original_points[i][0] - x) * ELASTICITY
        force_y = (original_points[i][1] - y) * ELASTICITY
        velocities[i] = (velocities[i][0] + force_x, velocities[i][1] + force_y)
    
    # Apply damping
    for i, (vx, vy) in enumerate(velocities):
        velocities[i] = (vx * DAMPING, vy * DAMPING)
    
    # Update position
    for i, (x, y) in enumerate(touch_points):
        vx, vy = velocities[i]
        touch_points[i] = (x + vx, y + vy)
    
    # Reset affected flags
    is_affected = [False] * len(touch_points)

def calculate_sensation(x, y, pressure, duration):
    sensation = 0
    for px, py in pain_points:
        distance = np.sqrt((x - px)**2 + (y - py)**2)
        if distance < 50:
            sensation -= (50 - distance) * pressure * (duration ** 1.5) / 50
    for px, py in pleasure_points:
        distance = np.sqrt((x - px)**2 + (y - py)**2)
        if distance < 50:
            sensation += (50 - distance) * pressure / 50
    return sensation

def on_tap(x, y, pressure, duration):
    global touch_points, velocities, is_affected
    
    for i, (tx, ty) in enumerate(touch_points):
        distance = ((tx - x)**2 + (ty - y)**2)**0.5
        if distance < 30:
            force_x = (tx - x) / distance
            force_y = (ty - y) / distance
            velocities[i] = (velocities[i][0] - force_x * 10 * pressure, velocities[i][1] - force_y * 10 * pressure)
            is_affected[i] = True
    
    sensation = calculate_sensation(x, y, pressure, duration)
    
    # Generate a description of the touch
    st.write(f"Touch at ({x:.2f}, {y:.2f}) with pressure {pressure:.2f} for {duration:.2f} seconds")
    st.write(f"Sensation: {sensation:.2f}")
    
    prompt = f"The user touched the screen at ({x:.2f}, {y:.2f}) with a pressure of {pressure:.2f} for {duration:.2f} seconds, resulting in a sensation of {sensation:.2f}. Describe the experience:"
    text = text_generator(prompt, max_length=100, num_return_sequences=1, do_sample=True, top_k=50, top_p=0.95, num_beams=1)[0]['generated_text']
    st.write(text)
    
    update_points()

# Initialize session state
if 'x' not in st.session_state:
    st.session_state.x = 0
if 'y' not in st.session_state:
    st.session_state.y = 0
if 'pressure' not in st.session_state:
    st.session_state.pressure = 1.0
if 'duration' not in st.session_state:
    st.session_state.duration = 0.0
if 'touch_start_time' not in st.session_state:
    st.session_state.touch_start_time = None

# Main app logic
fig, ax = plt.subplots(figsize=(6, 6))
ax.set_xlim(0, WIDTH)
ax.set_ylim(0, HEIGHT)

for i, (x, y) in enumerate(touch_points):
    color = "red" if is_affected[i] else "navy"
    ax.add_artist(plt.Circle((x, y), 5, color=color, alpha=0.5))

for x, y in pain_points:
    ax.add_artist(plt.Circle((x, y), 10, color="red", alpha=0.3))

for x, y in pleasure_points:
    ax.add_artist(plt.Circle((x, y), 10, color="green", alpha=0.3))

touch_container.pyplot(fig)

col1, col2 = st.columns(2)
with col1:
    st.session_state.x = st.slider("X coordinate", 0, WIDTH, st.session_state.x)
    st.session_state.y = st.slider("Y coordinate", 0, HEIGHT, st.session_state.y)
with col2:
    st.session_state.pressure = st.slider("Pressure", 0.1, 2.0, st.session_state.pressure)

if touch_container.button("Start Touch"):
    st.session_state.touch_start_time = time.time()

if touch_container.button("End Touch"):
    if st.session_state.touch_start_time is not None:
        st.session_state.duration = time.time() - st.session_state.touch_start_time
        on_tap(st.session_state.x, st.session_state.y, st.session_state.pressure, st.session_state.duration)
        st.session_state.touch_start_time = None
        st.session_state.duration = 0.0

update_points()