import torch
from torch.utils.data import DataLoader

from deepSTRF.datasets import NS1
from deepSTRF.model_zoo import DNet
from deepSTRF.metrics import performance_metrics

dataset = NS1(path='../datasets/NS1/data', stimuli=('nat',), neuron_indexes='all')
model = DNet(n_bands=49, temporal_window_size=5, n_hidden=20, adapTrans=True)

for neuron_idx in dataset.neuron_indices:
    
    dataset.select_neuron(neuron_idx)

    len_train, len_val, len_test = 0.7*len(dataset), 0.1*len(dataset), 0.2*len(dataset)
    train_set, valid_set, test_set = torch.utils.data.random_split(dataset, [len_train, len_val, len_test])
    train_dataloader = DataLoader(train_set, shuffle=True)
    valid_dataloader = DataLoader(valid_set, shuffle=False)
    test_dataloader = DataLoader(test_set, shuffle=False)

    for epoch in range(n_epochs):

        for spectrogram, responses, ccmax in train_dataloader:

            prediction = model(spectrogram)

            CCraw, CCnorm = performance_metrics(prediction, responses, ccmax)
            
            loss = MSE(prediction, responses)
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()