data_generator/data_generator.py

import json
import os
import sys
import random
import re
import concurrent.futures
from typing import List, Dict, Tuple
from datetime import datetime
import threading
from tqdm import tqdm
import requests

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from config import api_key, base_url

# OpenAI API configuration
API_KEY = api_key
API_PROVIDER = base_url

# Column combinations
COLUMN_COMBINATIONS = [
    "categorical + numerical",
    "categorical + numerical + categorical",
    "categorical + numerical + numerical",
    "categorical + numerical + numerical + categorical",
    "temporal + numerical",
    "temporal + numerical + categorical",
    "categorical + numerical + temporal"
]

# Thread-safe print function
print_lock = threading.Lock()
def thread_safe_print(*args, **kwargs):
    with print_lock:
        print(*args, **kwargs)

# LLM 响应时间统计
import time
llm_stats_lock = threading.Lock()
llm_response_times = []
llm_stats_running = True

def get_llm_stats():
    """获取 LLM 统计信息"""
    with llm_stats_lock:
        if not llm_response_times:
            return None
        total_calls = len(llm_response_times)
        avg_time = sum(llm_response_times) / total_calls
        min_time = min(llm_response_times)
        max_time = max(llm_response_times)
        # 最近 10 次调用的平均时间
        recent_times = llm_response_times[-10:]
        recent_avg = sum(recent_times) / len(recent_times)
        return {
            'total_calls': total_calls,
            'avg_time': avg_time,
            'min_time': min_time,
            'max_time': max_time,
            'recent_avg': recent_avg
        }

def llm_stats_reporter():
    """每 30 秒报告一次 LLM 统计信息"""
    global llm_stats_running
    while llm_stats_running:
        time.sleep(30)
        if not llm_stats_running:
            break
        stats = get_llm_stats()
        if stats:
            thread_safe_print(f"\n📊 [LLM Stats] 总调用: {stats['total_calls']} | "
                            f"平均: {stats['avg_time']:.2f}s | "
                            f"最近10次: {stats['recent_avg']:.2f}s | "
                            f"最小: {stats['min_time']:.2f}s | "
                            f"最大: {stats['max_time']:.2f}s")

def query_llm(prompt: str) -> str:
    """
    Query LLM API with a prompt
    Args:
        prompt: The prompt to send to LLM
    Returns:
        str: The response from LLM
    """
    headers = {
        'Authorization': f'Bearer {API_KEY}',
        'Content-Type': 'application/json'
    }
    
    data = {
        'model': 'deepseek-v3.2',
        'messages': [
            {
                'role': 'system', 
                'content': 'You are a senior data analyst and visualization expert with deep knowledge of real-world statistics, industry benchmarks, and data patterns. Generate realistic, diverse data that reflects authentic patterns found in published reports and research. Always return valid JSON format only when requested.'
            },
            {
                'role': 'user', 
                'content': prompt
            }
        ],
        'temperature': 0.85
    }
    
    start_time = time.time()
    
    try:
        response = requests.post(
            f'{API_PROVIDER}/chat/completions',
            headers=headers,
            json=data,
            timeout=120
        )
        response.raise_for_status()
        
        result = response.json()
        content = result['choices'][0]['message']['content'].strip()
        
        # 记录响应时间
        elapsed = time.time() - start_time
        with llm_stats_lock:
            llm_response_times.append(elapsed)
        
        return content
        
    except requests.exceptions.Timeout:
        elapsed = time.time() - start_time
        with llm_stats_lock:
            llm_response_times.append(elapsed)
        thread_safe_print("❌ LLM API 超时（120秒）")
        return None
    except requests.exceptions.HTTPError as e:
        elapsed = time.time() - start_time
        with llm_stats_lock:
            llm_response_times.append(elapsed)
        thread_safe_print(f"❌ LLM API HTTP 错误: {e}")
        if hasattr(e.response, 'text'):
            thread_safe_print(f"   响应: {e.response.text[:200]}")
        return None
    except requests.exceptions.RequestException as e:
        thread_safe_print(f"❌ LLM API 请求错误: {e}")
        return None
    except KeyError as e:
        thread_safe_print(f"❌ LLM API 响应格式错误: {e}")
        return None
    except Exception as e:
        thread_safe_print(f"❌ 查询 LLM 时出错: {e}")
        return None

