#!/usr/bin/env python3
"""
Generate 0001-style synthetic ECG images from multiple data sources.

Supports:
- PTB-XL (already done)
- Chapman-Shaoxing
- Georgia
- CPSC 2018
- Ningbo

Each source is kept in separate folders for tracking:
- synthetic_0001_{source}/raw/      - Raw generated images
- synthetic_0001_{source}/gt/       - Ground truth CSV files
- synthetic_0001_{source}_processed/raw/ - Stage0/1 processed images
- synthetic_0001_{source}_processed/gt/  - GT (symlinked from raw)

GT Format: CSV with 3926 lines, 4 columns (y-pixel coords for each row)
Image Format: 1700x2200 PNG (raw), processed through Stage0/1

Usage:
    python generate_synthetic_0001_multisource.py --source chapman --workers 8
    python generate_synthetic_0001_multisource.py --source georgia --workers 8
    python generate_synthetic_0001_multisource.py --source cpsc --workers 8
    python generate_synthetic_0001_multisource.py --source ningbo --workers 8
    python generate_synthetic_0001_multisource.py --source all --workers 8
"""

import os
import sys
import argparse
import json
import random
import subprocess
from pathlib import Path
from concurrent.futures import ProcessPoolExecutor, as_completed
from tqdm import tqdm
import numpy as np
import cv2
import shutil
import traceback

# Try to import wfdb
try:
    import wfdb
    from scipy import signal as scipy_signal
except ImportError:
    print("Installing required packages...")
    subprocess.run([sys.executable, '-m', 'pip', 'install', 'wfdb', 'scipy'], check=True)
    import wfdb
    from scipy import signal as scipy_signal

# =============================================================================
# Configuration
# =============================================================================
PROJECT_ROOT = Path(__file__).parent.parent
ECG_IMAGE_KIT = PROJECT_ROOT / 'ecg-image-kit' / 'codes' / 'ecg-image-generator'
BASELINE_PATH = PROJECT_ROOT / 'data' / 'hengck23-submit-physionet' / 'hengck23-submit-physionet'
DATA_ROOT = Path('/data/ecg-digitization')

# Data source directories
DATA_SOURCES = {
    'ptbxl': DATA_ROOT / 'ptbxl' / 'physionet.org' / 'files' / 'ptb-xl' / '1.0.3',
    'chapman': DATA_ROOT / 'chapman',
    'georgia': DATA_ROOT / 'georgia',
    'cpsc': DATA_ROOT / 'cpsc',
    'ningbo': DATA_ROOT / 'ningbo',
}

# Target dimensions (match competition raw format)
RAW_WIDTH = 2200
RAW_HEIGHT = 1700

# Target dimensions (match competition Stage 1 output)
TARGET_WIDTH = 4352
TARGET_HEIGHT = 1696

# GT parameters
T0 = 235
T1 = 4161
OUTPUT_GT_WIDTH = T1 - T0  # 3926

# Baseline y-positions for 4 rows (in full image coordinates)
ZERO_MV = np.array([703.5, 987.5, 1271.5, 1531.5])
MV_TO_PIXEL = 78.5  # pixels per mV

# Lead layout (3 columns x 4 rows, each column shows 2.5s)
LEAD_LAYOUT = [
    ['I', 'aVR', 'V1', 'V4'],      # Row indices for each column
    ['II', 'aVL', 'V2', 'V5'],
    ['III', 'aVF', 'V3', 'V6'],
]


# =============================================================================
# Record Discovery Functions
# =============================================================================
def find_records_ptbxl(data_dir):
    """Find PTB-XL record files (500Hz high-res)."""
    records = []
    data_dir = Path(data_dir)
    
    records500 = data_dir / 'records500'
    if records500.exists():
        for subdir in sorted(records500.iterdir()):
            if subdir.is_dir():
                for hea in subdir.glob('*_hr.hea'):
                    mat = hea.with_suffix('.dat')
                    if mat.exists():
                        records.append(str(hea))
    return records


def find_records_chapman(data_dir):
    """Find Chapman-Shaoxing record files."""
    records = []
    data_dir = Path(data_dir)
    
    # Chapman uses WFDBRecords/XX/XXX/JSXXXXX.hea structure
    for hea in data_dir.rglob('JS*.hea'):
        mat = hea.with_suffix('.mat')
        if mat.exists():
            records.append(str(hea))
    return sorted(records)


