Thematic Screeners

import os
from logging import Logger, getLogger
from typing import Dict, List, Optional

import matplotlib.gridspec as gridspec
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
import numpy as np
import pandas as pd
from pandas import DataFrame, merge
import plotly
import plotly.io as pio
import seaborn as sns
pio.renderers.default = 'colab' # Change this if you're running in a different environment
plotly.offline.init_notebook_mode()

from bigdata_client import Bigdata
from bigdata_client.models.entities import Company
from bigdata_client.models.search import DocumentType

from bigdata_research_tools.labeler.screener_labeler import ScreenerLabeler
from bigdata_research_tools.search.screener_search import search_by_companies
from bigdata_research_tools.themes import (
    SourceType,
    generate_theme_tree,
    stringify_label_summaries,
)

from bigdata_research_tools.screeners import ExecutiveNarrativeFactor
from bigdata_research_tools.screeners.utils import get_scored_df

# For Excel functionality
from bigdata_research_tools.excel import ExcelManager

# Define output file paths for our results
output_dir = "output"
os.makedirs(output_dir, exist_ok=True)

export_path = f"{output_dir}/thematic_screener_results.xlsx"

# ===== Theme Definition =====
main_theme = "Supply Chain Reshaping"

# ===== Company Universe (from Watchlist) =====
# Get Top Us 100 watchlist from Bigdata.com
top100_watchlist_id = "44118802-9104-4265-b97a-2e6d88d74893"
watchlist = bigdata.watchlists.get(top100_watchlist_id)
companies = bigdata.knowledge_graph.get_entities(watchlist.items)

# ===== LLM Specification =====
llm_model = "openai::gpt-4o-mini"

# ===== Transcript Configuration =====
document_type = DocumentType.TRANSCRIPTS
fiscal_year = 2024

# ===== Enable/Disable Reranker =====
rerank_threshold = None

# ===== Specify Time Range =====
start_date = "2024-03-01"
end_date = "2025-03-01"

theme_tree = generate_theme_tree(
    main_theme=main_theme,
    dataset=SourceType.CORPORATE_DOCS,
)

theme_tree.visualize()

node_summaries = theme_tree.get_summaries()
node_summaries

freq = "6M"
document_limit = 100
batch_size = 10

df_sentences = search_by_companies(
    companies=companies,
    sentences=node_summaries,
    start_date=start_date,
    end_date=end_date,
    scope=document_type,
    fiscal_year=fiscal_year,
    rerank_threshold=rerank_threshold,
    freq=freq,
    document_limit=document_limit,
    batch_size=batch_size,
)

df_sentences.head(5)

labeler = ScreenerLabeler(llm_model=llm_model)
df_labels = labeler.get_labels(
    main_theme=main_theme,
    labels=list(theme_tree.get_terminal_label_summaries().keys()),
    texts=df_sentences["masked_text"].tolist(),
)

# Merge and process results
df = merge(df_sentences, df_labels, left_index=True, right_index=True)
df = labeler.post_process_dataframe(df)

df_company, df_industry = DataFrame(), DataFrame()

df_company = get_scored_df(
    df, index_columns=["Company", "Ticker", "Industry"], pivot_column="Theme"
)
df_industry = get_scored_df(
    df, index_columns=["Industry"], pivot_column="Theme"
)

import pandas as pd
import numpy as np
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import plotly.io as pio
import plotly.offline
pio.renderers.default = 'colab'
plotly.offline.init_notebook_mode()


def create_thematic_exposure_dashboard(df_company, n_companies=10):
    """
    Creates five separate figures for analyzing thematic exposure of companies.

    Parameters:
    -----------
    df_company : pandas.DataFrame
        DataFrame containing company data with columns for 'Company', 'Industry',
        'Composite Score', and multiple thematic exposure columns.
    n_companies : int, default=10
        Number of companies to include in the analysis.

    Returns:
    --------
    tuple
        A tuple containing five Plotly figures:
        - fig_heatmap: Company-theme exposure heatmap
        - fig_total_scores: Total composite scores bar chart
        - fig_scatter_themes: Top themes scatter plot by company
        - fig_bar_themes: Thematic scores horizontal bar chart
        - fig_industry_heatmap: Industry-level analysis heatmap
    """
    # Select top n companies and reset index
    df = df_company[:n_companies].reset_index(drop=True).copy()

    # Extract theme column names (all columns between 'Industry' and 'Composite Score')
    theme_columns = list(df.iloc[:, 3:-1].columns)

    # Create all figures
    fig_heatmap = create_company_theme_heatmap(df, theme_columns)
    fig_total_scores = create_company_scores_bar_chart(df)
    fig_scatter_themes = create_top_themes_scatter_plot(df, theme_columns)
    fig_bar_themes = create_themes_summary_bar_chart(df, theme_columns)
    fig_industry_heatmap = create_industry_analysis_heatmap(df, theme_columns)

    return fig_heatmap, fig_total_scores, fig_scatter_themes, fig_bar_themes, fig_industry_heatmap


