#!/usr/bin/env python3
"""
Generate synthetic 0001-type ECG images directly using matplotlib.
Much faster than ecg-image-kit, generates clean baseline images.

Usage:
    python generate_synthetic_simple.py --num_samples 20000 --workers 32
"""

import os
import sys
import argparse
import random
import numpy as np
import cv2
import wfdb
import matplotlib
matplotlib.use('Agg')  # Non-interactive backend
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
from pathlib import Path
from concurrent.futures import ProcessPoolExecutor, as_completed
from tqdm import tqdm
from scipy import signal as scipy_signal

# Paths
PTBXL_DIR = Path('/data/ecg-digitization/ptbxl/physionet.org/files/ptb-xl/1.0.3')
OUTPUT_DIR = Path('/data/ecg-digitization/synthetic_0001_v2')

# Target dimensions
TARGET_WIDTH = 4352
TARGET_HEIGHT = 1696
T0, T1 = 235, 4161
OUTPUT_GT_WIDTH = T1 - T0  # 3926

# ECG layout parameters (matching competition)
ZERO_MV = [703.5, 987.5, 1271.5, 1531.5]  # Zero baseline Y positions for each row
MV_TO_PIXEL = 78.5  # Pixels per mV


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


def load_ecg_signal(hea_path):
    """Load 12-lead ECG signal from PTB-XL record."""
    try:
        record_path = hea_path.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 to standard lead order
        lead_map = {name: signals[:, i] for i, name in enumerate(sig_names)}
        
        return lead_map, 500
    except Exception as e:
        return None, 0


def create_ecg_image(lead_map, sample_rate=500):
    """Create ECG image with 0001-style (clean red grid)."""
    # Create figure with correct aspect ratio
    fig_width = TARGET_WIDTH / 100  # inches at 100 dpi
    fig_height = TARGET_HEIGHT / 100
    
    fig, ax = plt.subplots(figsize=(fig_width, fig_height), dpi=100)
    
    # White background
    ax.set_facecolor('white')
    fig.patch.set_facecolor('white')
    
    # Draw red grid (0001 style)
    grid_color = '#ffcccc'  # Light red
    grid_color_major = '#ff9999'  # Darker red
    
    # Minor grid (1mm = ~20 pixels at this scale)
    minor_step = 20
    for x in range(0, TARGET_WIDTH, minor_step):
        ax.axvline(x, color=grid_color, linewidth=0.3, alpha=0.5)
    for y in range(0, TARGET_HEIGHT, minor_step):
        ax.axhline(y, color=grid_color, linewidth=0.3, alpha=0.5)
    
    # Major grid (5mm = ~100 pixels)
    major_step = 100
    for x in range(0, TARGET_WIDTH, major_step):
        ax.axvline(x, color=grid_color_major, linewidth=0.5, alpha=0.7)
    for y in range(0, TARGET_HEIGHT, major_step):
        ax.axhline(y, color=grid_color_major, linewidth=0.5, alpha=0.7)
    
    # Lead layout: 4 rows, each with 3 lead segments (2.5s each)
    # Row 0: I, aVR, V1, V4 (rhythm strip shows full lead)
    # 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'],
    ]
    
    # Draw ECG traces
    segment_samples = int(2.5 * sample_rate)  # 2.5s at 500Hz = 1250 samples
    segment_width = (T1 - T0) // 3  # Pixels per segment
    
    gt_signal = np.zeros((4, OUTPUT_GT_WIDTH), dtype=np.float32)
    
    # Rows 0-2: 3 lead segments each
    for row_idx, leads in enumerate(row_leads):
        zero_y = ZERO_MV[row_idx]
        
        for col_idx, lead_name in enumerate(leads[:3]):  # First 3 leads per row
            if lead_name not in lead_map:
                continue
            
            lead_data = lead_map[lead_name]
            
            # Get 2.5s segment
            start = col_idx * segment_samples
            end = min(start + segment_samples, len(lead_data))
            segment = lead_data[start:end]
            
            # Pixel range for this segment
            col_start = T0 + col_idx * segment_width
            col_end = T0 + (col_idx + 1) * segment_width if col_idx < 2 else T1
            
            # Create x coordinates
            x_coords = np.linspace(col_start, col_end, len(segment))
            
            # Convert mV to pixels (y increases downward)
            y_coords = zero_y - segment * MV_TO_PIXEL
            
            # Draw trace
            ax.plot(x_coords, y_coords, 'k-', linewidth=0.8)
            
            # Store GT (resampled to output width)
            gt_start = col_idx * segment_width
            gt_end = (col_idx + 1) * segment_width if col_idx < 2 else OUTPUT_GT_WIDTH
            x_old = np.linspace(0, 1, len(segment))
            x_new = np.linspace(0, 1, gt_end - gt_start)
            gt_signal[row_idx, gt_start:gt_end] = np.interp(x_new, x_old, segment)
    
    # Row 3: Full Lead II (10s rhythm strip)
    if 'II' in lead_map:
        lead_data = lead_map['II']
        zero_y = ZERO_MV[3]
        
        full_samples = min(len(lead_data), int(10 * sample_rate))
        segment = lead_data[:full_samples]
        
        x_coords = np.linspace(T0, T1, len(segment))
        y_coords = zero_y - segment * MV_TO_PIXEL
        
        ax.plot(x_coords, y_coords, 'k-', linewidth=0.8)
        
        # Store GT
        x_old = np.linspace(0, 1, len(segment))
        x_new = np.linspace(0, 1, OUTPUT_GT_WIDTH)
        gt_signal[3, :] = np.interp(x_new, x_old, segment)
    
    # Set limits
    ax.set_xlim(0, TARGET_WIDTH)
    ax.set_ylim(TARGET_HEIGHT, 0)  # Inverted Y
    ax.axis('off')
    
    # Render to array
    fig.tight_layout(pad=0)
    fig.canvas.draw()
    
    # Get image from figure
    w, h = fig.canvas.get_width_height()
    img = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8).reshape(h, w, 3)
    
    plt.close(fig)
    
    # Resize to exact target
    if img.shape[:2] != (TARGET_HEIGHT, TARGET_WIDTH):
        img = cv2.resize(img, (TARGET_WIDTH, TARGET_HEIGHT))
    
    return img, gt_signal


