#!/usr/bin/env python3
"""
Fast Parallel Synthetic ECG Generator using ecg-image-kit

Generates synthetic ECG images from PTB-XL records in parallel.
Run with venv: source ecgkit_venv/bin/activate && python generate_synthetic_parallel.py

Usage:
    python generate_synthetic_parallel.py --n_samples 50000 --n_workers 16
"""

import os
import sys
import argparse
import subprocess
import random
import json
import shutil
import numpy as np
import wfdb
from scipy import signal as scipy_signal
from pathlib import Path
from multiprocessing import Pool, cpu_count
from tqdm import tqdm

# Constants for GT extraction
TARGET_HEIGHT = 1696
ZERO_MV = np.array([703.5, 987.5, 1271.5, 1531.5])
MV_TO_PIXEL = 78.5
T0, T1 = 235, 4161
OUTPUT_GT_WIDTH = T1 - T0  # 3926

# Paths
PROJECT_ROOT = Path(__file__).parent.parent
ECG_IMAGE_KIT = PROJECT_ROOT / 'ecg-image-kit' / 'codes' / 'ecg-image-generator'
PTBXL_DIR = Path('/data/ecg-digitization/ptbxl/physionet.org/files/ptb-xl/1.0.3')

# Grid color mappings (0001-0012 styles)
STYLE_CONFIGS = [
    {'grid_color': 5, 'augment': False, 'wrinkles': False},  # 0001: Clean red
    {'grid_color': 5, 'augment': True, 'wrinkles': True},    # 0002: Aged red
    {'grid_color': 2, 'augment': False, 'wrinkles': False},  # 0003: Green
    {'grid_color': 2, 'augment': True, 'wrinkles': True},    # 0004: Aged green
    {'grid_color': 1, 'augment': False, 'wrinkles': False},  # 0005: Blue
    {'grid_color': 3, 'augment': True, 'wrinkles': False},   # 0006: Grey
    {'grid_color': 4, 'augment': True, 'wrinkles': True},    # 0007: Yellow aged
    {'grid_color': 2, 'augment': True, 'wrinkles': False},   # 0008: Faded green
    {'grid_color': 6, 'augment': False, 'wrinkles': False},  # 0009: Orange
    {'grid_color': 5, 'augment': True, 'wrinkles': False},   # 0010: High contrast
    {'grid_color': 5, 'augment': True, 'wrinkles': True},    # 0011: Low contrast
    {'grid_color': 0, 'augment': False, 'wrinkles': False},  # 0012: No grid
]


def find_ptbxl_records():
    """Find all PTB-XL record files."""
    records = []
    
    for records_dir in ['records500', 'records100']:
        rec_path = PTBXL_DIR / records_dir
        if rec_path.exists():
            for subdir in rec_path.iterdir():
                if subdir.is_dir():
                    for hea in subdir.glob('*.hea'):
                        dat = hea.with_suffix('.dat')
                        if dat.exists():
                            records.append((str(hea), str(dat)))
            if records:
                break
    
    return records


def extract_gt_signal(hea_file):
    """Extract ground truth signal from PTB-XL record as pixel Y coordinates."""
    try:
        record_path = hea_file.replace('.hea', '')
        record = wfdb.rdrecord(record_path)
        
        signals = record.p_signal  # [samples, leads]
        sig_names = [n.upper() for n in record.sig_name]
        fs = record.fs
        
        # Resample to 500Hz if needed
        if fs != 500:
            num_samples = int(signals.shape[0] * 500 / fs)
            signals = scipy_signal.resample(signals, num_samples, axis=0)
        
        # Map lead names to indices
        lead_signals = {}
        for i, name in enumerate(sig_names):
            lead_signals[name] = signals[:, i]
        
        # 4-row layout matching Kaggle format
        # Row 0: I, aVR, V1, V4 (each 2.5s segment)
        # Row 1: II, aVL, V2, V5
        # Row 2: III, aVF, V3, V6
        # Row 3: Full Lead II (10s rhythm strip)
        row_leads = [
            ['I', 'AVR', 'V1', 'V4'],
            ['II', 'AVL', 'V2', 'V5'],
            ['III', 'AVF', 'V3', 'V6'],
        ]
        
        segment_width = OUTPUT_GT_WIDTH // 4  # ~981 pixels per segment
        gt_pixels = np.zeros((4, OUTPUT_GT_WIDTH), dtype=np.float32)
        
        # First 3 rows: 4 leads each, 2.5s segments
        for row_idx in range(3):
            for col_idx in range(4):
                lead_name = row_leads[row_idx][col_idx]
                if lead_name not in lead_signals:
                    continue
                    
                lead_data = lead_signals[lead_name]
                # 2.5s at 500Hz = 1250 samples
                start = col_idx * 1250
                end = min(start + 1250, len(lead_data))
                segment_mv = lead_data[start:end]
                
                if len(segment_mv) == 0:
                    continue
                
                col_start = col_idx * segment_width
                col_end = (col_idx + 1) * segment_width if col_idx < 3 else OUTPUT_GT_WIDTH
                
                # Resample to pixel width
                x_old = np.linspace(0, 1, len(segment_mv))
                x_new = np.linspace(0, 1, col_end - col_start)
                segment_resampled = np.interp(x_new, x_old, segment_mv)
                
                # Convert mV to pixel Y coordinates
                pixel_y = ZERO_MV[row_idx] - segment_resampled * MV_TO_PIXEL
                gt_pixels[row_idx, col_start:col_end] = pixel_y
        
        # Row 4: Full Lead II rhythm strip (10s)
        if 'II' in lead_signals:
            lead_ii = lead_signals['II']
            x_old = np.linspace(0, 1, len(lead_ii))
            x_new = np.linspace(0, 1, OUTPUT_GT_WIDTH)
            lead_ii_resampled = np.interp(x_new, x_old, lead_ii)
            gt_pixels[3, :] = ZERO_MV[3] - lead_ii_resampled * MV_TO_PIXEL
        
        return gt_pixels
        
    except Exception as e:
        return None


