#!/usr/bin/env python3
"""
Generate 0001-style synthetic ECG images directly (no ecg-image-kit).

Uses matplotlib to render ECG signals on a red grid background,
matching the 0001 competition format.

This is simpler and more reliable than ecg-image-kit.

Output:
- synthetic_0001_{source}/raw/*.png  (1700x2200)
- synthetic_0001_{source}/gt/*.csv   (3926 lines x 4 columns)

Usage:
    python generate_synthetic_0001_direct.py --source georgia --workers 8
    python generate_synthetic_0001_direct.py --source all --workers 8
"""

import os
import sys
import argparse
import json
from pathlib import Path
from concurrent.futures import ProcessPoolExecutor, as_completed
from tqdm import tqdm
import numpy as np
import cv2
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas

try:
    import wfdb
    from scipy import signal as scipy_signal
    from scipy.ndimage import gaussian_filter1d
except ImportError:
    import subprocess
    subprocess.run([sys.executable, '-m', 'pip', 'install', 'wfdb', 'scipy'], check=True)
    import wfdb
    from scipy import signal as scipy_signal
    from scipy.ndimage import gaussian_filter1d

# =============================================================================
# Configuration
# =============================================================================
DATA_ROOT = Path('/data/ecg-digitization')

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',
}

# Raw image dimensions (match competition format)
RAW_WIDTH = 2200
RAW_HEIGHT = 1700

# After Stage 1 processing dimensions
TARGET_WIDTH = 4352
TARGET_HEIGHT = 1696

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

# Baseline y-positions (in 1696-height image)
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
# Row 0: I,    II,   III  (+ lead II full in column 4)
# Row 1: aVR,  aVL,  aVF
# Row 2: V1,   V2,   V3
# Row 3: V4,   V5,   V6
LEAD_LAYOUT = [
    ['I', 'aVR', 'V1', 'V4'],      # Column leads at row indices
    ['II', 'aVL', 'V2', 'V5'],
    ['III', 'aVF', 'V3', 'V6'],
]

# 0001 style: Red grid on white/cream background
GRID_COLOR_MAJOR = '#FFB6B6'  # Light red for major grid
GRID_COLOR_MINOR = '#FFD0D0'  # Lighter red for minor grid
TRACE_COLOR = '#000000'  # Black trace
BG_COLOR = '#FFFFFF'  # White background


# =============================================================================
# Record Discovery
# =============================================================================
def find_records(source, data_dir):
    """Find all ECG record files for a given source."""
    records = []
    data_dir = Path(data_dir)
    
    if source == 'ptbxl':
        records500 = data_dir / 'records500'
        if records500.exists():
            for hea in records500.rglob('*_hr.hea'):
                if hea.with_suffix('.dat').exists():
                    records.append(str(hea))
    elif source == 'chapman':
        for hea in data_dir.rglob('JS*.hea'):
            if hea.with_suffix('.mat').exists():
                records.append(str(hea))
    elif source == 'georgia':
        for hea in data_dir.rglob('E*.hea'):
            if hea.with_suffix('.mat').exists():
                records.append(str(hea))
    elif source == 'cpsc':
        for hea in data_dir.rglob('A*.hea'):
            if hea.with_suffix('.mat').exists():
                records.append(str(hea))
    elif source == 'ningbo':
        for hea in data_dir.rglob('*.hea'):
            if hea.with_suffix('.mat').exists():
                records.append(str(hea))
    
    return sorted(records)


# =============================================================================
# Signal Loading and Processing
# =============================================================================
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 - normalize variations
        lead_signals = {}
        for i, name in enumerate(sig_names):
            # Handle various naming conventions
            name_norm = name.upper()
            if name_norm == 'AVR':
                name_norm = 'aVR'
            elif name_norm == 'AVL':
                name_norm = 'aVL'
            elif name_norm == 'AVF':
                name_norm = 'aVF'
            lead_signals[name_norm] = signals[:, i]
        
        return lead_signals, fs
    except Exception as e:
        return None, None


def get_lead_signal(lead_signals, lead_name):
    """Get signal for a lead, handling name variations."""
    variants = [lead_name, lead_name.upper(), lead_name.lower()]
    if lead_name.startswith('aV'):
        variants.extend([lead_name.upper(), 'A' + lead_name[1:].upper()])
    
    for v in variants:
        if v in lead_signals:
            return lead_signals[v]
    return None


