import os
import torch
import random
import warnings
import numpy as np
from torch.utils.data import Dataset

'''
there are 2 versions of selection of segments that feature vectors are created from
- DeepDataset that takes first 500 valid segments for feature vector extraction
- DeepDatasetV2 stores all valid segments and randomly selects 500 of them for feature vector extraction

'''

class DeepDataset(Dataset):
    def __init__(self, data_path, patient_ids, fs, lead):
        self.data_path = data_path
        self.patient_ids = patient_ids
        self.id_mapped = {pid: idx for idx, pid in enumerate(patient_ids)}
        self.fs = fs
        self.lead = lead
        self.vectors = []
        self.labels = []
        self.m = 8
        self.segment_length = 10  # In seconds
        self.num_segments = int(15*3600 / self.segment_length)  # Total segments in 15 hours
        self.max_segments_per_patient = 500
        
        for patient_id in self.patient_ids:
            valid_segments_count = 0
            signal_path = os.path.join(self.data_path, f"{patient_id}_signal.npy")
            qrs_path = os.path.join(self.data_path, f"{patient_id}_qrs.npy")
            signal = np.load(signal_path)[self.lead, :15*3600*self.fs]
            qrs_indices = np.load(qrs_path)

            for segment_idx in range(self.num_segments):
                if valid_segments_count >= self.max_segments_per_patient:
                    break  # Stop if we have enough valid segments

                start = segment_idx * self.segment_length * self.fs
                end = start + self.segment_length * self.fs
                qrs_in_segment = qrs_indices[(qrs_indices >= start) & (qrs_indices < end)] - start

                if len(qrs_in_segment) > 0:  # Check for valid segments
                    segment = signal[start:end]
                    vector_v = self.extract_feature_vector(segment, qrs_in_segment)
                    self.vectors.append(vector_v)
                    self.labels.append(self.id_mapped[patient_id])
                    valid_segments_count += 1

            if valid_segments_count < self.max_segments_per_patient:
                warnings.warn(f"Patient ID {patient_id} has less than 500 valid segments: {valid_segments_count} segments found.")
        
        self.vectors = np.array(self.vectors)
        self.labels = np.array(self.labels)


    def __len__(self):
        return len(self.vectors)

    def __getitem__(self, index):
        vector_v = self.vectors[index]
        label = self.labels[index]
        return torch.tensor(vector_v, dtype=torch.float).unsqueeze(0), torch.tensor(label, dtype=torch.long)


    def extract_feature_vector(self, segment, qrs_indices):
        qrs_complexes = []
        half_window = int(0.125 * self.fs / 2)
        window_length = 2 * half_window

        for idx in qrs_indices:
            start_idx = max(0, idx - half_window)
            end_idx = min(len(segment), idx + half_window)
            qrs_complex = segment[start_idx:end_idx]
            if len(qrs_complex) < window_length:
                padding = window_length - len(qrs_complex)
                qrs_complex = np.pad(qrs_complex, (0, padding), 'constant')

            qrs_complexes.append(qrs_complex)

        average_qrs = np.mean(qrs_complexes, axis=0) if qrs_complexes else np.zeros(window_length)
        correlations = [np.correlate(qrs_complex, average_qrs, 'valid')[0] for qrs_complex in qrs_complexes]
        top_m_indices = np.argsort(correlations)[-self.m:]
        vector_v = np.concatenate([qrs_complexes[idx] for idx in top_m_indices], axis=0)

        while len(vector_v) < self.m * window_length:
            vector_v = np.concatenate([vector_v, qrs_complexes[top_m_indices[-1]]])
        return vector_v[:self.m * window_length] 



class DeepDatasetV2(Dataset):
    def __init__(self, data_path, patient_ids, fs, lead):
        self.data_path = data_path
        self.patient_ids = patient_ids
        self.id_mapped = {pid: idx for idx, pid in enumerate(patient_ids)}
        self.fs = fs
        self.lead = lead
        self.vectors = []
        self.labels = []
        self.m = 8
        self.segment_length = 10  # seconds
        self.num_segments = int(15*3600 / self.segment_length)  # total segments in 15 hours
        self.max_segments_per_patient = 500 
        
        for patient_id in self.patient_ids:
            valid_segment_indices = [] 

            signal_path = os.path.join(self.data_path, f"{patient_id}_signal.npy")
            qrs_path = os.path.join(self.data_path, f"{patient_id}_qrs.npy")
            signal = np.load(signal_path)[self.lead, :15*3600*self.fs]
            qrs_indices = np.load(qrs_path)

            for segment_idx in range(self.num_segments):
                start = segment_idx * self.segment_length * self.fs
                end = start + self.segment_length * self.fs
                qrs_in_segment = qrs_indices[(qrs_indices >= start) & (qrs_indices < end)] - start

                if len(qrs_in_segment) > 0:  # Check for valid segments
                    valid_segment_indices.append(segment_idx)

            selected_indices = random.sample(valid_segment_indices, min(len(valid_segment_indices), self.max_segments_per_patient))
            if len(valid_segment_indices) < self.max_segments_per_patient:
                warnings.warn(f"Patient ID {patient_id} has less than 500 valid segments: {len(valid_segment_indices)} segments found.")

            for segment_idx in selected_indices: #extract V vectors only from selected segments
                start = segment_idx * self.segment_length * self.fs
                end = start + self.segment_length * self.fs
                qrs_in_segment = qrs_indices[(qrs_indices >= start) & (qrs_indices < end)] - start
                segment = signal[start:end]
                vector_v = self.extract_feature_vector(segment, qrs_in_segment)
                self.vectors.append(vector_v)
                self.labels.append(self.id_mapped[patient_id])
        
        self.vectors = np.array(self.vectors)
        self.labels = np.array(self.labels)

    def __len__(self):
        return len(self.vectors)

    def __getitem__(self, index):
        vector_v = self.vectors[index]
        label = self.labels[index]
        return torch.tensor(vector_v, dtype=torch.float).unsqueeze(0), torch.tensor(label, dtype=torch.long)
    
    def extract_feature_vector(self, segment, qrs_indices):
        qrs_complexes = []
        half_window = int(0.125 * self.fs / 2)
        window_length = 2 * half_window

        for idx in qrs_indices:
            start_idx = max(0, idx - half_window)
            end_idx = min(len(segment), idx + half_window)
            qrs_complex = segment[start_idx:end_idx]
            if len(qrs_complex) < window_length:
                padding = window_length - len(qrs_complex)
                qrs_complex = np.pad(qrs_complex, (0, padding), 'constant')

            qrs_complexes.append(qrs_complex)

        average_qrs = np.mean(qrs_complexes, axis=0) if qrs_complexes else np.zeros(window_length)
        correlations = [np.correlate(qrs_complex, average_qrs, 'valid')[0] for qrs_complex in qrs_complexes]
        top_m_indices = np.argsort(correlations)[-self.m:]
        vector_v = np.concatenate([qrs_complexes[idx] for idx in top_m_indices], axis=0)

        while len(vector_v) < self.m * window_length:
            vector_v = np.concatenate([vector_v, qrs_complexes[top_m_indices[-1]]])
        return vector_v[:self.m * window_length]