import wfdb
from wfdb import processing
import numpy as np
import joblib
import pywt
import os
import cv2
from pdf2image import convert_from_path
import warnings
import pickle
from scipy import signal as sg
warnings.filterwarnings('ignore')


def extract_hrv_features(rr_intervals):
    """
    Extract heart rate variability features from RR intervals.
    
    Args:
        rr_intervals (numpy.ndarray): RR intervals in seconds
        
    Returns:
        list: Four HRV features [sdnn, rmssd, pnn50, tri_index]
    """
    if len(rr_intervals) < 2:
        return [0, 0, 0, 0]
    
    sdnn = np.std(rr_intervals)
    diff_rr = np.diff(rr_intervals)
    rmssd = np.sqrt(np.mean(diff_rr**2)) if len(diff_rr) > 0 else 0
    pnn50 = 100 * np.sum(np.abs(diff_rr) > 0.05) / len(diff_rr) if len(diff_rr) > 0 else 0
    
    if len(rr_intervals) > 2:
        bin_width = 1/128
        bins = np.arange(min(rr_intervals), max(rr_intervals) + bin_width, bin_width)
        n, _ = np.histogram(rr_intervals, bins=bins)
        tri_index = len(rr_intervals) / np.max(n) if np.max(n) > 0 else 0
    else:
        tri_index = 0
    
    return [sdnn, rmssd, pnn50, tri_index]


def extract_qrs_features(signal, r_peaks, fs):
    """
    Extract QRS complex features from ECG signal and detected R peaks.
    
    Args:
        signal (numpy.ndarray): ECG signal
        r_peaks (numpy.ndarray): Array of R peak indices
        fs (int): Sampling frequency in Hz
        
    Returns:
        list: Three QRS features [qrs_width_mean, qrs_width_std, qrs_amplitude_mean]
    """
    if len(r_peaks) < 2:
        return [0, 0, 0]
    
    qrs_width = []
    for i in range(len(r_peaks)):
        r_pos = r_peaks[i]
        window_before = max(0, r_pos - int(0.1 * fs))
        window_after = min(len(signal) - 1, r_pos + int(0.1 * fs))
        
        if r_pos > window_before:
            q_pos = window_before + np.argmin(signal[window_before:r_pos])
        else:
            q_pos = window_before
            
        if r_pos < window_after:
            s_pos = r_pos + np.argmin(signal[r_pos:window_after])
        else:
            s_pos = r_pos
        
        if s_pos > q_pos:
            qrs_width.append((s_pos - q_pos) / fs)
    
    qrs_width_mean = np.mean(qrs_width) if qrs_width else 0
    qrs_width_std = np.std(qrs_width) if qrs_width else 0
    qrs_amplitude_mean = np.mean([signal[r] for r in r_peaks]) if r_peaks.size > 0 else 0
    
    return [qrs_width_mean, qrs_width_std, qrs_amplitude_mean]