# Data facts dictionary
datafacts = {
    "trend": {
        "increase": "Increasing trend",
        "decrease": "Decreasing trend",
        "stable": "Stable trend",
        "increase_then_decrease": "Increase then decrease",
        "decrease_then_increase": "Decrease then increase",
        "fluctuation": "Fluctuating trend"
    },
    "proportion": {
        "majority": "Majority",
        "minority": "Minority"
    },
    "value": {
        "total": "Total",
        "average": "Average",
        "maximum": "Maximum",
        "minimum": "Minimum"
    },
    "comparison": {
        "average_higher": "Higher average value compared to others",
        "average_lower": "Lower average value compared to others",
        "significant_difference": "Significant difference compared to other categories"
    },
    "change": {
        "sudden_increase": "Significantly higher value compared to the previous value",
        "sudden_decrease": "Significantly lower value compared to the previous value",
    },
    "correlation": {
        "positive": "Positive correlation",
        "negative": "Negative correlation"
    },
    "rank": {
        "first": "First",
        "second": "Second",
        "third": "Third",
        "last": "Last"
    }
}

def load_themes(file_path: str) -> List[Dict]:
    """Load themes from JSON file"""
    with open(file_path, 'r', encoding='utf-8') as f:
        return json.load(f)
def generate_scenarios_for_theme(theme: Dict, num_scenarios: int = 5) -> List[str]:
    """Step 1: Generate specific scenarios for a given theme
    
    Args:
        theme: 主题字典
        num_scenarios: 要生成的场景数量,默认为5个
    """
    prompt = f"""As a senior data journalist creating infographics for major publications (Bloomberg, The Economist, NYT), generate {num_scenarios} compelling visualization scenarios for: "{theme['theme']}" ({theme['description']})

    **SCENARIO DIVERSITY MATRIX** (each scenario should hit different cells):
    
    | Dimension | Options to Vary |
    |-----------|-----------------|
    | Geographic Scope | Global comparison, Regional (EU/Asia/Americas), National, City-level, Local |
    | Time Frame | Historical (10+ years), Recent trends (2-5 years), Current snapshot, Future projection |
    | Subject Type | Countries, Companies, Industries, Demographics, Products, Behaviors |
    | Analysis Angle | Ranking, Comparison, Distribution, Change over time, Correlation, Breakdown |
    | Audience | Executives, Policymakers, Consumers, Researchers, General public |
    | Data Source Type | Government stats, Industry reports, Surveys, Academic research, Financial data |
    
    **SCENARIO QUALITY REQUIREMENTS:**
    1. Each scenario must have a clear "story hook" - what makes this data interesting or surprising?
    2. Specify concrete analysis subjects (e.g., "Fortune 500 companies" not just "companies")
    3. Include realistic data sources or contexts (e.g., "based on WHO 2023 data")
    4. Make scenarios timely - reference recent events, emerging trends, or evergreen insights
    5. Each scenario should be 20-35 words with specific details
    
    **AVOID:**
    - Generic scenarios without specific subjects
    - Repetitive geographic or temporal scopes
    - Abstract or theoretical framings
    - Scenarios that couldn't be backed by real data
    
    **EXAMPLE GOOD SCENARIOS:**
    - "Comparing semiconductor manufacturing capacity across Taiwan, South Korea, US, and China from 2018-2024, showing the impact of CHIPS Act investments"
    - "How Gen Z vs Millennials allocate monthly entertainment budgets across streaming, gaming, concerts, and dining based on 2023 consumer spending surveys"
    - "European cities ranked by cost of living vs quality of life index, highlighting affordable livable alternatives to London and Paris"
    
    FORMAT: Return ONLY a numbered list (1-{num_scenarios}), one scenario per line. No explanations.
    """
    
    response = query_llm(prompt)
    if not response:
        return []
    
    scenarios = []
    for line in response.split('\n'):
        line = line.strip()
        if line and (line[0].isdigit() or line.lower().startswith('- ')):
            scenario = line.lstrip('0123456789.- ').strip()
            if scenario:
                scenarios.append(scenario)
    
    return scenarios[:num_scenarios]

