US Tariffs: Risks & Strategies

import os
from types import SimpleNamespace

from IPython.display import display, HTML
import pandas as pd
import numpy as np
from datetime import datetime, timedelta

from src.bigdata_rest import load_universe, company_ids_from_universe
from src.mindmap.generate_trees import generate_themes_tree_dict, get_most_granular_elements
from src.mindmap.themes import print_tree
from src.search.content_retrieval import DataRetriever
from src.label.label_process import LabelProcessor
from src.report_generator import GenerateReport

import plotly.io as pio

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

# ===== Customizable Parameters =====

# Company Universe: CSV with RP_ENTITY_ID + COMPANY_NAME (see Thematic_Screener_CLI)
universe_path = "../Thematic_Screener_CLI/40_companies.csv"
universe_df = load_universe(universe_path).head(5).reset_index(drop=True)
company_ids = company_ids_from_universe(universe_df)
id_to_name = dict(zip(universe_df["RP_ENTITY_ID"], universe_df["COMPANY_NAME"]))
list_entities = [SimpleNamespace(id=eid, name=name) for eid, name in id_to_name.items()]

# Main Analysis Theme
main_theme = 'US Import Tariffs Corporate Risk Impact Analysis'
focus = "Provide a detailed taxonomy of risks describing how new American import tariffs will impact worldwide companies, their operations and strategy."

# LLM Model Configuration (gpt-5.6-luna: omit temperature/top_p)
llm_model = "gpt-5.6-luna"

# Time Range Configuration
end_date = datetime.now().strftime("%Y-%m-%d")
start_date = (datetime.now() - timedelta(days=30)).strftime("%Y-%m-%d")
freq = 'M'  # Monthly search frequency

# Enable/Disable Reranker 
rerank_threshold = None

# Document Retrieval Limits
document_limit_news = 10
document_limit_filings = 5
batch_size = 1

# Toggle fallback to News for company responses
response_from_news = True

themes_tree_dict = generate_themes_tree_dict(main_theme, focus)
risk_tree = themes_tree_dict[main_theme]
terminal_labels = get_most_granular_elements(risk_tree, 'Label')
risk_summaries = get_most_granular_elements(risk_tree, 'Summary')

print(f"✅ Taxonomy generated with {len(terminal_labels)} leaf risk categories")
print_tree(risk_tree)

data_retriever_news = DataRetriever(
    company_ids=company_ids,
    id_to_name=id_to_name,
    document_limit=document_limit_news,
    sortby="relevance",
    search_freq=freq,
    start_date_query=start_date,
    end_date_query=end_date,
)

df_sentences_semantic = data_retriever_news.retrieve(
    themes_tree_dict=themes_tree_dict,
    list_specific_themes=[main_theme],
    document_type="news",
)

if df_sentences_semantic is None:
    df_sentences_semantic = pd.DataFrame()

df_sentences_semantic = df_sentences_semantic.head(document_limit_news)

label_processor = LabelProcessor(
    list_entities=list_entities,
    themes_tree_dict=themes_tree_dict,
    list_specific_themes=[main_theme],
    api_key=OPENAI_API_KEY,
)

if df_sentences_semantic.empty:
    df_labeled = pd.DataFrame()
else:
    df_labeled = label_processor.run_label_process(df_sentences=df_sentences_semantic)
    if df_labeled is None:
        df_labeled = pd.DataFrame()

report_generator = GenerateReport(
        universe_df=universe_df,
        main_theme=main_theme,
        focus=focus,
        llm_model=llm_model,
        api_key=OPENAI_API_KEY,
        start_date=start_date,
        end_date=end_date,
        search_frequency=freq,
        document_limit_news=document_limit_news,
        document_limit_filings=document_limit_filings,
        batch_size=batch_size,
        themes_tree_dict=themes_tree_dict
)

# Generate the risk report data
report = report_generator.generate_report(
    df_labeled=df_labeled,
    news_search_fallback = response_from_news, # Use response_from_news to enable/disable News fallback
    import_from_path=None,
    export_to_path=output_dir,
)

from src.html_report import generate_html_report, prepare_data_report_0

# Extract report data for processing
df_by_theme = report.report_by_theme
df_by_company_with_responses = report.report_by_company

# Prepare data with executive summary formatting
top_by_theme, top_by_company = prepare_data_report_0(df_by_theme, df_by_company_with_responses)

# Generate executive-style HTML report
html_content = generate_html_report(top_by_theme, top_by_company, 'US Import Tariffs: Corporate Risk Impact Analysis')

# Save the executive report
report_filename = f'{output_dir}/tariffs_executive_report.html'
with open(report_filename, 'w') as file:
     file.write(html_content)

from src.html_report import generate_html_report_v1, prepare_data_report_1

# Prepare data with detailed analysis formatting
top_by_theme, top_by_company = prepare_data_report_1(df_by_theme, df_by_company_with_responses)

# Generate detailed analysis HTML report
html_content_detailed = generate_html_report_v1(top_by_theme, top_by_company, 'US Import Tariffs: Comprehensive Risk Analysis')

# Save the detailed report
detailed_filename = f'{output_dir}/tariffs_detailed_analysis.html'
with open(detailed_filename, 'w') as file:
     file.write(html_content_detailed)

# Optional: Export structured data for external analysis
try:
    # Export the core datasets
    df_by_theme.to_csv(f'{output_dir}/tariffs_risks_by_theme.csv', index=False)
    df_by_company_with_responses.to_csv(f'{output_dir}/tariffs_risks_by_company.csv', index=False)
    
    print("✅ Data exported successfully:")
    print(f"   - Thematic analysis: {output_dir}/tariffs_risks_by_theme.csv")
    print(f"   - Company analysis: {output_dir}/tariffs_risks_by_company.csv")
    print(f"   - Executive report: {output_dir}/tariffs_executive_report.html")
    print(f"   - Detailed analysis: {output_dir}/tariffs_detailed_analysis.html")
    
except Exception as e:
    print(f"Warning: Export failed - {e}")