import plotly.graph_objects as go
import numpy as np

# Load the data
data = {
  "object_status": ["Yes", "No"],
  "habit_status": ["Yes", "No"],
  "renss_scores": [
    [35.77, 35.83],
    [35.76, 35.67]
  ],
  "counts": [
    [74, 36],
    [21, 9]
  ]
}

# Create the heatmap
fig = go.Figure(data=go.Heatmap(
    z=data['renss_scores'],
    x=data['object_status'],
    y=data['habit_status'],
    colorscale='Blues',
    text=[[f"{data['renss_scores'][i][j]:.2f} (n={data['counts'][i][j]})" 
           for j in range(len(data['object_status']))] 
          for i in range(len(data['habit_status']))],
    texttemplate="%{text}",
    textfont={"size":12},
    hoverinfo='text',
    colorbar=dict(title='RENSS Score')
))

# Update layout
fig.update_layout(
    title='RENSS Scores by Habit & Object Status'
)

# Update axes - removed cliponaxis property
fig.update_xaxes(title='Object Status')
fig.update_yaxes(title='Habit Status')

# Save the figure
fig.write_image("renss_heatmap.png")

fig