def select_relevant_datafacts(theme: Dict, scenario: str) -> List[Dict]:
    """Step 2: Select relevant datafacts for the theme and scenario"""
    # Convert datafacts to a flat list for easier processing
    flat_datafacts = []
    for category, facts in datafacts.items():
        for key, description in facts.items():
            flat_datafacts.append({
                "category": category,
                "key": key,
                "description": description
            })
    
    prompt = f"""As a data storytelling expert, select the 5 most compelling data facts to highlight in this visualization:

    **CONTEXT:**
    - THEME: {theme['theme']}
    - SCENARIO: {scenario}
    
    **TASK:**
    Select exactly 5 data facts that would create the most impactful and insightful visualization. Consider:
    
    1. **Story Arc**: Choose facts that together tell a coherent narrative
    2. **Diversity**: Mix different types (trends, comparisons, values, rankings)
    3. **Relevance**: Facts should directly support the scenario's key message
    4. **Visual Impact**: Facts that translate well into compelling visuals
    5. **Insight Value**: Prioritize facts that reveal non-obvious patterns
    
    **SELECTION STRATEGY by Scenario Type:**
    - Time-based analysis → Prioritize: trend, change, value facts
    - Comparison analysis → Prioritize: comparison, rank, proportion facts
    - Distribution analysis → Prioritize: proportion, value, comparison facts
    - Correlation analysis → Prioritize: correlation, trend, change facts
    
    **AVAILABLE FACTS:**
    {json.dumps(flat_datafacts, indent=2)}
    
    **FORMAT:**
    Return ONLY a numbered list (1-5):
    1. [Category]: [Description]
    
    No explanations.
    """
    
    response = query_llm(prompt)
    if not response:
        return []
    
    selected_facts = []
    for line in response.split('\n'):
        line = line.strip()
        if line:
            for fact in flat_datafacts:
                if fact['description'].lower() in line.lower():
                    selected_facts.append(fact)
                    break
    
    return selected_facts[:3]

def extract_json_from_response(response: str) -> str:
    """Extract JSON from LLM response using regex"""
    if not response:
        return "{}"
        
    # Try to find JSON content between triple backticks
    json_match = re.search(r'```(?:json)?\s*([\s\S]*?)\s*```', response)
    if json_match:
        extracted = json_match.group(1).strip()
        return json.loads(extracted)
    
    # Try to find content that looks like a JSON object
    json_match = re.search(r'(\{[\s\S]*\})', response)
    if json_match:
        extracted = json_match.group(1).strip()
        return json.loads(extracted)
    
    # Return the original response if no JSON pattern found
    return response.strip()

def parse_json_safely(text: str) -> Dict:
    """Parse JSON from text safely with error handling"""
    if not text:
        thread_safe_print(f"Warning: Empty text to parse as JSON")
        return {}
        
    return extract_json_from_response(text)

def validate_generated_data(generated_data, column_recommendation):
    """验证生成的数据是否有效"""
    validation = {
        "is_valid": True,
        "issues": []
    }
    
    # 检查数据是否为空
    if not generated_data or "data" not in generated_data or not generated_data["data"]:
        validation["is_valid"] = False
        validation["issues"].append("Empty data array")
        return validation
    
    # 获取期望的列
    expected_columns = [col["name"] for col in column_recommendation.get("columns", [])]
    
    # 检查每行数据
    for i, row in enumerate(generated_data["data"]):
        # 检查是否包含所有列
        for col in expected_columns:
            if col not in row:
                validation["is_valid"] = False
                validation["issues"].append(f"Row {i} missing column '{col}'")
        
        # 检查数值列的数据类型
        for col in column_recommendation.get("columns", []):
            if col["name"] in row and col["data_type"] == "numerical":
                if not isinstance(row[col["name"]], (int, float)):
                    validation["is_valid"] = False
                    validation["issues"].append(f"Row {i}, column '{col['name']}' has non-numeric value: {row[col['name']]}")
    
    return validation