# =============================================================================
# Image Generation
# =============================================================================
def render_ecg_image(lead_signals, fs=500):
    """
    Render ECG signals to a 0001-style image.
    
    Returns:
        image: numpy array (RAW_HEIGHT, RAW_WIDTH, 3)
        gt_data: numpy array (OUTPUT_GT_WIDTH, 4) - y-pixel coords
    """
    # Create figure with exact pixel dimensions
    dpi = 100
    fig_width = RAW_WIDTH / dpi
    fig_height = RAW_HEIGHT / dpi
    
    fig, ax = plt.subplots(1, 1, figsize=(fig_width, fig_height), dpi=dpi)
    fig.patch.set_facecolor(BG_COLOR)
    ax.set_facecolor(BG_COLOR)
    
    # Set axis limits
    ax.set_xlim(0, RAW_WIDTH)
    ax.set_ylim(RAW_HEIGHT, 0)  # Flip y-axis (0 at top)
    ax.set_aspect('equal')
    ax.axis('off')
    
    # Draw grid
    # Major grid every 5mm (200 pixels at this scale)
    major_spacing = 40  # Approximate for 5mm
    minor_spacing = 8   # 1mm
    
    for x in range(0, RAW_WIDTH + 1, minor_spacing):
        lw = 0.5 if x % major_spacing == 0 else 0.2
        color = GRID_COLOR_MAJOR if x % major_spacing == 0 else GRID_COLOR_MINOR
        ax.axvline(x, color=color, linewidth=lw)
    
    for y in range(0, RAW_HEIGHT + 1, minor_spacing):
        lw = 0.5 if y % major_spacing == 0 else 0.2
        color = GRID_COLOR_MAJOR if y % major_spacing == 0 else GRID_COLOR_MINOR
        ax.axhline(y, color=color, linewidth=lw)
    
    # Scale factors for raw image
    # In raw image (1700x2200), we need to map signals appropriately
    raw_scale_y = RAW_HEIGHT / TARGET_HEIGHT  # ~1.002
    raw_scale_x = RAW_WIDTH / TARGET_WIDTH * 2  # Need to fit into half width
    
    # Baseline positions in raw image coordinates
    raw_baselines = ZERO_MV * raw_scale_y
    
    # Time parameters: 3 columns, 2.5s each = 7.5s visible
    samples_per_column = int(2.5 * fs)  # 1250 samples
    column_width = RAW_WIDTH / 4  # 4 columns (3 short + 1 long lead II)
    
    # Pixel per mV in raw coordinates
    raw_mv_to_pixel = MV_TO_PIXEL * raw_scale_y
    
    # GT data array (will be computed for TARGET dimensions)
    gt_data = np.zeros((OUTPUT_GT_WIDTH, 4), dtype=np.float32)
    pixels_per_column = OUTPUT_GT_WIDTH // 3
    
    # Draw each lead
    for col_idx in range(3):
        for row_idx in range(4):
            lead_name = LEAD_LAYOUT[col_idx][row_idx]
            signal = get_lead_signal(lead_signals, lead_name)
            
            if signal is None:
                # Fill GT with baseline
                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] = ZERO_MV[row_idx]
                continue
            
            # Extract segment (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] = ZERO_MV[row_idx]
                continue
            
            # Smooth signal slightly to reduce noise
            if len(segment) > 10:
                segment = gaussian_filter1d(segment, sigma=1)
            
            # X coordinates in raw image
            x_start = col_idx * column_width + 50  # Small margin
            x_end = (col_idx + 1) * column_width - 10
            x = np.linspace(x_start, x_end, len(segment))
            
            # Y coordinates: baseline - signal (inverted because y=0 at top)
            y = raw_baselines[row_idx] - segment * raw_mv_to_pixel
            y = np.clip(y, 50, RAW_HEIGHT - 50)
            
            # Draw trace
            ax.plot(x, y, color=TRACE_COLOR, linewidth=0.8)
            
            # Compute GT for this segment (in TARGET 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_gt_pixels = col_end - col_start
            
            # Resample signal to GT resolution
            x_old = np.linspace(0, 1, len(segment))
            x_new = np.linspace(0, 1, num_gt_pixels)
            signal_resampled = np.interp(x_new, x_old, segment)
            
            # Convert to pixel coordinates (TARGET dimensions)
            y_pixels = ZERO_MV[row_idx] - signal_resampled * MV_TO_PIXEL
            y_pixels = np.clip(y_pixels, 0, TARGET_HEIGHT - 1)
            gt_data[col_start:col_end, row_idx] = y_pixels
    
    # Draw Lead II full (column 4, row 1 - taking full 10s)
    lead_ii = get_lead_signal(lead_signals, 'II')
    if lead_ii is not None:
        # Use full 10s (or available)
        segment = lead_ii[:int(10 * fs)]
        if len(segment) > 10:
            segment = gaussian_filter1d(segment, sigma=1)
        
        x_start = 3 * column_width + 20
        x_end = RAW_WIDTH - 20
        x = np.linspace(x_start, x_end, len(segment))
        y = raw_baselines[1] - segment * raw_mv_to_pixel  # Row 1 for Lead II full
        y = np.clip(y, 50, RAW_HEIGHT - 50)
        ax.plot(x, y, color=TRACE_COLOR, linewidth=0.8)
    
    # Convert figure to image
    canvas = FigureCanvas(fig)
    canvas.draw()
    
    # Get the image as numpy array
    buf = canvas.buffer_rgba()
    image = np.asarray(buf)
    image = cv2.cvtColor(image, cv2.COLOR_RGBA2BGR)
    
    plt.close(fig)
    
    # Ensure exact dimensions
    if image.shape[:2] != (RAW_HEIGHT, RAW_WIDTH):
        image = cv2.resize(image, (RAW_WIDTH, RAW_HEIGHT))
    
    return image, gt_data


