#!/usr/bin/env python3
"""
Advanced ECG Post-Processing Module

Research-backed signal processing techniques for ECG digitization:
1. Adaptive Binarization (variant-specific)
2. Bilateral Filtering + Savgol Hybrid
3. Per-Lead Adaptive Smoothing
4. Wavelet Coefficient Thresholding
5. Viterbi Path Tracing (Dynamic Programming)
6. QRS Peak Detection + Preservation
7. Baseline Drift Correction

References:
- PhysioNet 2024 Winner (Krones et al., 2024): Hough Transform + DP
- PMC 2024: Enhanced Wavelet-Based Medical Image Denoising
- Wu et al. 2022: Fully-automated paper ECG digitisation
"""

import numpy as np
import cv2
from scipy.signal import savgol_filter, find_peaks, medfilt
from scipy.ndimage import gaussian_filter
from typing import Dict, Optional, Tuple
import warnings
warnings.filterwarnings('ignore')

# Try to import pywt for wavelet denoising
try:
    import pywt
    HAS_PYWT = True
except ImportError:
    HAS_PYWT = False
    print("WARNING: pywt not available, wavelet denoising disabled")


# =============================================================================
# Constants
# =============================================================================
LEAD_LAYOUT = [
    ['I', 'aVR', 'V1', 'V4'],
    ['II', 'aVL', 'V2', 'V5'],
    ['III', 'aVF', 'V3', 'V6'],
]

# Lead-specific Savgol window sizes (empirically tuned)
# Lower values = less smoothing, preserve more detail
# All using window=7, which is the baseline anyway
LEAD_SMOOTHING_WINDOWS = {
    'I':      7,   # Clean in training data
    'II':     7,   # Second column segment
    'III':    7,   # Often small amplitude
    'aVR':    7,   # Keep same as others
    'aVL':    7,   # Keep same (9 was hurting it before)
    'aVF':    7,   # Moderate noise
    'V1':     7,   # Precordial leads
    'V2':     7,
    'V3':     7,
    'V4':     7,
    'V5':     7,
    'V6':     7,
    'II_rhythm': 7  # Same as column II
}

# Variant-specific adaptive thresholding block sizes
VARIANT_BLOCK_SIZES = {
    '0001': 11,   # Synthetic - fine details matter
    '0003': 21,   # Scanned color - larger local regions
    '0004': 21,   # Scanned B&W - similar to 0003
    '0005': 25,   # Mobile color - more variation
    '0006': 25,   # Mobile screen - uneven lighting
    '0009': 31,   # Stained/soaked - heavily degraded
    '0010': 31,   # Extensive damage
    '0011': 31,   # Moldy color
    '0012': 31,   # Moldy B&W
}

# Variant-specific wavelet threshold factors
VARIANT_WAVELET_FACTORS = {
    '0001': 1.2,  # Synthetic - less aggressive
    '0003': 1.5,
    '0004': 1.5,
    '0005': 1.6,
    '0006': 1.6,
    '0009': 2.0,  # Heavily degraded - more aggressive
    '0010': 2.0,
    '0011': 2.0,
    '0012': 2.0,
}


# =============================================================================
# Individual Processing Functions
# =============================================================================
def adaptive_binarize_and_filter(image_pred_pixels: np.ndarray, 
                                  variant_type: str = '0001') -> np.ndarray:
    """
    Convert continuous predictions to binary with variant-specific tuning.
    
    Args:
        image_pred_pixels: Float array [0, 1] from model
        variant_type: '0001', '0003', '0004', '0005', etc.
    
    Returns:
        Filtered binary image as float [0, 1]
    """
    # Convert to 8-bit for OpenCV
    img_8bit = (np.clip(image_pred_pixels, 0, 1) * 255).astype(np.uint8)
    
    block_size = VARIANT_BLOCK_SIZES.get(variant_type, 21)
    if block_size % 2 == 0:
        block_size += 1  # Must be odd
    
    # Adaptive Gaussian threshold (better for uneven lighting)
    binary = cv2.adaptiveThreshold(
        img_8bit,
        maxValue=255,
        adaptiveMethod=cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
        thresholdType=cv2.THRESH_BINARY,
        blockSize=block_size,
        C=2  # Constant subtracted
    )
    
    # Morphological denoising
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
    binary = cv2.morphologyEx(binary, cv2.MORPH_OPEN, kernel, iterations=1)
    binary = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel, iterations=1)
    
    return binary / 255.0


