import pandas as pd
import plotly.graph_objects as go
import numpy as np

# Create DataFrame from provided data
data = {
  "scales": ["RENSS_Total", "RENSS_1", "RENSS_2", "RENSS_3", "RENSS_4", "RENSS_5", 
             "RENSS_6", "RENSS_7", "RENSS_8", "RENSS_9", "RENSS_10", "RENSS_11", "RENSS_12"],
  "differences": [-1.55, 0.60, 0.05, -0.10, -0.20, -0.50, -0.95, 0.40, -0.25, 0.75, -0.85, -0.35, -0.15],
  "p_values": [0.196879, 0.248213, 0.885888, 0.873834, 0.508744, 0.252679, 
              0.036223, 0.431250, 0.563367, 0.158081, 0.072703, 0.579946, 0.732539]
}

df = pd.DataFrame(data)

# Define colors and create a category column for the legend
colors = []
categories = []
for p in df["p_values"]:
    if p < 0.05:
        colors.append("#1FB8CD")  # Significant
        categories.append("p < 0.05")
    elif p < 0.1:
        colors.append("#FFC185")  # Marginally significant
        categories.append("p < 0.1")
    else:
        colors.append("#ECEBD5")  # Not significant
        categories.append("p ≥ 0.1")

df["category"] = categories

# Sort data for better visualization
df = df.iloc[::-1]  # Reverse order to have RENSS_Total at the bottom
colors = colors[::-1]  # Reverse colors to match

# Create figure
fig = go.Figure()

# Create separate traces for each significance category to build proper legend
for category, color in [("p < 0.05", "#1FB8CD"), ("p < 0.1", "#FFC185"), ("p ≥ 0.1", "#ECEBD5")]:
    df_cat = df[df["category"] == category]
    if not df_cat.empty:
        fig.add_trace(go.Bar(
            y=df_cat["scales"],
            x=df_cat["differences"],
            orientation='h',
            marker_color=color,
            name=category,
            cliponaxis=False,
            hovertemplate='<b>%{y}</b><br>Diff: %{x:.2f}<br>p-value: %{customdata:.3f}<extra></extra>',
            customdata=df_cat["p_values"]
        ))

# Add vertical line at x=0
fig.add_shape(
    type="line",
    x0=0, y0=-0.5,
    x1=0, y1=len(df)-0.5,
    line=dict(color="red", width=2, dash="dash")
)

# Update layout
fig.update_layout(
    title="Change in RENSS Scores from Day 1 to Day 7",
    xaxis_title="Mean Diff (D7-D1)",
    yaxis_title="RENSS Items"
)

# Set x-axis range to focus on the relevant region
fig.update_xaxes(range=[-1.6, 1.0])

# Improve bar thickness for better readability
fig.update_traces(width=0.7)

# Center the legend with horizontal orientation
fig.update_layout(
    legend=dict(orientation='h', yanchor='bottom', y=1.05, xanchor='center', x=0.5),
    bargap=0.15  # Add some gap between bars
)

# Save the figure as PNG
fig.write_image("renss_wilcoxon_results.png")

fig.show()