def create_company_theme_heatmap(df, theme_columns):
    """
    Creates a heatmap showing thematic exposure scores for each company.

    Parameters:
    -----------
    df : pandas.DataFrame
        The data frame containing company and theme data
    theme_columns : list
        List of column names representing thematic categories

    Returns:
    --------
    plotly.graph_objects.Figure
        Heatmap figure with color scale
    """
    fig = go.Figure()
    heatmap_z = df[theme_columns].values
    heatmap_x = theme_columns
    heatmap_y = df['Company'].tolist()

    fig.add_trace(
        go.Heatmap(
            z=heatmap_z,
            x=heatmap_x,
            y=heatmap_y,
            colorscale='YlGnBu',
            text=heatmap_z.astype(int),
            texttemplate="%{text}",
            showscale=True,
        )
    )

    fig.update_layout(
        title='Company Thematic Exposure Heatmap (Raw Scores)',
        height=600,
        width=1200,
        margin=dict(l=60, r=50, t=80, b=50),
        xaxis=dict(tickangle=45, tickfont=dict(size=9)),
        yaxis=dict(title="Company")
    )

    return fig


def create_company_scores_bar_chart(df):
    """
    Creates a horizontal bar chart showing total composite scores for companies.

    Parameters:
    -----------
    df : pandas.DataFrame
        The data frame containing company scores

    Returns:
    --------
    plotly.graph_objects.Figure
        Bar chart figure showing total composite scores
    """
    fig = go.Figure()
    companies = df['Company'].tolist()
    total_scores = df['Composite Score'].tolist()

    # Sort by score for better visualization (highest first)
    sorted_indices = np.argsort(total_scores)[::-1]
    sorted_companies = [companies[i] for i in sorted_indices]
    sorted_scores = [total_scores[i] for i in sorted_indices]

    fig.add_trace(
        go.Bar(
            y=sorted_companies,
            x=sorted_scores,
            orientation='h',
            marker=dict(
                color=sorted_scores,
                colorscale='Viridis',
                showscale=False
            ),
            text=sorted_scores,
            textposition='outside',
            textfont=dict(size=10),
        )
    )

    fig.update_layout(
        title='Company Total Composite Scores',
        height=600,
        width=1200,
        margin=dict(l=60, r=50, t=80, b=50),
        xaxis=dict(title="Total Composite Score"),
        yaxis=dict(title="Company")
    )

    return fig


def create_top_themes_scatter_plot(df, theme_columns):
    """
    Creates a scatter plot showing the top 3 thematic exposures for each company.

    Parameters:
    -----------
    df : pandas.DataFrame
        The data frame containing company and theme data
    theme_columns : list
        List of column names representing thematic categories

    Returns:
    --------
    plotly.graph_objects.Figure
        Scatter plot showing top themes by company
    """
    fig = go.Figure()
    
    max_score = df[theme_columns].values.max()
    companies_unique = df['Company'].unique()
    
    for i, company in enumerate(companies_unique):
        company_data = df[df['Company'] == company]
        if len(company_data) == 0:
            continue
        
        company_row = company_data.iloc[0]
        company_scores = company_row[theme_columns].values
        top_indices = np.argsort(company_scores)[-3:]
        
        x_values = []
        y_values = []
        sizes = []
        hover_texts = []
        
        for idx in top_indices:
            if company_scores[idx] > 0:
                theme = theme_columns[idx]
                score = company_scores[idx]
                size = (score / max_score) * 80
                x_values.append(company)
                y_values.append(theme)
                sizes.append(size)
                hover_texts.append(f"{company}<br>{theme}: {int(score)}")
        
        if len(x_values) > 0:
            fig.add_trace(
                go.Scatter(
                    x=x_values,
                    y=y_values,
                    mode='markers',
                    marker=dict(
                        size=sizes,
                        sizemode='area',
                        sizeref=0.15,
                        color=i,
                        colorscale='Turbo',
                        showscale=False,
                        opacity=0.7,
                        line=dict(width=1, color='DarkSlateGrey'),
                    ),
                    text=hover_texts,
                    hoverinfo='text',
                    name=company,
                )
            )
    
    fig.update_layout(
        height=600,
        width=1200,
        title_text="Top 3 Thematic Exposures by Company",
        showlegend=False,
        margin=dict(l=60, r=50, t=80, b=50),
    )
    
    fig.update_xaxes(title_text="Company")
    fig.update_yaxes(title_text="Theme")
    
    return fig