def generate_single_sample(args):
    """Generate a single synthetic sample."""
    idx, hea_path, source, output_dir = 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"
    
    try:
        # Load signals
        lead_signals, fs = load_ecg_signals(hea_path)
        if lead_signals is None:
            return False, idx, "failed to load signals"
        
        # Generate image and GT
        image, gt_data = render_ecg_image(lead_signals, fs)
        
        # Save image
        raw_dir.mkdir(parents=True, exist_ok=True)
        cv2.imwrite(str(out_img), image)
        
        # Save GT as CSV
        gt_dir.mkdir(parents=True, exist_ok=True)
        np.savetxt(out_gt, gt_data, delimiter=',', fmt='%.2f')
        
        return True, idx, "success"
        
    except Exception as e:
        return False, idx, f"error: {str(e)[:80]}"


# =============================================================================
# Main
# =============================================================================
def main():
    parser = argparse.ArgumentParser(description='Generate 0001-style synthetic ECGs')
    parser.add_argument('--source', type=str, required=True,
                       choices=['ptbxl', 'chapman', 'georgia', 'cpsc', 'ningbo', 'all'])
    parser.add_argument('--workers', type=int, default=8)
    parser.add_argument('--limit', type=int, default=None)
    parser.add_argument('--output_base', type=str, default='/data/ecg-digitization')
    args = parser.parse_args()
    
    output_base = Path(args.output_base)
    
    if args.source == 'all':
        sources = ['chapman', 'georgia', 'cpsc', 'ningbo']
    else:
        sources = [args.source]
    
    for source in sources:
        print(f"\n{'='*60}")
        print(f"Processing: {source}")
        print(f"{'='*60}")
        
        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_dir = output_base / f'synthetic_0001_{source}'
        output_dir.mkdir(parents=True, exist_ok=True)
        
        print(f"Finding records...")
        records = find_records(source, data_dir)
        print(f"Found {len(records)} records")
        
        if len(records) == 0:
            continue
        
        if args.limit:
            records = records[:args.limit]
            print(f"Limited to {len(records)} records")
        
        tasks = [(i, hea, source, output_dir) for i, hea in enumerate(records)]
        
        success = exists = failed = 0
        
        with ProcessPoolExecutor(max_workers=args.workers) as executor:
            futures = {executor.submit(generate_single_sample, t): t for t in tasks}
            
            for future in tqdm(as_completed(futures), total=len(futures), desc=source):
                try:
                    ok, idx, msg = future.result()
                    if ok:
                        if msg == "exists":
                            exists += 1
                        else:
                            success += 1
                    else:
                        failed += 1
                        if failed <= 3:
                            tqdm.write(f"  Failed {idx}: {msg}")
                except Exception as e:
                    failed += 1
        
        print(f"\n✓ {source}: New={success}, Existing={exists}, Failed={failed}")
        
        # Create manifest
        manifest = []
        for img in sorted((output_dir / 'raw').glob('*.png')):
            gt = output_dir / 'gt' / img.name.replace('.png', '.csv')
            if gt.exists():
                manifest.append({
                    'sample_id': img.stem,
                    'source': source,
                    'raw': str(img),
                    'gt': str(gt)
                })
        
        with open(output_dir / 'manifest.json', 'w') as f:
            json.dump(manifest, f, indent=2)
        
        print(f"  Manifest: {len(manifest)} samples")


if __name__ == '__main__':
    main()