def recommend_columns(theme: Dict, scenario: str, selected_fact: Dict) -> Dict:
    """Step 3: Recommend column structure based on theme, scenario and selected facts"""
    facts_str = f"- {selected_fact['category']}: {selected_fact['description']}"
    
    prompt = f"""As a data architect designing schemas for business intelligence dashboards, recommend the optimal column structure:

    **CONTEXT:**
    - THEME: {theme['theme']}
    - SCENARIO: {scenario}
    - KEY DATA FACT: {facts_str}
    
    **COLUMN COMBINATION OPTIONS:**
    | Combination | Best For | Example |
    |-------------|----------|---------|
    | categorical + numerical | Rankings, comparisons | Countries by GDP |
    | categorical + numerical + numerical | Multi-metric comparisons | Companies: Revenue vs Profit |
    | categorical + numerical + categorical | Grouped comparisons | Products by Revenue by Category |
    | temporal + numerical | Time series, trends | Monthly sales 2020-2024 |
    | temporal + numerical + categorical | Multi-series trends | Quarterly revenue by region |
    | categorical + numerical + temporal | Snapshot comparisons | Country metrics across periods |
    
    **DECISION LOGIC:**
    1. Does the scenario involve change over TIME? → Use temporal column
    2. Does it compare CATEGORIES at a single point? → categorical + numerical
    3. Does it need MULTIPLE METRICS per item? → Add second numerical column
    4. Does it need to SEGMENT by groups? → Add categorical column
    5. PREFER SIMPLER combinations - only add columns if truly needed
    
    **COLUMN NAMING RULES - KEEP IT SHORT:**
    
    ✓ GOOD column names (short, clear):
    - "Country", "Company", "Region", "City", "Industry"
    - "Revenue", "Growth Rate", "Market Share", "GDP", "Population"
    - "Corruption Index", "Democracy Index", "Happiness Score"
    - "Year", "Quarter", "Month"
    
    ✗ BAD column names (too long, avoid):
    - "Corruption Perceptions Index 2023 (Transparency International, 0-100)" → Use "Corruption Index"
    - "EIU Democracy Index 2023 (score 0-10)" → Use "Democracy Index"
    - "Annual Revenue in USD Millions" → Use "Revenue"
    
    **DESCRIPTION RULES - ONE SHORT SENTENCE:**
    - Keep descriptions brief and informative (5-15 words max)
    - Example: "Annual company revenue" not "Annual revenue figures for Fortune 500 companies based on fiscal year 2023 reports"
    - Example: "Corruption perception score by country" not "Transparency International's Corruption Perceptions Index measuring public sector corruption levels"
    
    **UNIT RULES - SYMBOL ONLY:**
    - Use ONLY unit symbols: $, €, £, ¥, %, K, M, B, km, kg, °C, kWh, ms, etc.
    - For dimensionless metrics (index, score, ratio): leave unit as empty string ""
    - ✗ WRONG: "index", "score", "points", "0-100", "0-10", "USD millions"
    - ✓ CORRECT: "", "%", "$", "M", "B"
    
    **FORMAT:**
    Return ONLY valid JSON (no markdown):
    {{
        "selected_combination": "combination_name",
        "columns": [
            {{
                "name": "ShortName",
                "description": "Brief one-sentence description",
                "data_type": "categorical/numerical/temporal",
                "unit": "$ or % or M or empty string"
            }}
        ]
    }}
    
    **EXAMPLE OUTPUT:**
    {{
        "selected_combination": "categorical + numerical",
        "columns": [
            {{"name": "Country", "description": "G20 member countries", "data_type": "categorical", "unit": ""}},
            {{"name": "GDP", "description": "Gross domestic product", "data_type": "numerical", "unit": "B"}}
        ]
    }}
    """
    
    response = query_llm(prompt)
    if not response:
        return {"selected_combination": COLUMN_COMBINATIONS[0], "columns": []}
    
    result = parse_json_safely(response)
    try:
        # Validate the structure to ensure it has required keys
        if "selected_combination" not in result:
            result["selected_combination"] = COLUMN_COMBINATIONS[0]
        if "columns" not in result:
            result["columns"] = []
        # Add unit field if missing
        for col in result["columns"]:
            if "unit" not in col and col["data_type"] == "numerical":
                # Extract unit from description if possible
                desc = col["description"]
                if "%" in desc:
                    col["unit"] = "%"
                elif "$" in desc:
                    col["unit"] = "$"
                elif "£" in desc:
                    col["unit"] = "£"
                elif "€" in desc:
                    col["unit"] = "€"
                else:
                    col["unit"] = ""
            elif "unit" not in col:
                col["unit"] = ""
        return result
    except Exception as e:
        thread_safe_print(f"Error parsing column recommendation: {e}")
        print("response: ", result)
        # Return a default structure
        return {
            "selected_combination": COLUMN_COMBINATIONS[0],
            "columns": [
                {
                    "name": "Category",
                    "description": "Main category for the data",
                    "data_type": "categorical",
                    "unit": ""
                },
                {
                    "name": "Value",
                    "description": "Numerical value",
                    "data_type": "numerical",
                    "unit": ""
                }
            ]
        }

