import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import networkx as nx
import numpy as np
import json
import streamlit as st
from utils import client

def analyze_character(thread_id, additional_context=None):
    run = client.beta.threads.runs.create(
        thread_id=thread_id,
        assistant_id="asst_2xl7bqCuNlvfawVBCkSSRbIP"
    )
    
    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_character_analysis(analysis):
    try:
        # Print raw data for debugging
        st.write("Raw data:")
        st.write(analysis)

        # Parse JSON data
        character_data = json.loads(analysis)

        # Create a DataFrame for character attributes
        char_df = pd.DataFrame([
            {
                'character': char,
                'screentime': data['screentime'],
                'motivation': data['motivation']['score'],
                'internal_conflict': data['internal_conflict']['score'],
                'backstory': data['backstory']['score'],
                'character_arc': data['character_arc']['score']
            }
            for char, data in character_data.items()
        ])

        # Display the DataFrame
        st.write("Character Attributes:")
        st.dataframe(char_df)

        # Screentime Bar Chart
        st.write("### Character Screentime")
        fig, ax = plt.subplots(figsize=(10, 6))
        sns.barplot(x='character', y='screentime', data=char_df, ax=ax)
        plt.title("Character Screentime")
        plt.xticks(rotation=45, ha='right')
        st.pyplot(fig)

        # Character Attributes Heatmap
        st.write("### Character Attributes Heatmap")
        attr_df = char_df.set_index('character')[['motivation', 'internal_conflict', 'backstory', 'character_arc']]
        fig, ax = plt.subplots(figsize=(12, 8))
        sns.heatmap(attr_df, annot=True, cmap="YlGnBu", ax=ax)
        plt.title("Character Attributes Heatmap")
        st.pyplot(fig)

        # Radar Charts for each character
        st.write("### Character Radar Charts")
        attributes = ['motivation', 'internal_conflict', 'backstory', 'character_arc']
        for _, row in char_df.iterrows():
            values = row[attributes].values
            angles = np.linspace(0, 2*np.pi, len(attributes), endpoint=False)
            values = np.concatenate((values, [values[0]]))
            angles = np.concatenate((angles, [angles[0]]))
            
            fig, ax = plt.subplots(figsize=(6, 6), subplot_kw=dict(projection='polar'))
            ax.plot(angles, values, 'o-', linewidth=2)
            ax.fill(angles, values, alpha=0.25)
            ax.set_xticks(angles[:-1])
            ax.set_xticklabels(attributes)
            ax.set_ylim(0, 1)
            ax.set_title(f"{row['character']} Attributes")
            st.pyplot(fig)

        # Character Relationships Network Graph
        st.write("### Character Relationships Network")
        G = nx.Graph()
        for char, data in character_data.items():
            G.add_node(char)
            if 'relationships' in data:
                for rel, rel_data in data['relationships'].items():
                    G.add_edge(char, rel, weight=rel_data['strength'])

        fig, ax = plt.subplots(figsize=(12, 8))
        pos = nx.spring_layout(G)
        nx.draw(G, pos, with_labels=True, node_color='lightblue', 
                node_size=3000, font_size=10, font_weight='bold')
        edge_labels = nx.get_edge_attributes(G, 'weight')
        nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels)
        plt.title("Character Relationships Network")
        st.pyplot(fig)

        # Display character details
        st.write("### Character Details")
        for char, data in character_data.items():
            st.write(f"**{char}**")
            st.write(f"Screentime: {data['screentime']*100:.1f}%")
            for attr in attributes:
                st.write(f"{attr.capitalize()}: {data[attr]['score']} - {data[attr]['reason']}")
            if 'relationships' in data:
                st.write("Relationships:")
                for rel, rel_data in data['relationships'].items():
                    st.write(f"- {rel}: Strength {rel_data['strength']} - {rel_data['reason']}")
            st.write("---")

    except Exception as e:
        st.error(f"Error processing data for Character Analysis: {e}")
        st.write("Please check the structure of the JSON data:")
        st.json(analysis)