import torch
from opendc.utils.configuration import create_config, create_args
from opendc.utils.utils import load_checkpoint
from opendc.data_loader.flexible_loader import create_loader
from opendc.model.model import create_model
from opendc.model.criterion import create_criterion, create_scheduler
from opendc.train.optimizer import create_optimizer, create_fp16_scaler
from opendc.train.train import train

# configuration: get the arguments first and update the configurations for the system.
args = create_args()
cfg = create_config(args)

# dataset + loader
train_loader = create_loader(cfg, is_training=True)
val_loader = create_loader(cfg, is_training=False)	
  
# model - considering checkpoint
model = create_model(cfg)

# criterion, optimizer
optimizer = create_optimizer(model, cfg)
scheduler = create_scheduler(cfg)
criterion = create_criterion(cfg)

# for mixed precision training
fp16_scaler = create_fp16_scaler(cfg)

load_checkpoint(cfg, model=model, optimizer=optimizer, fp16_scaler=fp16_scaler)

# train - save the checkpoint in the train - test
train(model, train_loader, val_loader, optimizer, scheduler, fp16_scaler, criterion, cfg)

# post-training
# memory bank to investigate all data points.