def generate_single(args):
    """Generate a single synthetic image."""
    idx, hea_file, dat_file, output_dir, style_idx = args
    
    try:
        sample_id = f'syn_{idx:08d}'
        style_str = f'{(style_idx % 12) + 1:04d}'
        style_config = STYLE_CONFIGS[style_idx % 12]
        
        temp_dir = Path(f'/tmp/ecg_gen_{idx}_{os.getpid()}')
        temp_dir.mkdir(parents=True, exist_ok=True)
        
        # Build command
        # Note: -st (start_index) should be 0, -se (seed) controls randomization
        cmd = [
            'python', 'gen_ecg_image_from_data.py',
            '-i', dat_file,
            '-hea', hea_file,
            '-o', str(temp_dir),
            '-st', '0',
            '-se', str(idx),
            '-r', '200',
            '--num_columns', '4',
            '--full_mode', 'II',
        ]
        
        if style_config['grid_color'] > 0:
            cmd.extend(['--standard_grid_color', str(style_config['grid_color'])])
        else:
            cmd.extend(['--random_grid_present', '0'])
        
        if style_config['augment']:
            cmd.append('--augment')
        
        if style_config['wrinkles']:
            cmd.append('--wrinkles')
        
        # Set MPLCONFIGDIR to avoid lock conflicts
        env = os.environ.copy()
        env['MPLCONFIGDIR'] = f'/tmp/mpl_{idx}_{os.getpid()}'
        
        result = subprocess.run(
            cmd,
            cwd=str(ECG_IMAGE_KIT),
            capture_output=True,
            text=True,
            timeout=120,
            env=env
        )
        
        # Find and move generated image
        gen_images = list(temp_dir.glob('*.png'))
        output_path = None
        gt_path = None
        
        if gen_images:
            output_path = Path(output_dir) / 'raw' / f'{sample_id}-{style_str}.png'
            shutil.move(str(gen_images[0]), str(output_path))
            
            # Extract and save GT
            gt_pixels = extract_gt_signal(hea_file)
            if gt_pixels is not None:
                gt_path = Path(output_dir) / 'gt' / f'{sample_id}-{style_str}.csv'
                gt_path.parent.mkdir(parents=True, exist_ok=True)
                # Save as CSV with shape [N, 4] (pixel Y coordinates)
                np.savetxt(gt_path, gt_pixels.T, delimiter=',', fmt='%.2f')
        
        # Clean up
        shutil.rmtree(temp_dir, ignore_errors=True)
        mpl_dir = Path(f'/tmp/mpl_{idx}_{os.getpid()}')
        if mpl_dir.exists():
            shutil.rmtree(mpl_dir, ignore_errors=True)
        
        if output_path and output_path.exists() and gt_path and gt_path.exists():
            return {
                'sample_id': sample_id,
                'style': style_str,
                'path': str(output_path),
                'gt_path': str(gt_path),
                'source': hea_file,
            }
        return None
        
    except Exception as e:
        return None


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--output_dir', type=str, default='/data/ecg-digitization/synthetic_ecgkit')
    parser.add_argument('--n_samples', type=int, default=50000)
    parser.add_argument('--n_workers', type=int, default=16)
    parser.add_argument('--resume', action='store_true', help='Resume from existing progress')
    args = parser.parse_args()
    
    output_dir = Path(args.output_dir)
    (output_dir / 'raw').mkdir(parents=True, exist_ok=True)
    (output_dir / 'gt').mkdir(parents=True, exist_ok=True)
    
    # Find PTB-XL records
    print("Finding PTB-XL records...")
    records = find_ptbxl_records()
    print(f"Found {len(records)} records")
    
    if not records:
        print("No PTB-XL records found!")
        return
    
    # Check existing files for resume
    existing_indices = set()
    if args.resume:
        raw_dir = output_dir / 'raw'
        for f in raw_dir.glob('syn_*.png'):
            try:
                # Extract index from syn_00000123-0001.png
                idx = int(f.stem.split('_')[1].split('-')[0])
                existing_indices.add(idx)
            except:
                pass
        print(f"Found {len(existing_indices)} existing samples, resuming...")
    
    # Prepare tasks (skip existing if resuming)
    print(f"Preparing generation tasks...")
    tasks = []
    random.seed(42)  # Deterministic for resume
    for i in range(args.n_samples):
        hea, dat = random.choice(records)
        style_idx = 0  # Always style 0001 (clean red grid)
        if i not in existing_indices:
            tasks.append((i, hea, dat, str(output_dir), style_idx))
    
    print(f"Tasks to generate: {len(tasks)} (skipped {len(existing_indices)} existing)")
    
    if not tasks:
        print("All samples already generated!")
        return
    
    # Generate in parallel
    print(f"Generating with {args.n_workers} workers...")
    results = []
    
    with Pool(args.n_workers) as pool:
        for result in tqdm(pool.imap_unordered(generate_single, tasks), 
                          total=len(tasks), desc="Generating"):
            if result:
                results.append(result)
    
    print(f"\nSuccessfully generated {len(results)} images")
    
    # Save manifest
    manifest_path = output_dir / 'manifest.json'
    with open(manifest_path, 'w') as f:
        json.dump(results, f, indent=2)
    
    print(f"Manifest saved to {manifest_path}")
    print(f"Raw images in {output_dir / 'raw'}")


if __name__ == '__main__':
    main()