def generate_data(theme: Dict, scenario: str, selected_facts: List[Dict], column_recommendation: Dict, times = 1) -> List[Dict]:
    """生成数据"""
    results = []
    for _ in range(times):
        facts_str = "\n".join([f"- {fact['category']}: {fact['description']}" for fact in selected_facts])
        columns_str = "\n".join([f"- {col['name']} ({col['data_type']}): {col['description']}" for col in column_recommendation['columns']])
        
        # Determine data size constraints based on column combination using ranges
        combination = column_recommendation['selected_combination']
        constraints = []
        
        # Generate range constraints for different combinations
        if combination == "categorical + numerical" or combination == "categorical + numerical + numerical":
            constraints.append("First categorical column should have between 8-20 unique values")
        
        elif combination == "categorical + numerical + categorical" or combination == "categorical + numerical + numerical + categorical":
            constraints.append("First categorical column should have between 8-20 unique values")
            constraints.append("Second categorical column should have between 2-6 unique values")
            constraints.append("Total unique combinations should not exceed 60")
        
        elif combination == "temporal + numerical":
            constraints.append("Number of time points should be between 8-20")
        
        elif combination == "temporal + numerical + categorical":
            constraints.append("Number of time points should be between 5-20")
            constraints.append("Number of categories should be between 2-7")
        
        elif combination == "categorical + numerical + temporal":
            constraints.append("First categorical column should have between 5-20 unique values")
            constraints.append("Number of time points should be between 2-4")
        
        constraints_str = "\n".join([f"- {constraint}" for constraint in constraints])
        
        # 在提示中强调组合完整性和真实性
        prompt = f"""You are a statistician at a major research institution creating synthetic data that mirrors real-world patterns for this visualization:
        
        **VISUALIZATION CONTEXT:**
        - THEME: {theme['theme']}
        - SCENARIO: {scenario}
        - KEY DATA FACTS TO HIGHLIGHT: {facts_str}
        
        **COLUMN STRUCTURE:**
        {columns_str}
        
        **DATA SIZE REQUIREMENTS:**
        {constraints_str}
        
        **REALISM REQUIREMENTS - CRITICAL:**
        
        1. **Categorical Values - Use REAL names:**
           - Countries: Use actual country names (USA, Germany, Japan, Brazil, etc.)
           - Companies: Use real company names (Apple, Toyota, Samsung, Nestlé, etc.)
           - Cities: Use real city names (Tokyo, New York, London, Shanghai, etc.)
           - Industries: Use standard industry names (Healthcare, Technology, Finance, etc.)
           - Products: Use realistic product categories or actual brands
           - Demographics: Use realistic age groups, income brackets, education levels
        
        2. **Numerical Values - Match real-world magnitudes:**
           - GDP: Trillions for large countries, billions for smaller ones
           - Population: Match actual country/city scales
           - Percentages: Realistic ranges (market share 1-40%, growth rates -5% to 15%)
           - Prices: Match real-world price points for the category
           - Revenue: Match industry benchmarks (tech companies in billions, local businesses in millions)
           - Include natural variance - avoid round numbers for most values (use 47.3 not 50)
        
        3. **Temporal Values:**
           - Use YYYY, YYYY-MM, or YYYY-MM-DD format strictly
           - Choose appropriate time spans (economic trends: 2015-2024, recent events: 2022-2024)
           - Ensure chronological consistency in trends
        
        4. **Data Patterns - Reflect reality:**
           - Include outliers naturally (one market leader, one laggard)
           - Show regional/cultural patterns (Asian countries often cluster, Nordic countries cluster)
           - Respect known facts (USA/China usually top GDP, Nordic countries top happiness indices)
           - Add natural noise - real data isn't perfectly smooth
        
        **COMBINATION REQUIREMENTS:**
        - When both temporal and categorical columns exist, generate data for ALL combinations
        - Example: Years=[2020,2021,2022] × Countries=[USA,China] = 6 data points
        
        **TITLE REQUIREMENTS:**
        - main_title: Compelling headline that could appear in The Economist or Bloomberg (8-15 words)
        - sub_title: Contextual detail with time period, data source style, or key finding
        
        **FORMAT:**
        Return ONLY valid JSON (no markdown, no explanation):
        {{
            "data": [
                {{"column_name1": "value1", "column_name2": value2, ...}}
            ],
            "main_insight": "One clear sentence stating the single most important finding from this data",
            "description": "A comprehensive paragraph (80-150 words) that tells the complete data story. Include: (1) What the data shows - the key patterns and findings, (2) Why it matters - the significance and implications, (3) Notable outliers or surprises in the data, (4) Context that helps interpret the numbers. Write as if explaining the visualization to someone who hasn't seen it. This should read like a data journalism paragraph that could accompany the chart in a publication.",
            "titles": {{
                "main_title": "Publication-quality headline",
                "sub_title": "Contextual subtitle with specifics"
            }}
        }}
        """
        
        response = query_llm(prompt)
        if not response:
            continue
        
        try:
            result = parse_json_safely(response)
            results.append(result)
        except Exception as e:
            thread_safe_print(f"Error parsing response: {e}")
            continue
    return results