def digitize_ecg_from_pdf(pdf_path, output_file=None):
    """
    Process an ECG PDF file and convert it to a .dat signal file.
    
    Args:
        pdf_path (str): Path to the ECG PDF file
        output_file (str, optional): Path to save the output .dat file
    
    Returns:
        tuple: (path to the created .dat file, list of paths to segment files)
    """
    if output_file is None:
        output_file = 'calibrated_ecg.dat'
    
    images = convert_from_path(pdf_path)
    temp_image_path = 'temp_ecg_image.jpg'
    images[0].save(temp_image_path, 'JPEG')
    
    img = cv2.imread(temp_image_path, cv2.IMREAD_GRAYSCALE)
    height, width = img.shape
    
    calibration = {
        'seconds_per_pixel': 2.0 / 197.0,
        'mv_per_pixel': 1.0 / 78.8,
    }
    
    layer1_start = int(height * 35.35 / 100)
    layer1_end = int(height * 51.76 / 100)
    layer2_start = int(height * 51.82 / 100)
    layer2_end = int(height * 69.41 / 100)
    layer3_start = int(height * 69.47 / 100)
    layer3_end = int(height * 87.06 / 100)
    
    layers = [
        img[layer1_start:layer1_end, :],
        img[layer2_start:layer2_end, :],
        img[layer3_start:layer3_end, :]
    ]
    
    signals = []
    time_points = []
    layer_duration = 10.0
    
    for i, layer in enumerate(layers):
        _, binary = cv2.threshold(layer, 200, 255, cv2.THRESH_BINARY_INV)
        
        contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        waveform_contour = max(contours, key=cv2.contourArea)
        
        sorted_contour = sorted(waveform_contour, key=lambda p: p[0][0])
        x_coords = np.array([point[0][0] for point in sorted_contour])
        y_coords = np.array([point[0][1] for point in sorted_contour])
        
        isoelectric_line_y = layer.shape[0] * 0.6
        
        x_min, x_max = np.min(x_coords), np.max(x_coords)
        time = (x_coords - x_min) / (x_max - x_min) * layer_duration
        
        signal_mv = (isoelectric_line_y - y_coords) * calibration['mv_per_pixel']
        signal_mv = signal_mv - np.mean(signal_mv)
        
        time_points.append(time)
        signals.append(signal_mv)
    
    total_duration = layer_duration * len(layers)
    sampling_frequency = 500
    num_samples = int(total_duration * sampling_frequency)
    combined_time = np.linspace(0, total_duration, num_samples)
    combined_signal = np.zeros(num_samples)
    
    for i, (time, signal) in enumerate(zip(time_points, signals)):
        start_time = i * layer_duration
        mask = (combined_time >= start_time) & (combined_time < start_time + layer_duration)
        relevant_times = combined_time[mask]
        interpolated_signal = np.interp(relevant_times, start_time + time, signal)
        combined_signal[mask] = interpolated_signal
    
    combined_signal = combined_signal - np.mean(combined_signal)
    signal_peak = np.max(np.abs(combined_signal))
    target_amplitude = 2.0
    
    if signal_peak > 0 and (signal_peak < 0.5 or signal_peak > 4.0):
        scaling_factor = target_amplitude / signal_peak
        combined_signal = combined_signal * scaling_factor
    
    adc_gain = 1000.0
    int_signal = (combined_signal * adc_gain).astype(np.int16)
    int_signal.tofile(output_file)
    
    if os.path.exists(temp_image_path):
        os.remove(temp_image_path)
    
    segment_files = []
    samples_per_segment = int(layer_duration * sampling_frequency)
    
    base_name = os.path.splitext(output_file)[0]
    for i in range(3):
        start_idx = i * samples_per_segment
        end_idx = (i + 1) * samples_per_segment
        segment = combined_signal[start_idx:end_idx]
        
        segment_file = f"{base_name}_segment{i+1}.dat"
        (segment * adc_gain).astype(np.int16).tofile(segment_file)
        segment_files.append(segment_file)
    
    return output_file, segment_files


def split_dat_into_segments(file_path, segment_duration=10.0):
    """
    Split a DAT file into equal segments.
    
    Args:
        file_path (str): Path to the DAT file (without extension)
        segment_duration (float): Duration of each segment in seconds
        
    Returns:
        list: Paths to the segment files
    """
    signal_all_leads, fs = load_dat_signal(file_path)
    
    if signal_all_leads.shape[1] == 1:
        lead_index = 0
    else:
        lead_priority = [1, 0]  # Try Lead II (index 1), then I (index 0)
        lead_index = next((i for i in lead_priority if i < signal_all_leads.shape[1]), 0)
        
    signal = signal_all_leads[:, lead_index]
    
    samples_per_segment = int(segment_duration * fs)
    total_samples = len(signal)
    num_segments = total_samples // samples_per_segment
    
    segment_files = []
    
    base_name = os.path.splitext(file_path)[0]
    
    for i in range(num_segments):
        start_idx = i * samples_per_segment
        end_idx = (i + 1) * samples_per_segment
        segment = signal[start_idx:end_idx]
        
        segment_file = f"{base_name}_segment{i+1}.dat"
        segment.reshape(-1, 1).tofile(segment_file)
        segment_files.append(segment_file)
            
    return segment_files


