import os

import torch

import random

from torch import nn

from torch.utils.data import DataLoader, Subset

from sklearn.model_selection import train_test_split

import neural

from class_ecgdataset import ECGDataset

from class_deepdataset import DeepDatasetV2

from train_and_eval import train_model, evaluate_model

def run(NUM_SEGMENTS, NUM_SECONDS, NUM_BATCH, LEADS, NUM_EPOCHS, DATA_PATH, FS, NUM_PATIENTS=None):

    '''

    Function that:

    - loads ECG data, prepares the dataset, performs a stratified train-test split

    - trains a DeepECG model and evaluates its performance

    It operates on a specified number of subjects (if provided), segments, and leads.

    Returns:

    - accuracy (float): the accuracy of the model on the test set

    - total_train_time (int): the total training time in seconds.

    - total_eval_time (int): the total evaluation time in seconds.

    '''

    torch.manual_seed(44)

    random.seed(44)

    device = "cuda" if torch.cuda.is_available() else "cpu" 

    print(f'Used device: {device}')

    print('Loading signals...')

    patient_ids_path = os.path.join(DATA_PATH, 'patient_ids.txt')

    with open(patient_ids_path, 'r') as f:

        PATIENT_IDS = [line.strip() for line in f.readlines()]

    if NUM_PATIENTS is not None:

        PATIENT_IDS = PATIENT_IDS[:NUM_PATIENTS]

    full_dataset = ECGDataset(DATA_PATH, PATIENT_IDS, FS, NUM_SEGMENTS, NUM_SECONDS, LEADS)

    # stratified split to maintain equal distribution of each patient's data

    train_indices, test_indices = train_test_split(range(len(full_dataset)), test_size=0.5, stratify=full_dataset.labels)

    train_dataset = Subset(full_dataset, train_indices)

    test_dataset = Subset(full_dataset, test_indices)

    train_loader = DataLoader(train_dataset, batch_size=NUM_BATCH, shuffle=True)

    test_loader = DataLoader(test_dataset, batch_size=NUM_BATCH, shuffle=False)

    single_batch_segments, _ = next(iter(train_loader))

    output_shape = len(PATIENT_IDS) 

    hidden_units = 32 

    dummy_network = neural.DeepECG_DUMMY(len(LEADS), hidden_units, output_shape).to(device)

    single_batch_segments = single_batch_segments.to(device)

    final_features = dummy_network(single_batch_segments)

    model = neural.DeepECG(len(LEADS), hidden_units, output_shape, final_features).to(device)

    loss_fn = nn.CrossEntropyLoss()

    optimizer = torch.optim.SGD(params=model.parameters(), lr=0.01)

    total_train_time = train_model(model=model,

                                   data_loader=train_loader, 

                                   loss_fn=loss_fn,

                                   optimizer=optimizer,

                                   device=device, 

                                   num_epochs=NUM_EPOCHS,

                                   output_shape = output_shape

                                   )

    accuracy, total_eval_time = evaluate_model(model=model, 

                                               test_loader=test_loader,

                                               loss_fn=loss_fn,

                                               device=device,

                                               output_shape = output_shape

                                               )

    return accuracy, total_train_time, total_eval_time

def run_deepecg(NUM_BATCH, LEAD, NUM_EPOCHS, DATA_PATH, FS):

    '''

    Function that:

    - loads ECG data, prepares the dataset using DeepECG methodology, performs a stratified train-test split

    - trains a DeepECG model and evaluates its performance

    It operates on a specified lead of ECG recordings.

    Returns:

    - accuracy (float): the accuracy of the model on the test set

    - total_train_time (int): the total training time in seconds.

    - total_eval_time (int): the total evaluation time in seconds.

    '''

    torch.manual_seed(44)

    random.seed(44)

    device = "cuda" if torch.cuda.is_available() else "cpu" 

    print(f'Used device: {device}')

    print('Loading signals...')

    patient_ids_path = os.path.join(DATA_PATH, 'patient_ids.txt')

    with open(patient_ids_path, 'r') as f:

        PATIENT_IDS = [line.strip() for line in f.readlines()]

    full_dataset = DeepDatasetV2(DATA_PATH, PATIENT_IDS, FS, LEAD)

    # stratified split to maintain equal distribution of each patient's data

    train_indices, test_indices = train_test_split(range(len(full_dataset)), test_size=0.5, stratify=full_dataset.labels)

    train_dataset = Subset(full_dataset, train_indices)

    test_dataset = Subset(full_dataset, test_indices)

    train_loader = DataLoader(train_dataset, batch_size=NUM_BATCH, shuffle=True)

    test_loader = DataLoader(test_dataset, batch_size=NUM_BATCH, shuffle=False)

    single_batch_segments, _ = next(iter(train_loader))

    num_leads = 1

    output_shape = len(PATIENT_IDS) 

    hidden_units = 32 

    dummy_network = neural.DeepECG_DUMMY(num_leads, hidden_units, output_shape).to(device)

    single_batch_segments = single_batch_segments.to(device)

    final_features = dummy_network(single_batch_segments)

    model = neural.DeepECG(num_leads, hidden_units, output_shape, final_features).to(device)

    loss_fn = nn.CrossEntropyLoss()

    optimizer = torch.optim.SGD(params=model.parameters(), lr=0.01)

    total_train_time = train_model(model=model,

                                   data_loader=train_loader, 

                                   loss_fn=loss_fn,

                                   optimizer=optimizer,

                                   device=device, 

                                   num_epochs=NUM_EPOCHS,

                                   output_shape = output_shape

                                   )

    accuracy, total_eval_time = evaluate_model(model=model, 

                                               test_loader=test_loader,

                                               loss_fn=loss_fn,

                                               device=device,

                                               output_shape = output_shape

                                               )

    return accuracy, total_train_time, total_eval_time