import plotly.express as px
import plotly.graph_objects as go
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import pandas as pd
import numpy as np

# Sample Data
df = pd.DataFrame({
    'Group ID': ['A', 'A', 'B', 'B', 'C', 'C']*100,
    'Entity ID': ['M1', 'M2', 'M1', 'M2', 'M1', 'M2']*100,
    'Value': np.random.randn(600)
})

# Create a colormap
def get_hex_color(value):
    rgba_color = plt.get_cmap('coolwarm')(mcolors.Normalize(vmin=df['Value'].min(), vmax=df['Value'].max())(value))
    return mcolors.to_hex(rgba_color)
colormap = {value: get_hex_color(value) for value in df['Value'].unique()}

# ----------------------------------------

# Basic example
fig = px.strip(df, x='Group ID', y='Value', color='Value', color_discrete_map=colormap)
fig.update_traces(marker=dict(size=10, opacity=0.8, line=dict(width=1, color='DarkSlateGrey')))

# Desired behavior (this does not work):
fig.update_traces(jitter=0.1) # We would expect this to have a small amount of jitter; instead, there is a lot

# Let's try extracting the data from the plotly express figure
list(fig.data)
for box in fig.data:
    box.update(jitter=0.9) # We would expect this to have a lot of jitter, but there appears to be none
new_fig = go.Figure(fig.data)

# This is probably because the go.Box width is None; let's try adding a width param
add_width = go.Figure(fig.data).update_traces(width=20, jitter=1)

# ----------------------------------------

# Styling (ignore)
fig.update_layout(showlegend=False, title='Basic Colored Grouped Strip Plot', template='plotly_white', width=300, height=200, margin=dict(l=0, r=0, b=0, t=40), xaxis_visible=False)
fig.add_annotation(text='<b>update_traces(jitter=0.1)</b><br>We would expect this<br>to have a small amount of jitter;<br>instead, there is a lot', x=0.5, y=0.5, xref='paper', yref='paper', showarrow=False, bgcolor='white', bordercolor='#999', opacity=0.8)
new_fig.add_annotation(text='<b>for box in fig.data:<br>box.update(jitter=0.9)</b><br>We would expect this<br>to have a lot of jitter,<br>but there appears to be none', x=0.5, y=0.5, xref='paper', yref='paper', showarrow=False, bgcolor='white', bordercolor='#999', opacity=0.8)
new_fig.update_layout(showlegend=False, title='Fig.data extracted from px.strip', template='plotly_white', width=300, height=200, margin=dict(l=0, r=0, b=0, t=40), xaxis_visible=False)
add_width.layout=new_fig.layout
add_width.update_layout(title='Adding a width param to the traces')
# Remove the annotation:
add_width.layout.annotations = [new_fig.layout.annotations[0].update({'text':'<b>update_traces(width=20, jitter=1)</b><br>Giving the boxes a width<br>just shifts everything over'})]
for demo in [fig, new_fig, add_width]: demo.show(config = {'displayModeBar': False})