def load_dat_signal(file_path, n_leads=12, n_samples=5000, dtype=np.int16):
    """
    Load a DAT file containing ECG signal data.
    
    Args:
        file_path (str): Path to the DAT file (without extension)
        n_leads (int): Number of leads in the signal
        n_samples (int): Number of samples per lead
        dtype: Data type of the signal
        
    Returns:
        tuple: (numpy array of signal data, sampling frequency)
    """
    if file_path.endswith('.dat'):
        dat_path = file_path
    else:
        dat_path = file_path + '.dat'
        
    raw = np.fromfile(dat_path, dtype=dtype)
    
    if raw.size != n_leads * n_samples:
        if raw.size == n_samples:
            signal = raw.reshape(n_samples, 1)
            return signal, 500
            
        possible_leads = [1, 2, 3, 6, 12]
        for possible_lead_count in possible_leads:
            if raw.size % possible_lead_count == 0:
                actual_samples = raw.size // possible_lead_count
                signal = raw.reshape(actual_samples, possible_lead_count)
                return signal, 500
        
        signal = raw.reshape(-1, 1)
        return signal, 500
        
    signal = raw.reshape(n_samples, n_leads)
    return signal, 500


def extract_features_from_signal(signal):
    """
    Extract features from an ECG signal.
    
    Args:
        signal (numpy.ndarray): ECG signal
        
    Returns:
        list: Basic features extracted from the signal (32 features)
    """
    features = []
    features.append(np.mean(signal))
    features.append(np.std(signal))
    features.append(np.median(signal))
    features.append(np.min(signal))
    features.append(np.max(signal))
    features.append(np.percentile(signal, 25))
    features.append(np.percentile(signal, 75))
    features.append(np.mean(np.diff(signal)))

    coeffs = pywt.wavedec(signal, 'db4', level=5)
    for coeff in coeffs:
        features.append(np.mean(coeff))
        features.append(np.std(coeff))
        features.append(np.min(coeff))
        features.append(np.max(coeff))

    return features


