import streamlit as st
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image, ImageDraw, ImageFont
import time
from transformers import AutoModelForCausalLM, AutoTokenizer
import io

# Constants
AVATAR_WIDTH, AVATAR_HEIGHT = 400, 800

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

tokenizer, model = load_model()

# Simulated Sensor Classes
class Sensors:
    @staticmethod
    def measure_pressure(base_sensitivity, duration):
        return base_sensitivity * (1 - np.exp(-duration / 2))

    @staticmethod
    def measure_temperature(base_temp, duration):
        return base_temp + 5 * (1 - np.exp(-duration / 3))

    @staticmethod
    def measure_texture(x, y):
        textures = ["smooth", "rough", "bumpy", "silky", "grainy", "soft", "hard", "moist", "dry", "fuzzy"]
        return textures[hash((x, y)) % len(textures)]

    @staticmethod
    def measure_em_field(x, y):
        return np.sin(x/50) * np.cos(y/50) * 10

# Create more detailed sensation map for the avatar
def create_sensation_map(width, height):
    sensation_map = np.zeros((height, width, 8))  # pain, pleasure, pressure, temp, texture, em, tickle, itch
    for y in range(height):
        for x in range(width):
            # Head
            if 150 < x < 250 and 50 < y < 200:
                sensation_map[y, x] = [0.7, 0.5, 0.8, 0.6, 0.9, 0.9, 0.3, 0.4]
            # Face
            elif 160 < x < 240 and 80 < y < 180:
                sensation_map[y, x] = [0.9, 0.7, 1.0, 0.8, 1.0, 1.0, 0.5, 0.6]
            # Neck
            elif 175 < x < 225 and 200 < y < 250:
                sensation_map[y, x] = [0.8, 0.6, 0.9, 0.7, 0.8, 0.8, 0.7, 0.5]
            # Torso
            elif 150 < x < 250 and 250 < y < 500:
                sensation_map[y, x] = [0.5, 0.6, 0.7, 0.8, 0.6, 0.7, 0.5, 0.3]
            # Arms
            elif (100 < x < 150 or 250 < x < 300) and 250 < y < 500:
                sensation_map[y, x] = [0.6, 0.5, 0.9, 0.7, 0.8, 0.6, 0.7, 0.4]
            # Hands
            elif (75 < x < 125 or 275 < x < 325) and 450 < y < 525:
                sensation_map[y, x] = [0.8, 0.7, 1.0, 0.9, 1.0, 0.8, 0.9, 0.7]
            # Legs
            elif 150 < x < 250 and 500 < y < 700:
                sensation_map[y, x] = [0.7, 0.4, 0.8, 0.6, 0.7, 0.5, 0.6, 0.5]
            # Feet
            elif 150 < x < 250 and 700 < y < 800:
                sensation_map[y, x] = [0.9, 0.6, 1.0, 0.8, 0.9, 0.7, 1.0, 0.8]
            else:
                sensation_map[y, x] = [0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1]
    
    return sensation_map

avatar_sensation_map = create_sensation_map(AVATAR_WIDTH, AVATAR_HEIGHT)

# Create more detailed human-like avatar
def create_avatar():
    img = Image.new('RGB', (AVATAR_WIDTH, AVATAR_HEIGHT), color='white')
    draw = ImageDraw.Draw(img)
    
    # Head
    draw.ellipse([150, 50, 250, 200], fill='beige', outline='black')
    # Hair
    draw.polygon([(150, 120), (200, 30), (250, 120)], fill='brown')
    # Eyes
    draw.ellipse([175, 100, 195, 120], fill='white', outline='black')
    draw.ellipse([205, 100, 225, 120], fill='white', outline='black')
    draw.ellipse([182, 107, 188, 113], fill='blue')
    draw.ellipse([212, 107, 218, 113], fill='blue')
    # Nose
    draw.polygon([(200, 130), (190, 150), (210, 150)], fill='beige', outline='black')
    # Mouth
    draw.arc([185, 160, 215, 180], start=0, end=180, fill='red', width=2)
    
    # Neck
    draw.rectangle([175, 200, 225, 250], fill='beige', outline='black')
    
    # Body
    draw.rectangle([150, 250, 250, 500], fill='lightblue', outline='black')
    
    # Arms
    draw.rectangle([100, 250, 150, 500], fill='lightblue', outline='black')
    draw.rectangle([250, 250, 300, 500], fill='lightblue', outline='black')
    
    # Hands
    draw.ellipse([75, 450, 125, 525], fill='beige', outline='black')
    draw.ellipse([275, 450, 325, 525], fill='beige', outline='black')
    
    # Legs
    draw.rectangle([150, 500, 200, 700], fill='navy', outline='black')
    draw.rectangle([200, 500, 250, 700], fill='navy', outline='black')
    
    # Feet
    draw.ellipse([140, 700, 210, 800], fill='beige', outline='black')
    draw.ellipse([190, 700, 260, 800], fill='beige', outline='black')
    
    return img

avatar_image = create_avatar()

# Streamlit app
st.title("Advanced Humanoid Techno-Sensory Simulation")

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