def hybrid_smoothing(signal_mv: np.ndarray, window: int = 7, 
                     sigma_spatial: float = 0.5) -> np.ndarray:
    """
    Two-stage smoothing: Savgol + bilateral-like filtering (edge-preserving).
    
    Inspired by: PMC 2024 "Enhanced Wavelet-Based Medical Image Denoising"
    
    Args:
        signal_mv: 1D signal in mV
        window: Savgol window size
        sigma_spatial: Gaussian blur sigma
    
    Returns:
        Smoothed signal
    """
    if len(signal_mv) < window:
        return signal_mv.copy()
    
    # Stage 1: Savgol for structural smoothing
    if window % 2 == 0:
        window += 1
    smooth1 = savgol_filter(signal_mv, window, polyorder=2)
    
    # Stage 2: Bilateral-like filtering (preserve edges)
    # Simple approximation: weighted Gaussian based on local variance
    edges = np.abs(np.diff(smooth1, prepend=smooth1[0]))
    edge_threshold = np.median(edges) + np.std(edges)
    
    # Gaussian blur with reduced intensity at edges
    smooth2 = gaussian_filter(smooth1, sigma=sigma_spatial, mode='nearest')
    
    # Blend: use original more near detected edges
    edge_mask = (edges < edge_threshold).astype(float)
    edge_mask = gaussian_filter(edge_mask, sigma=1.0, mode='nearest')  # Smooth mask
    
    result = edge_mask * smooth1 + (1 - edge_mask) * smooth2
    
    return result


def apply_lead_specific_smoothing(signal_mv: np.ndarray, 
                                   lead_name: str) -> np.ndarray:
    """
    Apply per-lead smoothing with research-backed parameters.
    
    Args:
        signal_mv: 1D signal in mV
        lead_name: Lead name (I, II, III, aVR, etc.)
    
    Returns:
        Smoothed signal
    """
    window = LEAD_SMOOTHING_WINDOWS.get(lead_name, 7)
    if window % 2 == 0:
        window += 1
    
    if len(signal_mv) >= window:
        return savgol_filter(signal_mv, window, polyorder=2)
    return signal_mv.copy()


def wavelet_denoise_signal(signal_mv: np.ndarray, 
                           wavelet: str = 'db4',
                           decomp_level: int = 3,
                           threshold_factor: float = 1.5) -> np.ndarray:
    """
    Wavelet-based denoising: DWT -> soft thresholding -> IDWT.
    
    Research: Wiley 2021, PMC 2024 show 15-30 dB SNR improvements
    on medical signals.
    
    Args:
        signal_mv: 1D signal in mV
        wavelet: 'db4' (Daubechies) recommended for ECG
        decomp_level: Decomposition depth (3-4 typical)
        threshold_factor: Aggressive (1.0) to conservative (2.0)
    
    Returns:
        Denoised signal
    """
    if not HAS_PYWT:
        return signal_mv.copy()
    
    if len(signal_mv) < 2 ** decomp_level:
        return signal_mv.copy()
    
    try:
        # Step 1: Discrete Wavelet Transform
        coeffs = pywt.wavedec(signal_mv, wavelet, level=decomp_level)
        # coeffs = [cA, cD_n, cD_n-1, ..., cD_1]
        
        # Step 2: Estimate noise standard deviation from finest detail coeffs
        detail_coeffs_finest = coeffs[-1]
        sigma = np.median(np.abs(detail_coeffs_finest)) / 0.6745
        
        # Step 3: Compute threshold (Donoho-Johnstone universal threshold)
        threshold = threshold_factor * sigma * np.sqrt(2 * np.log(len(signal_mv)))
        
        # Step 4: Soft thresholding on detail coefficients only
        # (keep approximation coefficients untouched)
        thresholded = [coeffs[0]]  # Keep approximation
        for coeff in coeffs[1:]:   # Process detail coefficients
            thresholded.append(pywt.threshold(coeff, threshold, mode='soft'))
        
        # Step 5: Inverse Wavelet Transform
        denoised = pywt.waverec(thresholded, wavelet)
        
        # Ensure same length as input
        if len(denoised) != len(signal_mv):
            denoised = denoised[:len(signal_mv)]
        
        return denoised
        
    except Exception as e:
        # Fallback if wavelet fails
        return signal_mv.copy()


