#!/usr/bin/env python3
"""Simple serial synthetic image generation - 0001 style only."""

import os
import sys
import subprocess
from pathlib import Path
from tqdm import tqdm

# Paths
ECGKIT_DIR = '/home/azureuser/ecg-digitization/ecg-image-kit/codes/ecg-image-generator'
PTBXL_DIR = '/data/ecg-digitization/ptbxl/physionet.org/files/ptb-xl/1.0.3/records500'
OUTPUT_DIR = '/data/ecg-digitization/synthetic_0001'

# Create output dirs
os.makedirs(f'{OUTPUT_DIR}/raw', exist_ok=True)
os.makedirs(f'{OUTPUT_DIR}/gt', exist_ok=True)

# Set env to avoid matplotlib issues
os.environ['MPLBACKEND'] = 'Agg'
os.environ['MPLCONFIGDIR'] = f'{OUTPUT_DIR}/mplconfig'
os.makedirs(os.environ['MPLCONFIGDIR'], exist_ok=True)

# Get all .hea files
hea_files = sorted(Path(PTBXL_DIR).rglob('*.hea'))
print(f"Found {len(hea_files)} records")

# Check what's already done
existing = set(p.stem.replace('syn_', '').replace('-0001', '') for p in Path(f'{OUTPUT_DIR}/raw').glob('*.png'))
print(f"Already generated: {len(existing)}")

# Process each
for hea_path in tqdm(hea_files):
    record_id = hea_path.stem
    
    # Skip if done
    if record_id in existing:
        continue
    
    input_file = str(hea_path.with_suffix(''))  # Remove .hea extension
    output_file = f'{OUTPUT_DIR}/raw/syn_{record_id}-0001.png'
    
    # 0001 style: standard_grid_color=5 (red), NO augment, NO wrinkles
    cmd = [
        'python', f'{ECGKIT_DIR}/gen_ecg_image_from_data.py',
        '-i', input_file,
        '-o', output_file,
        '--input_directory', '',
        '--output_directory', '',
        '-r', '200',  # resolution
        '--pad_inches', '0',
        '--print_header', 'False',
        '--standard_grid_color', '5',  # Red grid (0001 style)
        '-se', '10',  # seed
        '--store_config', '2',
    ]
    
    try:
        result = subprocess.run(
            cmd,
            cwd=ECGKIT_DIR,
            capture_output=True,
            text=True,
            timeout=120,
            env={**os.environ, 'MPLBACKEND': 'Agg'}
        )
        if result.returncode != 0:
            print(f"Error {record_id}: {result.stderr[:200]}")
    except subprocess.TimeoutExpired:
        print(f"Timeout {record_id}")
    except Exception as e:
        print(f"Exception {record_id}: {e}")

print("Done!")