def find_records_georgia(data_dir):
    """Find Georgia record files."""
    records = []
    data_dir = Path(data_dir)
    
    # Georgia uses georgia/gX/EXXXXX.hea structure
    for hea in data_dir.rglob('E*.hea'):
        mat = hea.with_suffix('.mat')
        if mat.exists():
            records.append(str(hea))
    return sorted(records)


def find_records_cpsc(data_dir):
    """Find CPSC 2018 record files."""
    records = []
    data_dir = Path(data_dir)
    
    # CPSC uses cpsc_2018/gX/AXXXX.hea structure
    for hea in data_dir.rglob('A*.hea'):
        mat = hea.with_suffix('.mat')
        if mat.exists():
            records.append(str(hea))
    return sorted(records)


def find_records_ningbo(data_dir):
    """Find Ningbo record files."""
    records = []
    data_dir = Path(data_dir)
    
    # Ningbo uses similar structure
    for hea in data_dir.rglob('*.hea'):
        mat = hea.with_suffix('.mat')
        if mat.exists():
            records.append(str(hea))
    return sorted(records)


FIND_FUNCTIONS = {
    'ptbxl': find_records_ptbxl,
    'chapman': find_records_chapman,
    'georgia': find_records_georgia,
    'cpsc': find_records_cpsc,
    'ningbo': find_records_ningbo,
}


# =============================================================================
# Signal Extraction
# =============================================================================
def load_ecg_signals(hea_path):
    """Load ECG signals from WFDB record."""
    try:
        record_path = hea_path.replace('.hea', '')
        record = wfdb.rdrecord(record_path)
        
        signals = record.p_signal  # [samples, leads]
        sig_names = [name.upper() for name 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)
            fs = 500
        
        # Map lead names to signals
        lead_signals = {}
        for i, name in enumerate(sig_names):
            # Normalize lead names
            name_norm = name.replace('AVR', 'aVR').replace('AVL', 'aVL').replace('AVF', 'aVF')
            lead_signals[name_norm] = signals[:, i]
        
        return lead_signals, fs
        
    except Exception as e:
        return None, None


def create_gt_csv(lead_signals, fs=500):
    """
    Create ground truth CSV in the expected format.
    
    GT Format: 3926 lines, 4 columns
    Each line is x-position (from T0 to T1)
    Each column is y-pixel coordinate for that row
    
    The y-coordinate is in FULL image space (1696 height after resize).
    """
    if lead_signals is None:
        return None
    
    # 3 columns, each showing 2.5s of data
    samples_per_column = int(2.5 * fs)  # 1250 samples at 500Hz
    pixels_per_column = OUTPUT_GT_WIDTH // 3  # ~1308 pixels
    
    gt_data = np.zeros((OUTPUT_GT_WIDTH, 4), dtype=np.float32)
    
    # For each row (0-3), we need to get the appropriate leads
    for row_idx in range(4):
        row_baseline = ZERO_MV[row_idx]
        
        for col_idx in range(3):
            # Get lead name for this position
            lead_name = LEAD_LAYOUT[col_idx][row_idx]
            
            # Handle lead name variations
            lead_name_variants = [lead_name, lead_name.upper(), lead_name.lower()]
            if lead_name.startswith('aV'):
                lead_name_variants.extend([lead_name.upper(), lead_name[0].upper() + lead_name[1:]])
            
            signal = None
            for variant in lead_name_variants:
                if variant in lead_signals:
                    signal = lead_signals[variant]
                    break
            
            if signal is None:
                # Use baseline if lead not found
                col_start = col_idx * pixels_per_column
                col_end = (col_idx + 1) * pixels_per_column if col_idx < 2 else OUTPUT_GT_WIDTH
                gt_data[col_start:col_end, row_idx] = row_baseline
                continue
            
            # Extract segment (2.5s starting at col_idx * 2.5s)
            seg_start = col_idx * samples_per_column
            seg_end = min(seg_start + samples_per_column, len(signal))
            segment = signal[seg_start:seg_end]
            
            if len(segment) == 0:
                col_start = col_idx * pixels_per_column
                col_end = (col_idx + 1) * pixels_per_column if col_idx < 2 else OUTPUT_GT_WIDTH
                gt_data[col_start:col_end, row_idx] = row_baseline
                continue
            
            # Resample to pixel coordinates
            col_start = col_idx * pixels_per_column
            col_end = (col_idx + 1) * pixels_per_column if col_idx < 2 else OUTPUT_GT_WIDTH
            num_pixels = col_end - col_start
            
            x_old = np.linspace(0, 1, len(segment))
            x_new = np.linspace(0, 1, num_pixels)
            signal_resampled = np.interp(x_new, x_old, segment)
            
            # Convert mV to pixel coordinates
            # y = baseline - (signal_mV * MV_TO_PIXEL)
            # (negative because y increases downward)
            y_pixels = row_baseline - (signal_resampled * MV_TO_PIXEL)
            
            # Clip to valid range
            y_pixels = np.clip(y_pixels, 0, TARGET_HEIGHT - 1)
            
            gt_data[col_start:col_end, row_idx] = y_pixels
    
    return gt_data


