import torch

from src.utils import train_metric

def hybrid_fit(epochs, model, hybrid_train_loader, hybrid_val_loader, icdtype, opt_fn,loss_fn, learning_rate, device):

  optimizer = opt_fn(model.parameters(), lr=learning_rate)

  print('-'*10 + icdtype + '-'*10)

  for epoch in range(1,epochs+1):

    model.train()

    train_epoch_loss=0

    train_epoch_accuracy=0

    train_epoch_hammingloss=0

    train_epoch_f1score=0

    val_epoch_loss=0

    val_epoch_accuracy=0

    val_epoch_hammingloss=0

    val_epoch_f1score=0

    for rnn_x, cnn_x, y_dict in hybrid_train_loader:

      rnn_x = rnn_x.to(device)

      cnn_x = cnn_x.to(device)

      y = y_dict[icdtype]

      y = y.to(device)

      preds=model(rnn_x, cnn_x)

      optimizer.zero_grad()

      loss=loss_fn(preds,y)

      loss.backward()

      optimizer.step()

      accuracy, hammingloss, f1score  = train_metric(preds,y)

      train_epoch_loss+=loss.item()

      train_epoch_accuracy+=accuracy.item()

      train_epoch_hammingloss+=hammingloss

      train_epoch_f1score+=f1score

    model.eval()

    with torch.no_grad():

      for rnn_x, cnn_x, y_dict in hybrid_val_loader:

        rnn_x = rnn_x.to(device)

        cnn_x = cnn_x.to(device)

        y = y_dict[icdtype]

        y = y.to(device)

        preds=model(rnn_x, cnn_x)

        loss=loss_fn(preds,y)

        accuracy, hammingloss, f1score  = train_metric(preds,y)

        val_epoch_loss+=loss.item()

        val_epoch_accuracy+=accuracy.item()

        val_epoch_hammingloss+=hammingloss

        val_epoch_f1score+=f1score

    print("\n")

    print('-'*100)

    print('Epoch = {}/{}:\n train_loss = {:.4f}, train_accuracy = {:.4f}, train_hammingloss = {:.4f}, train_f1score = {:.4f}\n val_loss = {:.4f}, val_accuracy = {:.4f}, val_hammmingloss = {:.4f}, val_f1score = {:.4f}'.format(epoch

                                                              ,epochs

                                                              ,train_epoch_loss/len(hybrid_train_loader)

                                                              ,train_epoch_accuracy/len(hybrid_train_loader)

                                                              ,train_epoch_hammingloss/len(hybrid_train_loader)

                                                              ,train_epoch_f1score/len(hybrid_train_loader)

                                                              ,val_epoch_loss/len(hybrid_val_loader)

                                                              ,val_epoch_accuracy/len(hybrid_val_loader)

                                                              ,val_epoch_hammingloss/len(hybrid_val_loader)

                                                              ,val_epoch_f1score/len(hybrid_val_loader)

                                                              ))

    print('-'*100)

    print("\n")