scripts/process_data.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
ChartPipeline: Simplified implementation of modules 1-6
This script implements a simplified version of the first 6 modules in the ChartPipeline framework.
"""

import json
import random
import argparse
import logging
import re
import copy
import sys
import requests
from pathlib import Path
from typing import Dict, List, Any, Union, Optional
import base64
import os
from collections import Counter
import numpy as np
import torch
from transformers import AutoTokenizer, AutoModel

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger("ChartPipeline-Simplified")

# Initialize BERT model and tokenizer for embeddings
try:
    import torch
    from transformers import AutoTokenizer, AutoModel
    tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
    model = AutoModel.from_pretrained("bert-base-uncased")
    logger.info("Successfully loaded BERT model for embeddings")
    USE_BERT = True
except Exception as e:
    logger.warning(f"Failed to load BERT model: {str(e)}. Will use simplified embeddings instead.")
    USE_BERT = False
    # Make dummy imports to avoid errors
    class DummyModule:
        pass
    if 'torch' not in sys.modules:
        sys.modules['torch'] = DummyModule()
    if 'transformers' not in sys.modules:
        transformers_dummy = DummyModule()
        transformers_dummy.AutoTokenizer = DummyModule
        transformers_dummy.AutoModel = DummyModule
        sys.modules['transformers'] = transformers_dummy

# Cache for icon embeddings to avoid recomputing
ICON_EMBEDDINGS_CACHE = {}

# OpenAI API configuration
API_KEY = os.getenv("OPENAI_API_KEY") or os.getenv("AIHUBMIX_API_KEY", "")
API_PROVIDER = os.getenv("OPENAI_BASE_URL", "https://aihubmix.com")

def query_openai(prompt: str) -> str:
    """
    Query OpenAI API with a prompt
    Args:
        prompt: The prompt to send to OpenAI
    Returns:
        str: The response from OpenAI
    """
    headers = {
        'Authorization': f'Bearer {API_KEY}',
        'Content-Type': 'application/json'
    }
    
    data = {
        'model': 'gemini-2.0-flash',
        'messages': [
            {'role': 'system', 'content': 'You are a data visualization expert. Provide concise, specific answers.'},
            {'role': 'user', 'content': prompt}
        ],
        'temperature': 0.3,  # Lower temperature for more focused responses
        'max_tokens': 3000    # Limit response length
    }
    
    try:
        response = requests.post(f'{API_PROVIDER}/v1/chat/completions', headers=headers, json=data)
        response.raise_for_status()
        return response.json()['choices'][0]['message']['content'].strip()
    except Exception as e:
        print(f"Error querying OpenAI: {e}")
        return None

# Color palette
COLOR_PALETTE = ["#4269d0", "#efb118", "#ff725c", "#6cc5b0", "#3ca951", "#ff8ab7", "#a463f2", "#97bbf5"]

# Layout options
LAYOUT_OPTIONS = [
    {
        "title_to_chart": "TL",
        "image_to_chart": "R",
        "title_to_image": "TL",
        "chart_contains_title": False,
        "chart_contains_image": True
    },
    {
        "title_to_chart": "TL",
        "image_to_chart": "TR",
        "title_to_image": "L",
        "chart_contains_title": False,
        "chart_contains_image": False
    },
    {
        "title_to_chart": "TL",
        "image_to_chart": "L",
        "title_to_image": "T",
        "chart_contains_title": True,
        "chart_contains_image": True
    }
]

def parse_datafact_prompt(prompt: str, data_json: Dict[str, Any]) -> List[Dict[str, Any]]:
    """
    Prepare data for sending to a language model with the datafact_prompt.
    Uses the OpenAI API to generate data facts.
    
    Args:
        prompt: The datafact prompt template
        data_json: The input data JSON object
        
    Returns:
        List of data facts
    """
    logger.info("Preparing data for datafact generation")
    
    # Extract data for formatting
    metadata = data_json.get("metadata", {})
    columns = data_json.get("data", {}).get("columns", [])
    data_points = data_json.get("data", {}).get("data", [])
    
    if not columns or not data_points:
        logger.warning("No data columns or points found to generate data facts")
        return []
    
    # Format the data description for the prompt
    data_description = []
    
    # Add title and description if available
    if "title" in metadata:
        data_description.append(f"Title: {metadata['title']}")
    if "description" in metadata:
        data_description.append(f"Description: {metadata['description']}")
    
    # Add column information
    data_description.append("\nColumns:")
    for col in columns:
        col_info = []
        if "name" in col:
            col_info.append(f"Name: {col['name']}")
        if "description" in col:
            col_info.append(f"Description: {col['description']}")
        if "data_type" in col:
            col_info.append(f"Type: {col['data_type']}")
        if "role" in col:
            col_info.append(f"Role: {col['role']}")
        if "unit" in col and col["unit"] != "none":
            col_info.append(f"Unit: {col['unit']}")
        
        data_description.append(" | ".join(col_info))
    
    # Add sample data points (limited to first 10 for clarity)
    data_description.append("\nData:")
    max_samples = min(10, len(data_points))
    for i in range(max_samples):
        item = data_points[i]
        item_str = ", ".join([f"{k}: {v}" for k, v in item.items()])
        data_description.append(f"Row {i+1}: {item_str}")
    
    if len(data_points) > max_samples:
        data_description.append(f"...and {len(data_points) - max_samples} more rows")
    
    # Format the complete prompt with data
    formatted_data = "\n".join(data_description)
    final_prompt = prompt.replace("INPUT_TEXT", formatted_data)
    
    logger.info("Data formatted for datafact generation")
    
    # Call the OpenAI API to generate data facts
    try:
        logger.info("Calling language model API to generate data facts")
        response = query_openai(final_prompt)
        
        if not response:
            logger.warning("Failed to get response from language model API, falling back to simplified logic")
            return generate_fallback_datafacts(data_json)
        
        # Parse the JSON response
        try:
            # The model might wrap the JSON with markdown code blocks, remove them if present
            if response.startswith("```json"):
                response = response[7:]
            if response.endswith("```"):
                response = response[:-3]
            
            # Strip any leading/trailing whitespace
            response = response.strip()
            
            # Parse the JSON
            data_facts = json.loads(response)
            logger.info(f"Successfully parsed {len(data_facts)} data facts from API response")
            return data_facts
        except json.JSONDecodeError as e:
            logger.warning(f"Failed to parse data facts response as JSON: {e}")
            logger.warning(f"Raw response: {response}")
            return generate_fallback_datafacts(data_json)
    except Exception as e:
        logger.warning(f"Error calling language model API: {e}")
        return generate_fallback_datafacts(data_json)

def generate_fallback_datafacts(data_json: Dict[str, Any]) -> List[Dict[str, Any]]:
    """
    Generate fallback data facts when the API call fails
    
    Args:
        data_json: The input data JSON object
        
    Returns:
        List of data facts
    """
    logger.info("Using fallback logic to generate data facts")
    
    # Extract data for analysis
    columns = data_json.get("data", {}).get("columns", [])
    data_points = data_json.get("data", {}).get("data", [])
    
    # Find x and y columns
    x_column = next((col.get("name") for col in columns if col.get("role") == "x"), None)
    y_column = next((col.get("name") for col in columns if col.get("role") == "y"), None)
    
    if not x_column or not y_column:
        return []
    
    # Generate insights based on the data
    if any(col.get("data_type") == "numerical" for col in columns):
        sorted_data = sorted(data_points, key=lambda x: x.get(y_column, 0), reverse=True)
        
        # Find max and min values
        max_value = sorted_data[0]
        min_value = sorted_data[-1]
        
        # Calculate average
        total = sum(item.get(y_column, 0) for item in data_points)
        avg_value = total / len(data_points) if data_points else 0
        
        # Generate data facts
        data_facts = [
            {
                "type": "value",
                "score": 0.95,
                "annotation": f"{max_value.get(x_column)} has highest {y_column}",
                "reason": f"{max_value.get(x_column)} has the highest {y_column} at {max_value.get(y_column)}"
            },
            {
                "type": "value",
                "score": 0.85,
                "annotation": f"{min_value.get(x_column)} has lowest {y_column}",
                "reason": f"{min_value.get(x_column)} has the lowest {y_column} at {min_value.get(y_column)}"
            },
            {
                "type": "difference",
                "score": 0.80,
                "annotation": f"Gap between highest and lowest is {max_value.get(y_column) - min_value.get(y_column)}",
                "reason": f"The difference between the highest value ({max_value.get(y_column)}) and lowest value ({min_value.get(y_column)}) is {max_value.get(y_column) - min_value.get(y_column)}"
            },
            {
                "type": "overview",
                "score": 0.75,
                "annotation": f"Average {y_column} is {avg_value:.1f}",
                "reason": f"The average {y_column} across all {x_column} values is {avg_value:.1f}"
            }
        ]
        
        # If more than 3 data points, add trend insight
        if len(data_points) > 3:
            data_facts.append({
                "type": "trend",
                "score": 0.70,
                "annotation": f"{y_column} varies significantly across {x_column}",
                "reason": f"The {y_column} shows considerable variation across different {x_column} values, with a range of {max_value.get(y_column) - min_value.get(y_column)}"
            })
        
        # Add specific insights based on the data
        if len(data_points) > 2 and "Manchester" in max_value.get(x_column, "") and "Manchester" in sorted_data[1].get(x_column, ""):
            data_facts.append({
                "type": "comparison",
                "score": 0.65,
                "annotation": f"All Manchester terminals have long waiting times",
                "reason": f"Manchester airports occupy the top positions with waiting times of {', '.join([str(item.get(y_column)) for item in sorted_data[:3] if 'Manchester' in item.get(x_column, '')])}"
            })
        
        return data_facts
    
    return []

def module1_chart_type_recommender(data_json: Dict[str, Any]) -> Dict[str, Any]:
    """
    Module 1: Chart Type Recommender (simplified - returns empty)
    
    Args:
        data_json: The input data JSON object
        
    Returns:
        Updated JSON object with chart_type recommendations
    """
    logger.info("Module 1: Chart Type Recommender (simplified)")
    
    # Create a deep copy of the input data
    result = copy.deepcopy(data_json)
    
    # In the simplified version, this module is left empty as requested
    result["chart_type"] = []
    
    return result

def module2_datafact_generator(data_json: Dict[str, Any]) -> Dict[str, Any]:
    """
    Module 2: Data Fact Generator
    Uses the datafact_prompt to generate insights about the data
    
    Args:
        data_json: The input data JSON object
        
    Returns:
        Updated JSON object with datafacts
    """
    logger.info("Module 2: Data Fact Generator")
    
    # Create a deep copy of the input data
    result = copy.deepcopy(data_json)
    
    # Load the datafact prompt by importing it
    try:
        # Try to directly import the prompt from prompt.py
        from prompt import datafact_prompt
        logger.info("Successfully imported datafact_prompt from prompt.py")
    except ImportError as e:
        logger.warning(f"Failed to import datafact_prompt: {str(e)}")
        # Provide a minimal fallback prompt
        datafact_prompt = (
            "Analyze the following data and provide key insights.\n"
            "For each insight, include type, importance score, brief annotation, and detailed reason.\n"
            "Types can be: trend, proportion, outlier, difference, value, correlation, distribution, overview\n"
            "Respond in JSON format with an array of insight objects.\n"
            "INPUT_TEXT"
        )
    
    # Generate data facts
    data_facts = parse_datafact_prompt(datafact_prompt, data_json)
    
    # Add data facts to the result
    result["datafacts"] = data_facts
    logger.info(f"Generated {len(data_facts)} data facts")
    
    return result

def module3_title_generator(data_json: Dict[str, Any]) -> Dict[str, Any]:
    """
    Module 3: Title Generator
    Uses metadata title as main_title and description as sub_title
    
    Args:
        data_json: The input data JSON object
        
    Returns:
        Updated JSON object with titles
    """
    logger.info("Module 3: Title Generator")
    
    # Create a deep copy of the input data
    result = copy.deepcopy(data_json)
    
    # Get metadata
    metadata = data_json.get("metadata", {})
    
    # Set titles from metadata
    result["titles"] = {
        "main_title": metadata.get("title", ""),
        "sub_title": metadata.get("description", "")
    }
    
    return result

def module4_layout_recommender(data_json: Dict[str, Any]) -> Dict[str, Any]:
    """
    Module 4: Layout Recommender
    Randomly selects from predefined layout options and extracts variation from _extra
    
    Args:
        data_json: The input data JSON object
        
    Returns:
        Updated JSON object with layout and variation
    """
    logger.info("Module 4: Layout Recommender")
    
    # Create a deep copy of the input data
    result = copy.deepcopy(data_json)
    
    # Randomly select a layout from the predefined options
    layout = random.choice(LAYOUT_OPTIONS)
    
    # Extract variation from _extra if available
    variation = {}
    if "_extra" in data_json and "image_data" in data_json["_extra"] and "data" in data_json["_extra"]["image_data"]:
        extra_data = data_json["_extra"]["image_data"]["data"]
        
        variation = {
            "background": extra_data.get("background", "no"),
            "image_chart": extra_data.get("image_chart", "side"),
            "image_title": extra_data.get("image_title", "none"),
            "icon_mark": extra_data.get("icon_mark", "none"),
            "axis_label": extra_data.get("axis_label", "none"),
            "axes": {
                "x_axis": extra_data.get("axes", {}).get("x_axis", "yes"),
                "y_axis": extra_data.get("axes", {}).get("y_axis", "yes")
            }
        }
    else:
        # Default variation if not available in _extra
        variation = {
            "background": "no",
            "image_chart": "side",
            "image_title": "none",
            "icon_mark": "none",
            "axis_label": "none",
            "axes": {
                "x_axis": "yes",
                "y_axis": "yes"
            }
        }
    
    # Add layout and variation to the result
    result["layout"] = layout
    result["variation"] = variation
    
    return result

def module5_color_recommender(data_json: Dict[str, Any]) -> Dict[str, Any]:
    """
    Module 5: Color Recommender
    Assigns colors based on the specified rules:
    - If x+y, use a random color from palette as primary
    - If x+y+group, use group field for colors
    - Put remaining colors in available_colors
    
    Args:
        data_json: The input data JSON object
        
    Returns:
        Updated JSON object with color recommendations
    """
    logger.info("Module 5: Color Recommender")
    
    # Create a deep copy of the input data
    result = copy.deepcopy(data_json)
    
    # Extract column information
    columns = data_json.get("data", {}).get("columns", [])
    data_points = data_json.get("data", {}).get("data", [])
    
    # Find columns with roles
    x_column = next((col.get("name") for col in columns if col.get("role") == "x"), None)
    y_column = next((col.get("name") for col in columns if col.get("role") == "y"), None)
    group_column = next((col.get("name") for col in columns if col.get("role") == "group"), None)
    
    # Initialize colors object
    colors = {
        "field": {},
        "other": {},
        "available_colors": [],
        "background_color": "#FFFFFF",
        "text_color": "#000000"
    }
    
    # Create a copy of the color palette to work with
    available_colors = COLOR_PALETTE.copy()
    
    # Case: x+y+group
    if x_column and y_column and group_column:
        # Get unique group values
        group_values = set()
        for item in data_points:
            if group_column in item:
                group_values.add(item[group_column])
        
        # Assign colors to group values
        for i, group_value in enumerate(group_values):
            color_index = i % len(available_colors)
            colors["field"][group_value] = available_colors[color_index]
            # Remove used color
            available_colors.pop(color_index)
    
    # Case: x+y
    elif x_column and y_column:
        # Select a random color for primary
        primary_color_index = random.randint(0, len(available_colors) - 1)
        colors["other"]["primary"] = available_colors[primary_color_index]
        # Remove used color
        available_colors.pop(primary_color_index)
        
        # If there are more colors available, select one for secondary
        if available_colors:
            secondary_color_index = random.randint(0, len(available_colors) - 1)
            colors["other"]["secondary"] = available_colors[secondary_color_index]
            # Remove used color
            available_colors.pop(secondary_color_index)
    
    # Add remaining colors to available_colors
    colors["available_colors"] = available_colors
    
    # Add colors to the result
    result["colors"] = colors
    
    return result

def image_to_base64(image_path):
    """
    Convert an image file to base64-encoded string
    
    Args:
        image_path: Path to the image file
    
    Returns:
        Base64-encoded string
    """
    try:
        with open(image_path, "rb") as image_file:
            encoded_string = base64.b64encode(image_file.read()).decode('utf-8')
            return f"data:image/png;base64,{encoded_string}"
    except Exception as e:
        logger.warning(f"Failed to convert image to base64: {str(e)}")
        return None

def get_simplified_embedding(text):
    """
    Create a simplified "embedding" using word frequencies
    This is a very basic approximation of semantic similarity
    
    Args:
        text: Input text
    
    Returns:
        Dictionary of word frequencies
    """
    # Convert to lowercase and split into words
    words = text.lower().replace('_', ' ').replace('-', ' ').split()
    # Count word frequencies
    return Counter(words)

def calculate_similarity(embedding1, embedding2):
    """
    Calculate similarity between two simplified embeddings
    
    Args:
        embedding1: First embedding (Counter)
        embedding2: Second embedding (Counter)
    
    Returns:
        Similarity score
    """
    # Get common words
    common_words = set(embedding1.keys()) & set(embedding2.keys())
    
    if not common_words:
        return 0
    
    # Calculate dot product of common words
    similarity = sum(embedding1[word] * embedding2[word] for word in common_words)
    
    # Normalize
    norm1 = sum(val ** 2 for val in embedding1.values()) ** 0.5
    norm2 = sum(val ** 2 for val in embedding2.values()) ** 0.5
    
    if norm1 == 0 or norm2 == 0:
        return 0
    
    return similarity / (norm1 * norm2)

def get_bert_embedding(text):
    """
    Get BERT embedding for a given text
    
    Args:
        text: Input text
        
    Returns:
        Numpy array of embedding
    """
    if not USE_BERT:
        return get_simplified_embedding(text)
    
    try:
        # Add special tokens and convert to tensor
        inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=128)
        
        # Get model output (without gradient calculation for efficiency)
        with torch.no_grad():
            outputs = model(**inputs)
        
        # Use the [CLS] token embedding as the sentence embedding
        embedding = outputs.last_hidden_state[:, 0, :].numpy()
        return embedding[0]  # Return the first (and only) embedding
    except Exception as e:
        logger.warning(f"Error generating BERT embedding: {str(e)}. Using simplified embedding instead.")
        return get_simplified_embedding(text)

def precompute_icon_embeddings(icon_dir="test_data/icon"):
    """
    Precompute embeddings for all icons in the directory
    
    Args:
        icon_dir: Directory containing icons
        
    Returns:
        Dictionary mapping icon paths to their embeddings
    """
    global ICON_EMBEDDINGS_CACHE
    
    if not os.path.exists(icon_dir):
        logger.warning(f"Icon directory not found: {icon_dir}")
        return {}
    
    logger.info(f"Precomputing embeddings for icons in {icon_dir}")
    embeddings = {}
    
    for filename in os.listdir(icon_dir):
        if filename.lower().endswith('.png'):
            # Extract words from filename
            name_without_ext = os.path.splitext(filename)[0]
            # Get embedding
            if USE_BERT:
                embedding = get_bert_embedding(name_without_ext)
            else:
                embedding = get_simplified_embedding(name_without_ext)
            # Store in cache
            icon_path = os.path.join(icon_dir, filename)
            embeddings[icon_path] = embedding
    
    logger.info(f"Precomputed embeddings for {len(embeddings)} icons")
    ICON_EMBEDDINGS_CACHE = embeddings
    return embeddings

def calculate_bert_similarity(embedding1, embedding2):
    """
    Calculate cosine similarity between two BERT embeddings
    
    Args:
        embedding1: First embedding (numpy array)
        embedding2: Second embedding (numpy array)
        
    Returns:
        Similarity score
    """
    # If either embedding is a Counter (from simplified embedding), use the simplified similarity
    if isinstance(embedding1, Counter) or isinstance(embedding2, Counter):
        return calculate_similarity(embedding1, embedding2)
    
    # Calculate cosine similarity: dot(a, b) / (||a|| * ||b||)
    dot_product = np.dot(embedding1, embedding2)
    norm1 = np.linalg.norm(embedding1)
    norm2 = np.linalg.norm(embedding2)
    
    if norm1 == 0 or norm2 == 0:
        return 0
    
    return dot_product / (norm1 * norm2)

def find_best_matching_icon(text, icon_dir="test_data/icon"):
    """
    Find the best matching icon for a given text using BERT embeddings
    
    Args:
        text: Text to match
        icon_dir: Directory containing icons
        
    Returns:
        Path to the best matching icon
    """
    if not os.path.exists(icon_dir):
        logger.warning(f"Icon directory not found: {icon_dir}")
        return None
    
    # Get text embedding
    if USE_BERT:
        text_embedding = get_bert_embedding(text)
    else:
        text_embedding = get_simplified_embedding(text)
    
    # Ensure icon embeddings are precomputed
    global ICON_EMBEDDINGS_CACHE
    if not ICON_EMBEDDINGS_CACHE:
        precompute_icon_embeddings(icon_dir)
    
    best_match = None
    best_score = -1
    
    # Compare with all precomputed icon embeddings
    for icon_path, icon_embedding in ICON_EMBEDDINGS_CACHE.items():
        # Calculate similarity
        if USE_BERT:
            similarity = calculate_bert_similarity(text_embedding, icon_embedding)
        else:
            similarity = calculate_similarity(text_embedding, icon_embedding)
        
        if similarity > best_score:
            best_score = similarity
            best_match = icon_path
    
    # If no good match found, return a default icon
    if best_score < 0.2:  # Higher threshold for BERT embeddings
        # Return a random icon as fallback
        all_icons = list(ICON_EMBEDDINGS_CACHE.keys())
        if all_icons:
            return random.choice(all_icons)
    
    return best_match

def module6_image_recommender(data_json: Dict[str, Any]) -> Dict[str, Any]:
    """
    Module 6: Image Recommender
    Recommends images for only one attribute based on priority rules:
    - If x is temporal/numerical, assign images to group values
    - Otherwise, assign images to x values
    
    Args:
        data_json: The input data JSON object
        
    Returns:
        Updated JSON object with image recommendations
    """
    logger.info("Module 6: Image Recommender")
    
    # Create a deep copy of the input data
    result = copy.deepcopy(data_json)
    
    # Initialize images object
    images = {
        "field": {},
        "other": {}
    }
    
    # Check if icon directory exists
    icon_dir = "test_data/icon"
    if not os.path.exists(icon_dir):
        logger.warning(f"Icon directory not found: {icon_dir}. Cannot recommend images.")
        result["images"] = images
        return result
    
    # Extract column information and metadata
    columns = data_json.get("data", {}).get("columns", [])
    data_points = data_json.get("data", {}).get("data", [])
    metadata = data_json.get("metadata", {})
    
    # Ensure icon embeddings are precomputed for efficiency
    precompute_icon_embeddings(icon_dir)
    
    # Find x and group columns
    x_column = None
    x_is_temporal = False
    group_column = None
    
    for col in columns:
        if col.get("role") == "x":
            x_column = col.get("name")
            # Check if x is temporal
            if col.get("data_type") in ["temporal", "time", "date"]:
                x_is_temporal = True
        elif col.get("role") == "group":
            group_column = col.get("name")
    
    # Process field images based on priority rules
    if data_points:
        # Rule: If x is temporal, prioritize group values
        if x_is_temporal and group_column:
            logger.info(f"X column '{x_column}' is temporal, assigning images to group values")
            # Get unique group values
            group_values = []
            for item in data_points:
                if group_column in item and item[group_column] not in group_values:
                    group_values.append(item[group_column])
            
            # Find icons for each unique group value
            for group_value in group_values:
                logger.info(f"Finding icon for group value: {group_value}")
                icon_path = find_best_matching_icon(str(group_value), icon_dir)
                if icon_path:
                    base64_image = image_to_base64(icon_path)
                    if base64_image:
                        images["field"][str(group_value)] = base64_image
        
        # Rule: If x is not temporal or there's no group column, assign to x values
        elif x_column and not x_is_temporal:
            logger.info(f"Assigning images to non-temporal x values: {x_column}")
            # Get unique x values
            x_values = []
            for item in data_points:
                if x_column in item and item[x_column] not in x_values:
                    x_values.append(item[x_column])
            
            # Find icons for each unique x value
            for x_value in x_values:
                logger.info(f"Finding icon for x value: {x_value}")
                icon_path = find_best_matching_icon(str(x_value), icon_dir)
                if icon_path:
                    base64_image = image_to_base64(icon_path)
                    if base64_image:
                        images["field"][str(x_value)] = base64_image
    
    # Find icon for title
    if "title" in metadata:
        title_text = metadata["title"]
        logger.info(f"Finding icon for title: {title_text}")
        icon_path = find_best_matching_icon(title_text, icon_dir)
        if icon_path:
            base64_image = image_to_base64(icon_path)
            if base64_image:
                images["other"]["primary"] = base64_image
    
    # Add images to the result
    result["images"] = images
    logger.info(f"Generated {len(images['field'])} field images and {len(images['other'])} other images")
    
    return result

def process_modules(input_data: Dict[str, Any]) -> Dict[str, Any]:
    """
    Process data through all modules 1-6
    
    Args:
        input_data: The input data JSON object
        
    Returns:
        Processed data after running through all modules
    """
    # Apply each module in sequence
    result = module1_chart_type_recommender(input_data)
    result = module2_datafact_generator(result)
    result = module3_title_generator(result)
    result = module4_layout_recommender(result)
    result = module5_color_recommender(result)
    result = module6_image_recommender(result)
    
    return result

def main():
    """
    Main function to run the simplified modules 1-6
    
    If arguments are provided, process the specified input file to the specified output file.
    If no arguments are provided, process all files in test_data/new_data and save to test_data/data.
    """
    parser = argparse.ArgumentParser(description="Simplified ChartPipeline Modules 1-6")
    parser.add_argument("--input", help="Input JSON file path")
    parser.add_argument("--output", help="Output JSON file path")
    
    args = parser.parse_args()
    
    try:
        # If specific input and output are provided, process a single file
        if args.input and args.output:
            logger.info(f"Processing single file: {args.input} -> {args.output}")
            process_single_file(args.input, args.output)
        # Otherwise, process all files in test_data/new_data
        else:
            logger.info("No arguments provided, processing all files in test_data/new_data")
            process_directory()
        
        logger.info("Processing completed successfully")
        return 0
    
    except Exception as e:
        logger.error(f"Error processing data: {str(e)}")
        return 1

def process_single_file(input_path, output_path):
    """
    Process a single input file and save the result to the specified output path
    
    Args:
        input_path: Path to the input JSON file
        output_path: Path to save the output JSON file
    """
    # Load input data
    logger.info(f"Loading input data from {input_path}")
    with open(input_path, 'r', encoding='utf-8') as f:
        input_data = json.load(f)
    
    # Process data through modules 1-6
    output_data = process_modules(input_data)
    
    # Create output directory if it doesn't exist
    output_dir = os.path.dirname(output_path)
    if output_dir and not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # Save output data
    logger.info(f"Saving output data to {output_path}")
    with open(output_path, 'w', encoding='utf-8') as f:
        json.dump(output_data, f, indent=2, ensure_ascii=False)

def process_directory():
    """
    Process all JSON files in test_data/new_data directory and save results to test_data/data
    """
    input_dir = "test_data/new_data"
    output_dir = "test_data/data"
    
    # Check if input directory exists
    if not os.path.exists(input_dir):
        logger.warning(f"Input directory not found: {input_dir}")
        logger.info(f"Creating input directory: {input_dir}")
        os.makedirs(input_dir)
        return
    
    # Create output directory if it doesn't exist
    if not os.path.exists(output_dir):
        logger.info(f"Creating output directory: {output_dir}")
        os.makedirs(output_dir)
    
    # Get all JSON files in the input directory
    input_files = [f for f in os.listdir(input_dir) if f.endswith('.json')]
    
    if not input_files:
        logger.warning(f"No JSON files found in {input_dir}")
        return
    
    logger.info(f"Found {len(input_files)} JSON files to process")
    
    # Process each file
    for filename in input_files:
        input_path = os.path.join(input_dir, filename)
        output_path = os.path.join(output_dir, filename)
        
        try:
            logger.info(f"Processing file: {filename}")
            process_single_file(input_path, output_path)
        except Exception as e:
            logger.error(f"Error processing file {filename}: {str(e)}")

if __name__ == "__main__":
    sys.exit(main())