# =============================================================================
# Image Generation
# =============================================================================
def generate_single_sample(args):
    """Generate a single synthetic sample."""
    idx, hea_path, source, output_dir, temp_base = args
    
    sample_id = f'syn_{source}_{idx:08d}'
    style = '0001'
    
    raw_dir = output_dir / 'raw'
    gt_dir = output_dir / 'gt'
    
    out_img = raw_dir / f'{sample_id}-{style}.png'
    out_gt = gt_dir / f'{sample_id}-{style}.csv'
    
    # Skip if both exist
    if out_img.exists() and out_gt.exists():
        return True, idx, "exists"
    
    temp_dir = temp_base / f'temp_{idx}'
    
    try:
        # Load signals and create GT first (faster, can skip if fails)
        lead_signals, fs = load_ecg_signals(hea_path)
        if lead_signals is None:
            return False, idx, "failed to load signals"
        
        gt_data = create_gt_csv(lead_signals, fs)
        if gt_data is None:
            return False, idx, "failed to create GT"
        
        # Create temp directory
        temp_dir.mkdir(parents=True, exist_ok=True)
        
        # Get dat file path
        dat_path = hea_path.replace('.hea', '.mat')
        if not Path(dat_path).exists():
            dat_path = hea_path.replace('.hea', '.dat')
        
        # Generate image with ecg-image-kit
        cmd = [
            'python', 'gen_ecg_image_from_data.py',
            '-i', dat_path,
            '-hea', hea_path,
            '-o', str(temp_dir),
            '-st', str(idx),
            '-se', str(idx),
            '-r', '200',  # DPI
            '--num_columns', '4',
            '--full_mode', 'II',
            '--standard_grid_color', '5',  # Red grid (0001 style)
            '--store_config', '0',
        ]
        
        # Set matplotlib to use non-interactive backend
        env = os.environ.copy()
        env['MPLBACKEND'] = 'Agg'
        env['MPLCONFIGDIR'] = str(temp_dir)
        
        result = subprocess.run(
            cmd,
            cwd=str(ECG_IMAGE_KIT),
            capture_output=True,
            text=True,
            timeout=120,
            env=env
        )
        
        if result.returncode != 0:
            shutil.rmtree(temp_dir, ignore_errors=True)
            return False, idx, f"ecg-kit error: {result.stderr[:100]}"
        
        # Find generated image
        gen_imgs = list(temp_dir.glob('*.png'))
        if not gen_imgs:
            shutil.rmtree(temp_dir, ignore_errors=True)
            return False, idx, "no image generated"
        
        # Load and resize to raw target size (1700x2200)
        img = cv2.imread(str(gen_imgs[0]))
        if img is None:
            shutil.rmtree(temp_dir, ignore_errors=True)
            return False, idx, "failed to read image"
        
        img_resized = cv2.resize(img, (RAW_WIDTH, RAW_HEIGHT), interpolation=cv2.INTER_LINEAR)
        
        # Save raw image
        raw_dir.mkdir(parents=True, exist_ok=True)
        cv2.imwrite(str(out_img), img_resized)
        
        # Save GT as CSV
        gt_dir.mkdir(parents=True, exist_ok=True)
        np.savetxt(out_gt, gt_data, delimiter=',', fmt='%.2f')
        
        # Cleanup
        shutil.rmtree(temp_dir, ignore_errors=True)
        
        return True, idx, "success"
        
    except Exception as e:
        shutil.rmtree(temp_dir, ignore_errors=True)
        return False, idx, f"error: {str(e)[:100]}"