def preserve_qrs_peaks(signal_mv: np.ndarray, fs: int = 500, 
                       prominence_factor: float = 0.5) -> np.ndarray:
    """
    Detect QRS complexes and preserve their sharpness.
    
    QRS are the most diagnostically important features—over-smoothing
    can reduce SNR despite improving visual smoothness.
    
    Args:
        signal_mv: 1D signal in mV
        fs: Sampling frequency in Hz
        prominence_factor: Peak prominence threshold
    
    Returns:
        Signal with preserved QRS peaks
    """
    if len(signal_mv) < 10:
        return signal_mv.copy()
    
    try:
        # Estimate local peaks (QRS have high amplitude)
        peaks, properties = find_peaks(
            np.abs(signal_mv),
            height=np.percentile(np.abs(signal_mv), 75),  # Top 25% peaks
            distance=int(0.15 * fs),  # Minimum 150ms apart (ECG physiology)
            prominence=np.std(signal_mv) * prominence_factor
        )
        
        if len(peaks) == 0:
            return signal_mv.copy()
        
        # For each peak, estimate original "sharpness"
        sharpened = signal_mv.copy()
        for peak in peaks:
            if 1 < peak < len(signal_mv) - 2:
                # Curvature check
                curvature = abs(signal_mv[peak+1] - 2*signal_mv[peak] + signal_mv[peak-1])
                if curvature > np.std(signal_mv) * 0.3:
                    # Slightly increase sharpness (carefully!)
                    sharpened[peak] *= 1.02  # Amplify by 2%
        
        return sharpened
        
    except Exception:
        return signal_mv.copy()


def remove_baseline_drift(signal_mv: np.ndarray, 
                          order: Optional[int] = None,
                          aggressive: bool = False) -> np.ndarray:
    """
    Remove slow baseline drift using high-pass filtering.
    
    NOTE: For ECG digitization where we're predicting pixel positions,
    the "baseline" IS the signal. Only use this for very slow drifts,
    not the DC component.
    
    Args:
        signal_mv: 1D signal in mV
        order: Filter order
        aggressive: If True, remove more drift (use carefully)
    
    Returns:
        Detrended signal (or original if not needed)
    """
    # DISABLED: Baseline drift removal removes the actual signal
    # in ECG digitization context where baseline IS the signal value
    return signal_mv.copy()


def viterbi_signal_extraction(cost_map: np.ndarray, 
                               smoothness_penalty: float = 0.5,
                               max_jump: int = 10) -> np.ndarray:
    """
    Extract continuous signal using Viterbi algorithm (Dynamic Programming).
    
    Ensures smooth, continuous path even with broken/noisy segments.
    Reference: PhysioNet Challenge 2024 Winner [Krones et al. 2024]
    
    Args:
        cost_map: [H, W] cost map (lower = more likely signal)
        smoothness_penalty: Penalty for vertical jumps
        max_jump: Maximum allowed jump between adjacent columns
    
    Returns:
        signal_path: Continuous signal path [W] in y-coordinates
    """
    H, W = cost_map.shape
    
    # Dynamic Programming forward pass
    DP = np.full((H, W), np.inf)
    backtrack = np.zeros((H, W), dtype=np.int32)
    
    # Initialize: first column costs
    DP[:, 0] = cost_map[:, 0]
    
    # Forward pass: for each column, find best y-position
    for x in range(1, W):
        for y in range(H):
            # Check previous column: which y-positions are reachable?
            y_min = max(0, y - max_jump)
            y_max = min(H, y + max_jump + 1)
            
            for prev_y in range(y_min, y_max):
                transition_cost = smoothness_penalty * abs(y - prev_y)
                total_cost = DP[prev_y, x-1] + transition_cost + cost_map[y, x]
                
                if total_cost < DP[y, x]:
                    DP[y, x] = total_cost
                    backtrack[y, x] = prev_y
    
    # Backtrack: find best path
    signal_path = np.zeros(W, dtype=np.float32)
    
    # Start from rightmost column: find best y
    signal_path[-1] = np.argmin(DP[:, -1])
    
    # Backtrack to start
    for x in range(W - 2, -1, -1):
        signal_path[x] = backtrack[int(signal_path[x+1]), x+1]
    
    return signal_path


