import plotly.graph_objects as go
import numpy as np
import pandas as pd
from scipy import stats

# Create dataframe from the provided data
data = {
    "Participant": ["User 1", "User 10", "User 11", "User 12", "User 13", "User 14", "User 15",
                   "User 16", "User 17", "User 18", "User 19", "User 2", "User 20", "User 3",
                   "User 4", "User 5", "User 6", "User 7", "User 8", "User 9"],
    "Habit_Rate": [0.857143, 0.571429, 0.571429, 0.714286, 0.714286, 0.714286, 0.857143,
                  0.571429, 0.714286, 0.714286, 0.857143, 0.714286, 0.428571, 0.428571,
                  1.000000, 0.714286, 0.428571, 0.428571, 0.857143, 0.714286],
    "Object_Rate": [1.000000, 0.857143, 0.857143, 1.000000, 0.428571, 0.714286, 0.714286,
                   0.571429, 0.714286, 0.857143, 0.714286, 0.857143, 0.857143, 0.857143,
                   0.857143, 0.428571, 0.857143, 0.714286, 1.000000, 0.857143]
}

df = pd.DataFrame(data)

# Calculate correlation coefficient
correlation = np.corrcoef(df['Object_Rate'], df['Habit_Rate'])[0, 1]

# Create the scatter plot
fig = go.Figure()

# Add scatter points with proper color gradient
fig.add_trace(go.Scatter(
    x=df['Object_Rate'],
    y=df['Habit_Rate'],
    mode='markers',  # Remove text from markers, we'll add it separately
    marker=dict(
        color=df['Habit_Rate'],
        colorscale=[[0, '#C8DCFD'], [1, '#1F77B4']],  # Light to dark blue
        size=12,
        opacity=1.0,  # Fully opaque
        showscale=False  # Remove the color scale for clarity
    ),
    hovertemplate='<b>%{text}</b><br>Habit: %{y:.1%}<br>Object: %{x:.1%}<extra></extra>',
    text=df['Participant'],
    cliponaxis=False
))

# Add text labels separately with offset
fig.add_trace(go.Scatter(
    x=df['Object_Rate'],
    y=df['Habit_Rate'],
    mode='text',
    text=df['Participant'],
    textposition='top center',
    textfont=dict(size=9),
    hoverinfo='skip',
    showlegend=False,
    cliponaxis=False,
    dy=10  # Offset the text upward
))

# Add regression line
slope, intercept, r_value, p_value, std_err = stats.linregress(df['Object_Rate'], df['Habit_Rate'])
x_range = np.linspace(0.4, 1.0, 100)
y_range = slope * x_range + intercept

fig.add_trace(go.Scatter(
    x=x_range,
    y=y_range,
    mode='lines',
    line=dict(color='#1FB8CD', width=1.5, dash='dash'),  # Thinner, dashed line
    name='Trend',
    hoverinfo='skip',
    showlegend=False
))

# Update layout
title = f"Habit & Object Placement (r={correlation:.2f})"

fig.update_layout(
    title=title,
    plot_bgcolor='white'
)

# Format axes - Focus on the actual data range
fig.update_xaxes(
    title="Object Rate (%)",
    tickformat='.0%',
    range=[0.4, 1.0],
    zeroline=False
)

fig.update_yaxes(
    title="Habit Rate (%)",
    tickformat='.0%',
    range=[0.4, 1.0],
    zeroline=False
)

# Save the chart
fig.write_image("habit_object_relation.png")
fig.show()