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

st.set_page_config(layout="wide")

st.title("TwinTrack")
st.subheader("Observed vs Simulated Interaction Dynamics")

DATA_DIR = "/app/src/"

FILES = {
    "Human": f"{DATA_DIR}/BigFlame_turns.csv",
    "Minimalist": f"{DATA_DIR}/minimalist_turns.csv",
    "Cheerleader": f"{DATA_DIR}/cheerleader_turns.csv",
    "Poet": f"{DATA_DIR}/poet_turns.csv"
}

@st.cache_data
def load_data():
    human = pd.read_csv(FILES['Human'])
    minimalist = pd.read_csv(FILES['Minimalist'])
    cheerleader = pd.read_csv(FILES['Cheerleader'])
    poet = pd.read_csv(FILES['Poet'])
    return human, minimalist, cheerleader, poet

df_human, df_minimlist, df_cheerleader, df_poet = load_data()

def normalize_time(df):
    df["t_norm"] = df.index / (len(df) - 1)
    return df

df_human = normalize_time(df_human)
df_minimlist = normalize_time(df_minimlist)
df_cheerleader = normalize_time(df_cheerleader)
df_poet = normalize_time(df_poet)

synth_map = {
    "minimalist": df_minimlist,
    "cheerleader": df_cheerleader,
    "poet": df_poet
}

choice = st.selectbox("Select Synthetic Persona", synth_map.keys())
df_synth = synth_map[choice]

WINDOW = 0.05  # 5% of conversation

def add_rolling_variance(df, col="tokens_est"):
    window_size = int(len(df) * WINDOW)
    df = df.copy()
    df["roll_var"] = (
        df[col]
        .rolling(window=window_size, center=True)
        .std()
        .fillna(0)
    )
    return df

df_human = add_rolling_variance(df_human)
df_synth = add_rolling_variance(df_synth)

combined_max = max(
    df_human["roll_var"].max(),
    df_synth["roll_var"].max()
)

df_human["roll_var_norm"] = df_human["roll_var"] / combined_max
df_synth["roll_var_norm"] = df_synth["roll_var"] / combined_max

def make_band(df, thickness=20):
    z = [df["roll_var_norm"].values] * thickness
    return z

human_band = make_band(df_human)
synth_band = make_band(df_synth)

# --- REPLACEMENT STARTS HERE ---
from plotly.subplots import make_subplots

COMMON_LEN = 500
t_common = np.linspace(0, 1, COMMON_LEN)

# Interpolate to smooth the data
human_interp = np.interp(t_common, df_human["t_norm"].values, df_human["roll_var_norm"])
synth_interp = np.interp(t_common, df_synth["t_norm"].values, df_synth["roll_var_norm"])

# Create a subplot with 2 rows. 
# "vertical_spacing" is the magic number for your gap! (0.1 = 10% gap)
fig = make_subplots(
    rows=2, cols=1, 
    shared_xaxes=True, 
    vertical_spacing=0.15,  # <--- ADJUST THIS to make the gap bigger/smaller
    subplot_titles=("Human: Big Flame", f"Synthetic: {choice.title()}")
)

# Helper to make the bands "thick" so they look like strips, not thin lines
def make_thick_band(data_row, thickness=10):
    return [data_row] * thickness

# Add Human Trace (Top)
fig.add_trace(go.Heatmap(
    z=make_thick_band(human_interp),
    colorscale="Plasma",
    showscale=False,
    zmin=0, zmax=1  # Lock colors so they are comparable
), row=1, col=1)

# Add Synthetic Trace (Bottom)
fig.add_trace(go.Heatmap(
    z=make_thick_band(synth_interp),
    colorscale="Plasma",
    showscale=False,
    zmin=0, zmax=1
), row=2, col=1)

fig.update_layout(
    height=350,  # Made it slightly taller to accommodate the gap
    margin=dict(l=10, r=10, t=50, b=10),
    xaxis2_title="Normalized Interaction Time", # Label only the bottom axis
)

# Hide the y-axis ticks for a cleaner look
fig.update_yaxes(showticklabels=False, showgrid=False, zeroline=False)
fig.update_xaxes(showgrid=False, zeroline=False)

st.plotly_chart(fig, use_container_width=True)