def generate_single(args):
    """Generate a single synthetic ECG sample."""
    idx, hea_path, output_dir = args
    
    sample_id = f'syn2_{idx:08d}'
    out_img = output_dir / 'stage1' / f'{sample_id}-0001.png'
    out_gt = output_dir / 'gt' / f'{sample_id}.npy'
    
    if out_img.exists() and out_gt.exists():
        return True, idx, "exists"
    
    try:
        # Load ECG
        lead_map, fs = load_ecg_signal(hea_path)
        if lead_map is None:
            return False, idx, "load failed"
        
        # Generate image
        img, gt_signal = create_ecg_image(lead_map, fs)
        
        # Save
        out_img.parent.mkdir(parents=True, exist_ok=True)
        out_gt.parent.mkdir(parents=True, exist_ok=True)
        
        cv2.imwrite(str(out_img), cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
        np.save(out_gt, gt_signal)
        
        return True, idx, "success"
        
    except Exception as e:
        return False, idx, str(e)[:100]


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--num_samples', type=int, default=20000)
    parser.add_argument('--output_dir', type=str, default=str(OUTPUT_DIR))
    parser.add_argument('--workers', type=int, default=16)
    args = parser.parse_args()
    
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    (output_dir / 'stage1').mkdir(exist_ok=True)
    (output_dir / 'gt').mkdir(exist_ok=True)
    
    print(f"Finding PTB-XL records...")
    records = find_ptbxl_records()
    print(f"Found {len(records)} records")
    
    if not records:
        print("No records found!")
        return
    
    # Sample/repeat records as needed
    if len(records) > args.num_samples:
        records = random.sample(records, args.num_samples)
    elif len(records) < args.num_samples:
        records = records * (args.num_samples // len(records) + 1)
        records = records[:args.num_samples]
    
    print(f"Generating {len(records)} 0001-type synthetic images...")
    
    tasks = [(i, rec, output_dir) for i, rec in enumerate(records)]
    
    success, failed = 0, 0
    
    with ProcessPoolExecutor(max_workers=args.workers) as executor:
        futures = {executor.submit(generate_single, t): t for t in tasks}
        
        for future in tqdm(as_completed(futures), total=len(futures), desc="Generating"):
            ok, idx, msg = future.result()
            if ok:
                success += 1
            else:
                failed += 1
                if failed <= 5:
                    print(f"\n  Failed {idx}: {msg}")
    
    print(f"\n✓ Done! Success: {success}, Failed: {failed}")
    
    # Create manifest
    import json
    manifest = []
    for img in sorted((output_dir / 'stage1').glob('*.png')):
        sample_id = img.stem.replace('-0001', '')
        gt_path = output_dir / 'gt' / f'{sample_id}.npy'
        if gt_path.exists():
            manifest.append({
                'sample_id': sample_id,
                'style': '0001',
                'path': str(img),
                'gt': str(gt_path)
            })
    
    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()