# audience_reaction.py

import pandas as pd
import matplotlib.pyplot as plt
import streamlit as st
import json
from utils import client

import plotly.graph_objs as go
import plotly.express as px
from plotly.subplots import make_subplots

def analyze_audience_reaction(thread_id, additional_context=None):
    run = client.beta.threads.runs.create(
        thread_id=thread_id,
        assistant_id="asst_kr92jSrWpbdEI9wSHl2OIOA2"
    )
    
    while run.status in ['queued', 'in_progress', 'cancelling']:
        run = client.beta.threads.runs.retrieve(
            thread_id=thread_id,
            run_id=run.id
        )
    
    if run.status == 'completed':
        messages = client.beta.threads.messages.list(thread_id=thread_id)
        analysis = next((msg.content[0].text.value for msg in reversed(list(messages)) if msg.role == "assistant"), "")
        return analysis
    else:
        return f"Error: Run status is {run.status}"

def process_audience_reaction(analysis):
    try:
        # Parse JSON data
        data = json.loads(analysis)
        
        # Display raw data for debugging
        st.write("Raw data:")
        st.json(data)

        # Convert the JSON data to a DataFrame
        df = pd.DataFrame.from_dict(data, orient='index')
        df = df.reset_index()
        df.columns = ['Sequence', 'sequence_name', 'emotional_impact', 'excitement']
        
        # Sort the DataFrame by the order of sequences in the original JSON
        df['Sequence'] = pd.Categorical(df['Sequence'], categories=data.keys(), ordered=True)
        df = df.sort_values('Sequence')

        # 1. Emotional Impact and Excitement Line Chart
        st.write("### Emotional Impact and Excitement Throughout the Script")
        fig = go.Figure()
        fig.add_trace(go.Scatter(x=df['Sequence'], y=df['emotional_impact'], mode='lines+markers', name='Emotional Impact'))
        fig.add_trace(go.Scatter(x=df['Sequence'], y=df['excitement'], mode='lines+markers', name='Excitement'))
        fig.update_layout(title='Emotional Impact and Excitement Throughout the Script',
                          xaxis_title='Sequence', yaxis_title='Level',
                          legend_title='Metric')
        st.plotly_chart(fig, use_container_width=True)

        # 2. Emotional Impact vs Excitement Scatter Plot
        st.write("### Emotional Impact vs Excitement")
        fig = px.scatter(df, x='emotional_impact', y='excitement', text='Sequence',
                         title='Emotional Impact vs Excitement for Each Sequence',
                         labels={'emotional_impact': 'Emotional Impact', 'excitement': 'Excitement'})
        fig.update_traces(textposition='top center')
        st.plotly_chart(fig, use_container_width=True)

        # 3. Stacked Bar Chart of Emotional Impact and Excitement
        st.write("### Comparison of Emotional Impact and Excitement")
        fig = go.Figure(data=[
            go.Bar(name='Emotional Impact', x=df['Sequence'], y=df['emotional_impact']),
            go.Bar(name='Excitement', x=df['Sequence'], y=df['excitement'])
        ])
        fig.update_layout(barmode='group', title='Comparison of Emotional Impact and Excitement Across Sequences')
        st.plotly_chart(fig, use_container_width=True)

        # 4. Radar Chart of Emotional Impact and Excitement
        st.write("### Radar Chart of Emotional Impact and Excitement")
        fig = go.Figure(data=go.Scatterpolar(
          r=df['emotional_impact'].tolist() + [df['emotional_impact'].iloc[0]],
          theta=df['Sequence'].tolist() + [df['Sequence'].iloc[0]],
          fill='toself',
          name='Emotional Impact'
        ))
        fig.add_trace(go.Scatterpolar(
          r=df['excitement'].tolist() + [df['excitement'].iloc[0]],
          theta=df['Sequence'].tolist() + [df['Sequence'].iloc[0]],
          fill='toself',
          name='Excitement'
        ))
        fig.update_layout(
          polar=dict(radialaxis=dict(visible=True, range=[0, 1])),
          showlegend=True,
          title='Radar Chart of Emotional Impact and Excitement'
        )
        st.plotly_chart(fig, use_container_width=True)

        # Key Insights
        st.write("### Key Insights")
        peak_emotion = df.loc[df['emotional_impact'].idxmax()]
        peak_excitement = df.loc[df['excitement'].idxmax()]
        avg_emotion = df['emotional_impact'].mean()
        avg_excitement = df['excitement'].mean()
        
        st.write(f"1. The sequence with the highest emotional impact is '{peak_emotion['Sequence']}' with a score of {peak_emotion['emotional_impact']:.2f}.")
        st.write(f"2. The most exciting sequence is '{peak_excitement['Sequence']}' with an excitement level of {peak_excitement['excitement']:.2f}.")
        st.write(f"3. The average emotional impact across all sequences is {avg_emotion:.2f}.")
        st.write(f"4. The average excitement level across all sequences is {avg_excitement:.2f}.")
        if peak_emotion['Sequence'] == peak_excitement['Sequence']:
            st.write(f"5. '{peak_emotion['Sequence']}' is the most impactful sequence, peaking in both emotional impact and excitement.")
        
    except Exception as e:
        st.error(f"Error processing data for Audience Reaction Analysis: {e}")
        st.write("Please check the structure of the JSON data:")
        st.json(analysis)