#!/usr/bin/env python3
"""Create a simple Image2Struct_Latex test dataset for VLMEvalKit."""

import base64
import io
import os
import pandas as pd

# Try to use matplotlib to render LaTeX
try:
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    HAS_MATPLOTLIB = True
except ImportError:
    HAS_MATPLOTLIB = False


def latex_to_base64(latex_code: str, fontsize: int = 20) -> str:
    """Render LaTeX to base64-encoded PNG image."""
    if not HAS_MATPLOTLIB:
        raise ImportError("matplotlib is required to render LaTeX")
    
    fig, ax = plt.subplots(figsize=(6, 2))
    ax.axis('off')
    
    # Wrap in math mode if not already
    if not latex_code.startswith('$'):
        latex_code = f'${latex_code}$'
    
    ax.text(0.5, 0.5, latex_code, fontsize=fontsize, ha='center', va='center',
            transform=ax.transAxes)
    
    # Save to buffer
    buf = io.BytesIO()
    fig.savefig(buf, format='png', bbox_inches='tight', dpi=150,
                facecolor='white', edgecolor='none')
    plt.close(fig)
    
    # Encode to base64
    buf.seek(0)
    img_base64 = base64.b64encode(buf.read()).decode('utf-8')
    return img_base64


# Sample LaTeX formulas for testing
LATEX_SAMPLES = [
    # Simple formulas
    (r"E = mc^2", "E = mc^2"),
    (r"\frac{a}{b}", r"\frac{a}{b}"),
    (r"x^2 + y^2 = z^2", r"x^2 + y^2 = z^2"),
    (r"\sqrt{a^2 + b^2}", r"\sqrt{a^2 + b^2}"),
    (r"\int_0^1 x^2 dx", r"\int_0^1 x^2 dx"),
    
    # Medium complexity
    (r"\sum_{i=1}^{n} i = \frac{n(n+1)}{2}", r"\sum_{i=1}^{n} i = \frac{n(n+1)}{2}"),
    (r"\lim_{x \to 0} \frac{\sin x}{x} = 1", r"\lim_{x \to 0} \frac{\sin x}{x} = 1"),
    (r"e^{i\pi} + 1 = 0", r"e^{i\pi} + 1 = 0"),
    (r"\binom{n}{k} = \frac{n!}{k!(n-k)!}", r"\binom{n}{k} = \frac{n!}{k!(n-k)!}"),
    (r"\alpha + \beta = \gamma", r"\alpha + \beta = \gamma"),
    
    # More complex
    (r"\frac{d}{dx}\left(\frac{f}{g}\right) = \frac{f'g - fg'}{g^2}", 
     r"\frac{d}{dx}\left(\frac{f}{g}\right) = \frac{f'g - fg'}{g^2}"),
    (r"\nabla \cdot \vec{E} = \frac{\rho}{\epsilon_0}", 
     r"\nabla \cdot \vec{E} = \frac{\rho}{\epsilon_0}"),
]


def create_dataset(output_path: str):
    """Create Image2Struct_Latex dataset TSV file."""
    
    data = []
    prompt = (
        "Convert the image to LaTeX code. "
        "Output only the LaTeX code that reproduces the content shown in the image. "
        "Do not include \\documentclass or \\usepackage commands."
    )
    
    for idx, (latex_display, latex_answer) in enumerate(LATEX_SAMPLES):
        try:
            img_base64 = latex_to_base64(latex_display)
            data.append({
                'index': idx,
                'question': prompt,
                'image': img_base64,
                'answer': latex_answer,
                'domain': 'latex',
            })
            print(f"Created sample {idx}: {latex_display[:30]}...")
        except Exception as e:
            print(f"Failed to create sample {idx}: {e}")
    
    if not data:
        print("No samples created!")
        return
    
    # Create DataFrame and save
    df = pd.DataFrame(data)
    df.to_csv(output_path, sep='\t', index=False)
    print(f"\nSaved {len(data)} samples to {output_path}")


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--output', '-o', default='/home/chengyuew/LMUData/Image2Struct_Latex.tsv',
                        help='Output TSV file path')
    args = parser.parse_args()
    
    create_dataset(args.output)
