import os

import torch

import numpy as np

from torch.utils.data import Dataset

#----------------------------------------------------

#----------storing-signals-implementation------------

#----------------------------------------------------

class ECGDataset(Dataset):

    def __init__(self, data_path, patient_ids, fs, n_windows, n_seconds, leads=None):

        self.data_path = data_path

        self.patient_ids = patient_ids

        self.fs = fs

        self.n_windows = n_windows

        self.n_seconds = n_seconds

        self.leads = leads

        self.segments = []

        self.labels = []

        self.id_mapped = {pid: idx for idx, pid in enumerate(patient_ids)}

        # Precompute and store segments

        for patient_id in self.patient_ids:

            signal = self.load_signal(patient_id)

            start_points = self.generate_starts(signal)

            for start_point in start_points:

                end_point = start_point + self.n_seconds * self.fs

                segment = signal[:, start_point:end_point]

                self.segments.append(segment)

                self.labels.append(self.id_mapped[patient_id])

        self.segments = np.array(self.segments)

        self.labels = np.array(self.labels)

    def load_signal(self, patient_id):

        signal_path = os.path.join(self.data_path, f"{patient_id}_signal.npy")

        signal = np.load(signal_path)

        signal = signal[:, :15*3600*self.fs]

        if self.leads is not None:

            signal = signal[self.leads, :]

        return signal

    def generate_starts(self, signal):

        max_start = signal.shape[1] - self.n_seconds * self.fs

        all_starts = np.arange(0, max_start, self.n_seconds * self.fs)

        chosen_starts = np.random.choice(all_starts, self.n_windows, replace=False)

        return chosen_starts

    def __len__(self):

        return len(self.segments)

    def __getitem__(self, index):

        segment = self.segments[index]

        label = self.labels[index]

        return torch.tensor(segment, dtype=torch.float), torch.tensor(label, dtype=torch.long)

#----------------------------------------------------

#-------------on-the-go-implementation---------------

#----------------------------------------------------

class ECGDataset_on_the_fly(Dataset):

    def __init__(self, data_path, patient_ids, fs, n_windows, n_seconds, leads=None):

        self.data_path = data_path

        self.patient_ids = patient_ids

        self.fs = fs

        self.n_windows = n_windows

        self.n_seconds = n_seconds

        self.leads = leads

        self.signal_cache = {}

        self.id_mapped = {pid: idx for idx, pid in enumerate(patient_ids)}

        self.segment_starts = {pid: self.generate_starts(pid) for pid in patient_ids}

    def load_signal(self, patient_id):

        if patient_id not in self.signal_cache:

            signal_path = os.path.join(self.data_path, f"{patient_id}_signal.npy")

            signal = np.load(signal_path)

            signal = signal[:, :15*3600*self.fs]

            if self.leads is not None:

                signal = signal[self.leads, :15*3600*self.fs]

            self.signal_cache[patient_id] = signal

        return self.signal_cache[patient_id]

    def generate_starts(self, patient_id):

        signal = self.load_signal(patient_id)

        max_start = signal.shape[1] - self.n_seconds * self.fs

        # make a grid of points that are n_seconds apart

        all_starts = np.arange(0, max_start, self.n_seconds * self.fs) 

        # and choose n_windows of them as our starting pts

        chosen_starts = np.random.choice(all_starts, self.n_windows, replace=False)

        return chosen_starts

    def __len__(self):

        return len(self.patient_ids) * self.n_windows

    def __getitem__(self, index):

        patient_index = index // self.n_windows

        patient_id = self.patient_ids[patient_index]

        window_index = index % self.n_windows

        start_point = self.segment_starts[patient_id][window_index]

        end_point = start_point + self.n_seconds * self.fs

        signal = self.load_signal(patient_id)

        segment = signal[:, start_point:end_point]

        label = self.id_mapped[patient_id]  

        return torch.tensor(segment, dtype=torch.float), torch.tensor(label, dtype=torch.long)