def viterbi_refine_prediction(pred_y_normalized: np.ndarray,
                               row_height: int = 500,
                               smoothness: float = 1.0) -> np.ndarray:
    """
    Refine continuous prediction using Viterbi path finding.
    
    Args:
        pred_y_normalized: Model predictions [0, 1] for y-position per x
        row_height: Height of row crop in pixels
        smoothness: Smoothness penalty (higher = smoother)
    
    Returns:
        Refined y-positions [0, 1]
    """
    W = len(pred_y_normalized)
    H = row_height
    
    # Create cost map from predictions
    # Lower cost where prediction says signal should be
    cost_map = np.zeros((H, W), dtype=np.float32)
    
    for x in range(W):
        pred_y = pred_y_normalized[x] * H
        # Create Gaussian cost centered at prediction
        for y in range(H):
            dist = abs(y - pred_y)
            cost_map[y, x] = dist  # Linear cost
    
    # Run Viterbi
    signal_path = viterbi_signal_extraction(cost_map, 
                                            smoothness_penalty=smoothness,
                                            max_jump=15)
    
    # Normalize back to [0, 1]
    return signal_path / H


# =============================================================================
# Main Processor Class
# =============================================================================
class AdvancedECGProcessor:
    """
    Complete advanced ECG post-processing pipeline.
    
    Combines multiple research-backed techniques for maximum SNR improvement.
    """
    
    def __init__(self, variant_type: str = '0001',
                 enable_wavelet: bool = True,
                 enable_viterbi: bool = False,
                 enable_qrs_preservation: bool = True,
                 enable_baseline_correction: bool = True,
                 enable_bilateral: bool = True):
        """
        Initialize processor with variant-specific parameters.
        
        Args:
            variant_type: ECG image variant ('0001', '0003', etc.)
            enable_wavelet: Enable wavelet denoising
            enable_viterbi: Enable Viterbi path refinement
            enable_qrs_preservation: Enable QRS peak preservation
            enable_baseline_correction: Enable baseline drift removal
            enable_bilateral: Enable bilateral-like filtering
        """
        self.variant = variant_type
        self.enable_wavelet = enable_wavelet and HAS_PYWT
        self.enable_viterbi = enable_viterbi
        self.enable_qrs = enable_qrs_preservation
        self.enable_baseline = enable_baseline_correction
        self.enable_bilateral = enable_bilateral
        
        self._setup_parameters()
    
    def _setup_parameters(self):
        """Set variant-specific parameters."""
        self.binarize_block_size = VARIANT_BLOCK_SIZES.get(self.variant, 21)
        self.wavelet_threshold = VARIANT_WAVELET_FACTORS.get(self.variant, 1.5)
        
        # More aggressive smoothing for degraded variants
        if self.variant in ['0009', '0010', '0011', '0012']:
            self.base_smoothing_window = 9
            self.wavelet_level = 4
        else:
            self.base_smoothing_window = 7
            self.wavelet_level = 3
    
    def process_signal(self, signal_mv: np.ndarray, 
                       lead_name: str = 'I') -> np.ndarray:
        """
        Complete processing pipeline for a single lead signal.
        
        Args:
            signal_mv: 1D signal in mV
            lead_name: Lead name for lead-specific parameters
        
        Returns:
            Processed signal in mV
        """
        if len(signal_mv) < 10:
            return signal_mv.copy()
        
        result = signal_mv.copy()
        
        # 1. Per-lead adaptive smoothing (most beneficial)
        result = apply_lead_specific_smoothing(result, lead_name)
        
        # 2. Wavelet denoising - DISABLED for now, testing impact
        # Doesn't seem to help much and may add artifacts
        # if self.enable_wavelet:
        #     result = wavelet_denoise_signal(result,
        #                                     wavelet='db4',
        #                                     decomp_level=self.wavelet_level,
        #                                     threshold_factor=self.wavelet_threshold)
        
        # Final clamp
        return np.clip(result, -10, 10)
    
    def process_row_predictions(self, pred_mv: np.ndarray,
                                 row_idx: int) -> np.ndarray:
        """
        Process all leads in a row.
        
        Args:
            pred_mv: Full row predictions in mV [W]
            row_idx: Row index (0, 1, 2, 3)
        
        Returns:
            Processed row predictions
        """
        if row_idx < 3:
            # Standard rows with 4 segments
            lead_names = LEAD_LAYOUT[row_idx]
            segment_width = len(pred_mv) // 4
            result = pred_mv.copy()
            
            for seg_idx, lead_name in enumerate(lead_names):
                start = seg_idx * segment_width
                end = (seg_idx + 1) * segment_width
                
                segment = pred_mv[start:end]
                processed = self.process_signal(segment, lead_name)
                result[start:end] = processed
            
            return result
        else:
            # Rhythm strip (row 3) - full Lead II
            return self.process_signal(pred_mv, 'II_rhythm')
    
    def refine_with_viterbi(self, pred_y_normalized: np.ndarray,
                            row_height: int = 500) -> np.ndarray:
        """
        Optionally refine predictions using Viterbi path finding.
        
        Args:
            pred_y_normalized: Model predictions [0, 1]
            row_height: Height of row crop
        
        Returns:
            Refined predictions [0, 1]
        """
        if not self.enable_viterbi:
            return pred_y_normalized
        
        return viterbi_refine_prediction(pred_y_normalized,
                                         row_height=row_height,
                                         smoothness=1.0)