def create_themes_summary_bar_chart(df, theme_columns):
    """
    Creates a horizontal bar chart showing total scores across all themes.
    Uses a red color palette that never goes to white.

    Parameters:
    -----------
    df : pandas.DataFrame
        The data frame containing company and theme data
    theme_columns : list
        List of column names representing thematic categories

    Returns:
    --------
    plotly.graph_objects.Figure
        Horizontal bar chart showing thematic scores summary
    """
    fig = go.Figure()
    
    # Calculate theme totals across all companies
    theme_totals = df[theme_columns].sum()
    theme_names = theme_totals.index.tolist()
    theme_values = theme_totals.values.tolist()
    
    # Sort by values (descending)
    sorted_indices = np.argsort(theme_values)[::-1]
    sorted_themes = [theme_names[i] for i in sorted_indices]
    sorted_values = [theme_values[i] for i in sorted_indices]
    
    # Create custom color scale that never goes to white
    normalized_values = np.array(sorted_values)
    normalized_values = (normalized_values - normalized_values.min()) / (normalized_values.max() - normalized_values.min())
    # Scale to range 0.3 to 1.0 to avoid white colors
    color_values = 0.3 + (normalized_values * 0.7)
    
    fig.add_trace(
        go.Bar(
            y=sorted_themes,
            x=sorted_values,
            orientation='h',
            marker=dict(
                color=color_values,
                colorscale=[[0, '#8B0000'], [1, '#FF4500']],  # Dark red to orange-red, no white
                showscale=False
            ),
            text=sorted_values,
            textposition='outside',
            textfont=dict(size=10),
        )
    )
    
    fig.update_layout(
        height=600,
        width=1200,
        title_text="Total Thematic Scores Summary",
        showlegend=False,
        margin=dict(l=60, r=50, t=80, b=50),
    )
    
    fig.update_xaxes(title_text="Total Score Across All Companies")
    fig.update_yaxes(title_text="Theme")
    
    return fig


def create_industry_analysis_heatmap(df, theme_columns):
    """
    Creates a heatmap showing average thematic scores grouped by industry.

    Parameters:
    -----------
    df : pandas.DataFrame
        The data frame containing company and industry data
    theme_columns : list
        List of column names representing thematic categories

    Returns:
    --------
    plotly.graph_objects.Figure
        Heatmap figure showing industry-level thematic analysis
    """
    # Group by industry and calculate mean scores
    industry_data = []

    for industry, group in df.groupby('Industry'):
        for theme in theme_columns:
            industry_data.append({
                'Industry': industry,
                'Theme': theme,
                'Average_Score': group[theme].mean()
            })

    industry_df = pd.DataFrame(industry_data)

    # Create a pivot table for the heatmap
    industry_pivot = industry_df.pivot(index='Industry', columns='Theme', values='Average_Score')

    # Create the industry analysis figure
    fig = go.Figure(data=go.Heatmap(
        z=industry_pivot.values,
        x=industry_pivot.columns,
        y=industry_pivot.index,
        colorscale='YlOrRd',
        text=np.round(industry_pivot.values, 1),
        texttemplate="%{text}",
        showscale=True,
    ))

    # Format the industry analysis figure
    fig.update_layout(
        title='Industry-Level Thematic Exposure Analysis (Average Scores)',
        height=500,
        width=1200,
        margin=dict(l=60, r=50, t=80, b=50),
        xaxis=dict(tickangle=45, tickfont=dict(size=9)),
        yaxis=dict(title="Industry")
    )

    return fig


# Enable custom widget manager for Colab
from google.colab import output
output.enable_custom_widget_manager()

# Alternative display method using HTML
from IPython.display import HTML

# Create the dashboard with all figures
fig_heatmap, fig_total_scores, fig_scatter_themes, fig_bar_themes, fig_industry_heatmap = create_thematic_exposure_dashboard(df_company)

# Convert figures to HTML for display
html_heatmap = fig_heatmap.to_html(include_plotlyjs='cdn')
html_total_scores = fig_total_scores.to_html(include_plotlyjs='cdn')
html_scatter_themes = fig_scatter_themes.to_html(include_plotlyjs='cdn')
html_bar_themes = fig_bar_themes.to_html(include_plotlyjs='cdn')
html_industry_heatmap = fig_industry_heatmap.to_html(include_plotlyjs='cdn')

# Display all figures
print("\n")
display(HTML(html_heatmap))
print("\n")
display(HTML(html_total_scores))
print("\n")
display(HTML(html_scatter_themes))
print("\n")
display(HTML(html_bar_themes))
print("\n")
display(HTML(html_industry_heatmap))

try:
    # Create the Excel manager
    excel_manager = ExcelManager()

    # Define the dataframes and their sheet configurations
    df_args = [
        (df, "Semantic Labels", (0, 0)),
        (df_company, "By Company", (2, 4)),
        (df_industry, "By Industry", (2, 2))
    ]

    # Save the workbook
    excel_manager.save_workbook(df_args, export_path)

except Exception as e:
    print(f"Warning while exporting to excel: {e}")