def process_theme(theme: Dict, syn_data_dir: str) -> Dict:
    """处理单个主题、生成场景并保存数据"""
    thread_safe_print(f"Processing theme: '{theme['theme']}'")
    
    # 步骤1：生成场景
    scenarios = generate_scenarios_for_theme(theme)
    
    result = {
        'theme': theme['theme'],
        'base_description': theme['description'],
        'main_category': theme.get('main_category', ''),  # 添加main_category
        'scenarios': []
    }
    
    # 处理每个场景
    for i, scenario in enumerate(scenarios):
        scenario_num = i + 1
        thread_safe_print(f"  Scenario {scenario_num}/{len(scenarios)}: '{scenario[:50]}...'")
        
        selected_facts = select_relevant_datafacts(theme, scenario)
        for fact in selected_facts:
            try:
                column_recommendation = recommend_columns(theme, scenario, fact)
                
                # 如果没有有效的列结构，跳过
                if not column_recommendation or "columns" not in column_recommendation or not column_recommendation["columns"]:
                    thread_safe_print(f"    No valid column structure, skipping")
                    continue
                    
                # 步骤4：生成数据
                generated_datas = generate_data(theme, scenario, selected_facts, column_recommendation)
                
                # 如果没有数据，跳过
                for generated_data in generated_datas:
                    if not generated_data or "data" not in generated_data or not generated_data["data"]:
                        thread_safe_print(f"    No data generated, skipping")
                        continue
                    
                    # 准备场景结果
                    scenario_result = {
                        'description': generated_data.get('description', scenario),
                        'data': {
                            'data': generated_data.get('data', []),
                            'columns': column_recommendation.get('columns', []),
                            'type_combination': column_recommendation.get('selected_combination', '')
                        },
                        'metadata': {
                            'main_insight': generated_data.get('main_insight', ''),
                            'scenario': scenario,
                            'datafact': selected_facts
                        },
                        'titles': generated_data.get('titles', {'main_title': '', 'sub_title': ''})
                    }
                    
                    save_individual_data(theme['theme'], scenario_result, scenario_num, syn_data_dir, theme.get('main_category', None))
                    
                thread_safe_print(f"    ✓ Completed")
            
            except Exception as e:
                thread_safe_print(f"    ✗ Error: {str(e)}")
                continue
    
    return result

def process_theme_wrapper(args):
    """Wrapper for process_theme to be used with ProcessPoolExecutor"""
    theme, syn_data_dir, theme_idx, total_themes = args
    thread_safe_print(f"\nProcessing theme {theme_idx+1}/{total_themes}: '{theme['theme']}'")
    try:
        theme_result = process_theme(theme, syn_data_dir)
        thread_safe_print(f"✓ Completed processing for theme {theme_idx+1}/{total_themes}: '{theme['theme']}'")
        return theme['theme'], theme_result
    except Exception as e:
        thread_safe_print(f"✗ Error processing theme '{theme['theme']}': {e}")
        import traceback
        thread_safe_print(f"Stack trace: {traceback.format_exc()}")
        return theme['theme'], {
            'theme': theme['theme'],
            'base_description': theme['description'],
            'main_category': theme.get('main_category', ''),  # 添加main_category
            'scenarios': []
        }

def save_results(results: Dict, output_file: str):
    """Save generated results to a JSON file"""
    with print_lock:  # Use lock to prevent file corruption from multiple threads
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(results, f, indent=2, ensure_ascii=False)