# =============================================================================
# Utility Functions
# =============================================================================
def detect_variant_from_filename(filename: str) -> str:
    """
    Detect variant type from filename.
    
    Args:
        filename: Image filename (e.g., '1006427285-0003.png')
    
    Returns:
        Variant string ('0001', '0003', etc.)
    """
    import re
    match = re.search(r'-(\d{4})\.', filename)
    if match:
        return match.group(1)
    return '0001'  # Default


def create_processor_for_variant(variant: str, 
                                  aggressive: bool = False) -> AdvancedECGProcessor:
    """
    Factory function to create processor with variant-appropriate settings.
    
    Args:
        variant: Variant type string
        aggressive: Enable more aggressive processing for degraded images
    
    Returns:
        Configured AdvancedECGProcessor
    """
    # Degraded variants need more aggressive processing
    is_degraded = variant in ['0009', '0010', '0011', '0012']
    
    return AdvancedECGProcessor(
        variant_type=variant,
        enable_wavelet=True,
        enable_viterbi=aggressive and is_degraded,
        enable_qrs_preservation=True,
        enable_baseline_correction=True,
        enable_bilateral=True
    )


# =============================================================================
# Test Functions
# =============================================================================
def test_processing_chain():
    """Quick test of processing chain."""
    # Create synthetic ECG-like signal
    t = np.linspace(0, 2, 1000)
    signal = np.sin(2 * np.pi * 1.5 * t) + 0.3 * np.sin(2 * np.pi * 10 * t)
    signal += 0.1 * np.random.randn(len(signal))  # Add noise
    
    processor = AdvancedECGProcessor(variant_type='0003')
    processed = processor.process_signal(signal, lead_name='I')
    
    print(f"Input std: {np.std(signal):.4f}")
    print(f"Output std: {np.std(processed):.4f}")
    print(f"Processing complete!")
    
    return signal, processed


if __name__ == '__main__':
    test_processing_chain()