# Avatar display with crosshair
with col1:
    st.subheader("Humanoid Avatar")
    
    # Touch input
    touch_x = st.slider("Touch X coordinate", 0, AVATAR_WIDTH, AVATAR_WIDTH // 2)
    touch_y = st.slider("Touch Y coordinate", 0, AVATAR_HEIGHT, AVATAR_HEIGHT // 2)
    
    # Add crosshair to avatar image
    avatar_with_crosshair = avatar_image.copy()
    draw = ImageDraw.Draw(avatar_with_crosshair)
    draw.line((touch_x - 10, touch_y, touch_x + 10, touch_y), fill="red", width=2)
    draw.line((touch_x, touch_y - 10, touch_x, touch_y + 10), fill="red", width=2)
    
    # Display avatar with crosshair
    st.image(avatar_with_crosshair, use_column_width=True)

# Touch controls and output
with col2:
    st.subheader("Touch Controls")
    
    # Touch duration
    touch_duration = st.slider("Touch duration (seconds)", 0.1, 5.0, 1.0, 0.1)
    
    # Touch pressure
    touch_pressure = st.slider("Touch pressure", 0.1, 2.0, 1.0, 0.1)
    
    if st.button("Apply Touch"):
        sensation = avatar_sensation_map[touch_y, touch_x]
        pain, pleasure, pressure_sens, temp_sens, texture_sens, em_sens, tickle_sens, itch_sens = sensation

        measured_pressure = Sensors.measure_pressure(pressure_sens * touch_pressure, touch_duration)
        measured_temp = Sensors.measure_temperature(37, touch_duration)
        measured_texture = Sensors.measure_texture(touch_x, touch_y)
        measured_em = Sensors.measure_em_field(touch_x, touch_y) * em_sens

        # Calculate overall sensation
        pain_level = pain * measured_pressure
        pleasure_level = pleasure * (measured_temp - 37) / 5
        tickle_level = tickle_sens * (1 - np.exp(-touch_duration / 0.5))
        itch_level = itch_sens * (1 - np.exp(-touch_duration / 1.5))

        st.write(f"Touch applied at ({touch_x}, {touch_y}) for {touch_duration:.1f} seconds")
        st.write(f"Pressure: {measured_pressure:.2f}")
        st.write(f"Temperature: {measured_temp:.2f}°C")
        st.write(f"Texture: {measured_texture}")
        st.write(f"Electromagnetic field: {measured_em:.2f}")
        st.write(f"Pain level: {pain_level:.2f}")
        st.write(f"Pleasure level: {pleasure_level:.2f}")
        st.write(f"Tickle level: {tickle_level:.2f}")
        st.write(f"Itch level: {itch_level:.2f}")

        # Generate description
        prompt = f"""Human: Describe the sensation when touched at ({touch_x}, {touch_y}) for {touch_duration:.1f} seconds with these measurements:
        Pressure: {measured_pressure:.2f}
        Temperature: {measured_temp:.2f}°C
        Texture: {measured_texture}
        Electromagnetic field: {measured_em:.2f}
        Resulting in:
        Pain: {pain_level:.2f}, Pleasure: {pleasure_level:.2f}, Tickle: {tickle_level:.2f}, Itch: {itch_level:.2f}
        Avatar:"""
        
        input_ids = tokenizer.encode(prompt, return_tensors="pt")
        output = model.generate(input_ids, max_length=200, num_return_sequences=1, no_repeat_ngram_size=2, top_k=50, top_p=0.95, temperature=0.7)
        
        response = tokenizer.decode(output[0], skip_special_tokens=True).split("Avatar: ")[-1].strip()
        
        st.write("Avatar's response:")
        st.write(response)

# Visualize sensation map
st.subheader("Sensation Map Visualization")
fig, axs = plt.subplots(2, 4, figsize=(20, 10))
titles = ['Pain', 'Pleasure', 'Pressure', 'Temperature', 'Texture', 'EM Field', 'Tickle', 'Itch']

for i, title in enumerate(titles):
    ax = axs[i // 4, i % 4]
    im = ax.imshow(avatar_sensation_map[:, :, i], cmap='viridis')
    ax.set_title(title)
    fig.colorbar(im, ax=ax)

plt.tight_layout()
st.pyplot(fig)

st.write("The sensation map shows the sensitivity of different body parts to various stimuli. Brighter colors indicate higher sensitivity.")

# Add some context about the avatar's sensory capabilities
st.subheader("Avatar Sensory Capabilities")
st.write("""
This advanced humanoid avatar is equipped with cutting-edge sensory technology:

1. Pressure Sensors: Highly sensitive to touch, with increased sensitivity in hands, feet, and face.
2. Temperature Sensors: Can detect slight changes in temperature, simulating human thermal perception.
3. Texture Analysis: Capable of distinguishing between various textures, from smooth to rough, soft to hard, and more.
4. Electromagnetic Field Detection: Mimics the subtle EM sensitivity some humans report.
5. Pain and Pleasure Processing: Simulates the complex interplay of pain and pleasure responses.
6. Tickle Sensation: Replicates the unique tickle response, which can be pleasurable or uncomfortable.
7. Itch Simulation: Reproduces the sensation of itching, which can be triggered by light touch or prolonged contact.

The avatar's responses are generated using an advanced language model, attempting to describe the sensations in human-like terms. This simulation demonstrates the potential for creating highly responsive and realistic artificial sensory systems.
""")

# Footer
st.write("---")
st.write("Advanced Humanoid Techno-Sensory Simulation v2.0")
st.write("Disclaimer: This is a simulation and does not represent actual human sensory experiences.")