def save_individual_data(theme_name: str, scenario_data: Dict, index: int, syn_data_dir: str, main_category: str = None):
    """Save individual scenario data to separate JSON files in syn_data directory"""
    # Create safe filename from main_category (if available) or theme
    if main_category:
        prefix = "".join(c for c in main_category if c.isalnum() or c in [' ', '_']).strip().replace(' ', '_')
    else:
        prefix = "".join(c for c in theme_name if c.isalnum() or c in [' ', '_']).strip().replace(' ', '_')
    
    timestamp = datetime.now().strftime("%H%M%S")
    filename = f"{prefix}_scenario_{index}_{timestamp}_{random.randint(10000, 99999)}.json"
    filepath = os.path.join(syn_data_dir, filename)
    
    try:
        with print_lock:  # Use lock to prevent file corruption from multiple threads
            with open(filepath, 'w', encoding='utf-8') as f:
                json.dump(scenario_data, f, indent=2, ensure_ascii=False)
        return True
    except Exception as e:
        thread_safe_print(f"Error saving individual data file {filepath}: {e}")
        return False

def main():
    global llm_stats_running
    
    # File paths
    current_dir = os.path.dirname(os.path.abspath(__file__))
    theme_file = os.path.join(current_dir, 'theme_new.json')
    output_file = os.path.join(current_dir, 'theme_analysis.json')
    
    thread_safe_print(f"Starting data generation process")
    thread_safe_print(f"Theme file: {theme_file}")
    thread_safe_print(f"Output file: {output_file}")
    
    # 启动 LLM 统计报告线程
    stats_thread = threading.Thread(target=llm_stats_reporter, daemon=True)
    stats_thread.start()
    thread_safe_print("📊 LLM 响应时间监控已启动（每30秒报告一次）")
    
    # Create syn_data directory if it doesn't exist
    syn_data_dir = os.path.join(current_dir, 'syn_data')
    print('Creating syn_data directory...', syn_data_dir)
    if not os.path.exists(syn_data_dir):
        os.makedirs(syn_data_dir)
        thread_safe_print(f"Created directory: {syn_data_dir}")
    
    # Load themes
    try:
        thread_safe_print(f"Loading themes from {theme_file}...")
        themes = load_themes(theme_file)
        thread_safe_print(f"Loaded {len(themes)} themes")
    except Exception as e:
        thread_safe_print(f"Error loading themes: {e}")
        return
    
    # Prepare for parallel processing
    results = {}
    NUM_WORKERS = 20  # Number of concurrent threads
    thread_safe_print(f"Using {NUM_WORKERS} concurrent workers for processing")
    
    # Create arguments for each theme processing task
    theme_args = [(theme, syn_data_dir, idx, len(themes)) for idx, theme in enumerate(themes)]
    
    # Process themes in parallel
    with concurrent.futures.ThreadPoolExecutor(max_workers=NUM_WORKERS) as executor:
        # Submit all tasks and get futures
        future_to_theme = {executor.submit(process_theme_wrapper, args): args[0]['theme'] for args in theme_args}
        
        # Process results as they complete
        for future in tqdm(concurrent.futures.as_completed(future_to_theme), total=len(themes), desc="Processing themes"):
            theme_name = future_to_theme[future]
            try:
                theme_name, theme_result = future.result()
                results[theme_name] = theme_result
                
                # Save overall progress after each theme
                save_results(results, output_file)
                thread_safe_print(f"Updated overall progress file with theme '{theme_name}'")
                
            except Exception as e:
                thread_safe_print(f"Error processing theme '{theme_name}': {e}")
    
    # 停止 LLM 统计报告线程
    llm_stats_running = False
    
    # 输出最终 LLM 统计
    final_stats = get_llm_stats()
    
    thread_safe_print(f"\n========== SUMMARY ==========")
    thread_safe_print(f"Processed {len(results)} themes")
    total_scenarios = sum(len(theme_data['scenarios']) for theme_data in results.values())
    thread_safe_print(f"Generated {total_scenarios} scenarios in total")
    thread_safe_print(f"Full analysis saved to: {output_file}")
    thread_safe_print(f"Individual scenario data saved to: {syn_data_dir}")
    
    if final_stats:
        thread_safe_print(f"\n📊 LLM 调用统计:")
        thread_safe_print(f"   总调用次数: {final_stats['total_calls']}")
        thread_safe_print(f"   平均响应时间: {final_stats['avg_time']:.2f}s")
        thread_safe_print(f"   最小响应时间: {final_stats['min_time']:.2f}s")
        thread_safe_print(f"   最大响应时间: {final_stats['max_time']:.2f}s")
    
    thread_safe_print(f"============================\n")

if __name__ == "__main__":
    main()