def classify_new_ecg(file_path, model):
    """
    Classify a new ECG file.
    
    Args:
        file_path (str): Path to the ECG file (without extension)
        model: The trained model for classification
        
    Returns:
        str: Classification result ("Normal", "Abnormal", or error message)
    """
    signal_all_leads, fs = load_dat_signal(file_path)
    
    if signal_all_leads.shape[1] == 1:
        lead_index = 0
    else:
        lead_priority = [1, 0]  # Try Lead II (index 1), then I (index 0)
        lead_index = next((i for i in lead_priority if i < signal_all_leads.shape[1]), 0)

    signal = signal_all_leads[:, lead_index]
    signal = (signal - np.mean(signal)) / np.std(signal)
    
    try:
        r_peaks = processing.gqrs_detect(sig=signal, fs=fs)
    except:
        r_peaks = np.array([])

    if len(r_peaks) < 2:
        basic_features = extract_features_from_signal(signal)
        record_features = basic_features + [0] * (45 - len(basic_features))
    else:
        rr_intervals = np.diff(r_peaks) / fs
        qrs_durations = np.array([r_peaks[i] - r_peaks[i - 1] for i in range(1, len(r_peaks))])

        record_features = []
        basic_features = extract_features_from_signal(signal)
        record_features.extend(basic_features)
        
        record_features.extend([
            len(r_peaks),
            np.mean(rr_intervals) if len(rr_intervals) > 0 else 0,
            np.std(rr_intervals) if len(rr_intervals) > 0 else 0,
            np.median(rr_intervals) if len(rr_intervals) > 0 else 0,
            np.mean(qrs_durations) / fs if len(qrs_durations) > 0 else 0,
            np.std(qrs_durations) / fs if len(qrs_durations) > 0 else 0
        ])
        
        hrv_features = extract_hrv_features(rr_intervals)
        record_features.extend(hrv_features)
        
        qrs_features = extract_qrs_features(signal, r_peaks, fs)
        record_features.extend(qrs_features)
        
        if len(rr_intervals) >= 4:
            try:
                rr_times = np.cumsum(rr_intervals)
                rr_times = np.insert(rr_times, 0, 0)
                
                fs_interp = 4.0
                t_interp = np.arange(0, rr_times[-1], 1/fs_interp)
                rr_interp = np.interp(t_interp, rr_times[:-1], rr_intervals)
                
                freq, psd = sg.welch(rr_interp, fs=fs_interp, nperseg=min(256, len(rr_interp)))
                
                vlf_mask = (freq >= 0.0033) & (freq < 0.04)
                lf_mask = (freq >= 0.04) & (freq < 0.15)
                hf_mask = (freq >= 0.15) & (freq < 0.4)
                
                lf_power = np.trapz(psd[lf_mask], freq[lf_mask]) if np.any(lf_mask) else 0
                hf_power = np.trapz(psd[hf_mask], freq[hf_mask]) if np.any(hf_mask) else 0
                
                lf_hf_ratio = lf_power / hf_power if hf_power > 0 else 0
                normalized_lf = lf_power / (lf_power + hf_power) if (lf_power + hf_power) > 0 else 0
            except:
                lf_power = hf_power = lf_hf_ratio = normalized_lf = 0
        else:
            lf_power = hf_power = lf_hf_ratio = normalized_lf = 0
            
        record_features.extend([lf_power, hf_power, lf_hf_ratio, normalized_lf])
    
    if len(record_features) < 45:
        record_features.extend([0] * (45 - len(record_features)))
    elif len(record_features) > 45:
        record_features = record_features[:45]
        
    prediction = model.predict([record_features])[0]
    result = "Abnormal" if prediction == 1 else "Normal"
    
    return result


def classify_ecg(file_path, model, is_pdf=False):
    """
    Wrapper function that handles both PDF and DAT ECG files with segment voting.
    
    Args:
        file_path (str): Path to the ECG file (.pdf or without extension for .dat)
        model: The trained model for classification
        is_pdf (bool): Whether the input file is a PDF (True) or DAT (False)
        
    Returns:
        str: Classification result ("Normal", "Abnormal", or error message)
    """
    try:
        if model is None:
            return "Error: Model not loaded. Please check model compatibility."
            
        if is_pdf:
            base_name = os.path.splitext(os.path.basename(file_path))[0]
            output_dat = f"{base_name}_digitized.dat"
            
            dat_path, segment_files = digitize_ecg_from_pdf(
                pdf_path=file_path, 
                output_file=output_dat
            )
        else:
            segment_files = split_dat_into_segments(file_path)
            
            if not segment_files:
                return classify_new_ecg(file_path, model)
        
        segment_results = []
        
        for segment_file in segment_files:
            segment_path = os.path.splitext(segment_file)[0]
            result = classify_new_ecg(segment_path, model)
            segment_results.append(result)
            
            try:
                os.remove(segment_file)
            except:
                pass
        
        if segment_results:
            normal_count = segment_results.count("Normal")
            abnormal_count = segment_results.count("Abnormal")
            
            if abnormal_count > normal_count:
                final_result = "Abnormal"
            elif normal_count > abnormal_count:
                final_result = "Normal"
            else:
                final_result = "Inconclusive"
                
            return final_result
        else:
            return "Error: No valid segments to classify"
        
    except Exception as e:
        error_msg = f"Classification error: {str(e)}"
        return error_msg