# =============================================================================
# Main
# =============================================================================
def main():
    parser = argparse.ArgumentParser(description='Generate 0001-style synthetic ECGs from multiple sources')
    parser.add_argument('--source', type=str, required=True,
                       choices=['ptbxl', 'chapman', 'georgia', 'cpsc', 'ningbo', 'all'],
                       help='Data source to use')
    parser.add_argument('--workers', type=int, default=8,
                       help='Number of parallel workers')
    parser.add_argument('--limit', type=int, default=None,
                       help='Limit number of samples (for testing)')
    parser.add_argument('--output_base', type=str, default='/data/ecg-digitization',
                       help='Base output directory')
    args = parser.parse_args()
    
    output_base = Path(args.output_base)
    
    # Determine sources to process
    if args.source == 'all':
        sources = ['chapman', 'georgia', 'cpsc', 'ningbo']  # Skip ptbxl (already done)
    else:
        sources = [args.source]
    
    for source in sources:
        print(f"\n{'='*60}")
        print(f"Processing source: {source}")
        print(f"{'='*60}")
        
        # Check if data directory exists
        data_dir = DATA_SOURCES.get(source)
        if data_dir is None or not data_dir.exists():
            print(f"Data directory not found: {data_dir}")
            continue
        
        # Output directories
        output_dir = output_base / f'synthetic_0001_{source}'
        output_dir.mkdir(parents=True, exist_ok=True)
        (output_dir / 'raw').mkdir(exist_ok=True)
        (output_dir / 'gt').mkdir(exist_ok=True)
        
        temp_base = output_dir / 'temp'
        temp_base.mkdir(exist_ok=True)
        
        # Find records
        find_func = FIND_FUNCTIONS.get(source)
        if find_func is None:
            print(f"No find function for source: {source}")
            continue
        
        print(f"Finding records in {data_dir}...")
        records = find_func(data_dir)
        print(f"Found {len(records)} records")
        
        if len(records) == 0:
            print("No records found, skipping")
            continue
        
        # Apply limit if specified
        if args.limit and len(records) > args.limit:
            records = records[:args.limit]
            print(f"Limited to {len(records)} records")
        
        # Prepare tasks
        tasks = [
            (i, hea, source, output_dir, temp_base)
            for i, hea in enumerate(records)
        ]
        
        # Generate in parallel
        success = 0
        failed = 0
        exists = 0
        
        print(f"Generating {len(tasks)} 0001-style synthetic images...")
        
        with ProcessPoolExecutor(max_workers=args.workers) as executor:
            futures = {executor.submit(generate_single_sample, task): task for task in tasks}
            
            pbar = tqdm(as_completed(futures), total=len(futures), desc=f"[{source}]")
            for future in pbar:
                try:
                    ok, idx, msg = future.result()
                    if ok:
                        if msg == "exists":
                            exists += 1
                        else:
                            success += 1
                    else:
                        failed += 1
                        if failed <= 5:
                            tqdm.write(f"  Failed {idx}: {msg}")
                except Exception as e:
                    failed += 1
                    if failed <= 5:
                        tqdm.write(f"  Exception: {str(e)[:100]}")
                
                pbar.set_postfix({'ok': success, 'skip': exists, 'fail': failed})
        
        print(f"\n✓ {source} complete!")
        print(f"  New: {success}, Existing: {exists}, Failed: {failed}")
        print(f"  Raw images: {output_dir / 'raw'}")
        print(f"  GT files: {output_dir / 'gt'}")
        
        # Create manifest
        manifest = []
        for img in sorted((output_dir / 'raw').glob('*.png')):
            sample_id = img.stem
            gt_path = output_dir / 'gt' / f'{sample_id}.csv'
            if gt_path.exists():
                manifest.append({
                    'sample_id': sample_id,
                    'source': source,
                    'style': '0001',
                    'raw_path': str(img),
                    'gt_path': str(gt_path)
                })
        
        with open(output_dir / 'manifest.json', 'w') as f:
            json.dump(manifest, f, indent=2)
        
        print(f"  Manifest: {len(manifest)} samples")
        
        # Cleanup temp
        shutil.rmtree(temp_base, ignore_errors=True)
    
    print(f"\n{'='*60}")
    print("All sources complete!")
    print(f"{'='*60}")


if __name__ == '__main__':
    main()