import torch 
import torch.nn as nn
import numpy as np
import copy
import model
import argparse
import os

parser = argparse.ArgumentParser(description='Test for argparse')
parser.add_argument('--epochs', help='Number of epochs', default=10)
args = parser.parse_args()


epoch = int(args.epochs)

def train_model(model, train_dataset, val_dataset, n_epochs):
  optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
  criterion = nn.L1Loss(reduction='sum')
  #criterion = torch.nn.MSELoss(reduction='mean',reduce=True, size_average=True)
  history = dict(train=[], val=[])

  best_model_wts = copy.deepcopy(model.state_dict())
  best_loss = 10000.0
  
  for epoch in range(1, n_epochs + 1):
    model = model.train()

    train_losses = []
    for seq_true in train_dataset:
      optimizer.zero_grad()

      seq_true = seq_true
      seq_pred = model(seq_true)

      loss = criterion(seq_pred, seq_true)

      loss.backward()
      optimizer.step()

      train_losses.append(loss.item())

    val_losses = []
    model = model.eval()
    with torch.no_grad():
      for seq_true in val_dataset:

        seq_true = seq_true
        seq_pred = model(seq_true)

        loss = criterion(seq_pred, seq_true)
        val_losses.append(loss.item())

    train_loss = np.mean(train_losses)
    val_loss = np.mean(val_losses)

    history['train'].append(train_loss)
    history['val'].append(val_loss)

    if val_loss < best_loss:
      best_loss = val_loss
      best_model_wts = copy.deepcopy(model.state_dict())
    

    print(f'Epoch {epoch}: train loss {train_loss} val loss {val_loss}')

  model.load_state_dict(best_model_wts)
  return model.eval(), history



if __name__ == '__main__':


    model, history = train_model(
      model.model, 
      model.train_dataset, 
      model.val_dataset, 
      n_epochs=epoch
    )


    if not os.path.isdir("./model"):
        os.mkdir("./model")

    MODEL_PATH = './model/model.pth'
    torch.save(model, MODEL_PATH)




