#!/usr/bin/env python3
"""
Convert NPY ground truth files to CSV format for training.

NPY files contain mV values with shape [4, N].
CSV files should contain pixel Y-coordinates with N rows and 4 columns.

Conversion: pixel_y = ZERO_MV[row] - mv_value * MV_TO_PIXEL
"""

import os
import sys
import numpy as np
import pandas as pd
from pathlib import Path
from tqdm import tqdm
import argparse

# Constants from training script
ZERO_MV = np.array([703.5, 987.5, 1271.5, 1531.5])
MV_TO_PIXEL = 78.5
TARGET_HEIGHT = 1696


def convert_mv_to_pixels(mv_signal):
    """Convert mV values to pixel Y-coordinates.
    
    Args:
        mv_signal: Array of shape [4, N] with mV values
        
    Returns:
        pixel_signal: Array of shape [4, N] with pixel Y-coordinates
    """
    pixel_signal = np.zeros_like(mv_signal)
    for row in range(4):
        pixel_signal[row] = ZERO_MV[row] - mv_signal[row] * MV_TO_PIXEL
    
    # Clip to valid range
    pixel_signal = np.clip(pixel_signal, 0, TARGET_HEIGHT - 1)
    return pixel_signal


def convert_npy_to_csv(npy_path, csv_path):
    """Convert a single NPY file to CSV format."""
    # Load NPY file [4, N] in mV
    mv_signal = np.load(npy_path)
    
    # Convert to pixel coordinates
    pixel_signal = convert_mv_to_pixels(mv_signal)
    
    # Transpose to [N, 4] for CSV format
    # When training loads: gt_df.values.T gives [4, N]
    df = pd.DataFrame(pixel_signal.T, columns=['row0', 'row1', 'row2', 'row3'])
    
    # Save CSV
    df.to_csv(csv_path, index=False)
    return True


def main():
    parser = argparse.ArgumentParser(description='Convert NPY ground truth to CSV format')
    parser.add_argument('--input_dir', type=str, 
                       default='/data/ecg-digitization/synthetic_ecgkit/gt',
                       help='Directory containing NPY files')
    parser.add_argument('--output_dir', type=str, default=None,
                       help='Output directory for CSV files (default: same as input)')
    parser.add_argument('--stage1_dir', type=str,
                       default='/data/ecg-digitization/synthetic_ecgkit/stage1',
                       help='Stage1 images directory (for creating per-variant CSVs)')
    parser.add_argument('--per_variant', action='store_true',
                       help='Create separate CSV for each image variant')
    args = parser.parse_args()
    
    input_dir = Path(args.input_dir)
    output_dir = Path(args.output_dir) if args.output_dir else input_dir
    stage1_dir = Path(args.stage1_dir)
    
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # Find all NPY files
    npy_files = sorted(input_dir.glob('*.npy'))
    print(f"Found {len(npy_files)} NPY files")
    
    if args.per_variant:
        # Create CSV for each image variant
        # Image names: syn_00000000-0001.png -> need syn_00000000-0001.csv
        # NPY names: syn_00000000.npy
        
        png_files = sorted(stage1_dir.glob('*.png'))
        print(f"Found {len(png_files)} PNG files")
        
        converted = 0
        for png_path in tqdm(png_files, desc="Converting"):
            # Extract base ID from image name (e.g., syn_00000000-0001 -> syn_00000000)
            img_stem = png_path.stem  # syn_00000000-0001
            base_id = img_stem.rsplit('-', 1)[0]  # syn_00000000
            
            npy_path = input_dir / f"{base_id}.npy"
            if not npy_path.exists():
                continue
            
            csv_path = output_dir / f"{img_stem}.csv"
            if convert_npy_to_csv(npy_path, csv_path):
                converted += 1
        
        print(f"Converted {converted} files")
    else:
        # Simple 1:1 conversion
        converted = 0
        for npy_path in tqdm(npy_files, desc="Converting"):
            csv_path = output_dir / f"{npy_path.stem}.csv"
            if convert_npy_to_csv(npy_path, csv_path):
                converted += 1
        
        print(f"Converted {converted} files")
    
    # Verify a sample
    if npy_files:
        sample_npy = npy_files[0]
        sample_csv = output_dir / f"{sample_npy.stem}.csv"
        
        print(f"\nSample verification:")
        print(f"  NPY: {sample_npy}")
        
        mv_data = np.load(sample_npy)
        print(f"  NPY shape: {mv_data.shape}")
        print(f"  NPY (mV) range: [{mv_data.min():.2f}, {mv_data.max():.2f}]")
        
        if sample_csv.exists():
            csv_data = pd.read_csv(sample_csv)
            print(f"  CSV shape: {csv_data.shape}")
            print(f"  CSV (pixels) range: [{csv_data.values.min():.1f}, {csv_data.values.max():.1f}]")
            print(f"  Expected pixel range: ~[{ZERO_MV.min() - 2*MV_TO_PIXEL:.0f}, {ZERO_MV.max() + 2*MV_TO_PIXEL:.0f}]")


if __name__ == '__main__':
    main()