# --------------------------------------------------------------------------------
# 		Import
# --------------------------------------------------------------------------------
import torch
import torch.nn as nn
import argparse
import json
import numpy as np
import os
from torchsummary import summary
from utils.logger import Logger

# Transformer model and trainer
from ocr_models.lstm_model import LSTMModel
from trainers.lstm_trainer import LSTMTrainer

from data_loaders.iam_dataloader import IAMDataLoader
from data_loaders.synthesis04_dataloader import JPDataLoader as JPDataloader_1
from data_loaders.katakana_dataloader import JPDataLoader as JPDataloader_2


# --------------------------------------------------------------------------------
#       Parser
# --------------------------------------------------------------------------------
parser = argparse.ArgumentParser(description='Training model')

parser.add_argument('--config', type=str, default=None,
                    help='Path to config file')

parser.add_argument('--resume_path', type=str, default=None,
                    help='Path to checkpoint file to resume')

args = parser.parse_args()


# --------------------------------------------------------------------------------
#       Configuration
# --------------------------------------------------------------------------------
with open("configs/config_lstm_synthesis04.json") as json_file:
    config = json.load(json_file)
    trainer_config = config["trainer"]
    model_config = config["model"]
    data_config = config["data"]

# Setup files for data loader
data_files_dir = data_config["path_preprocessing_files"]
dictionary = os.path.join(data_files_dir, data_config["dictionary"])

# Path to data for training
path_data_train = data_config["path_data_train"]
labels_file_train = os.path.join(path_data_train, data_config["labels_file_train"])

# Path to data for validating
path_data_valid = data_config["path_data_valid"]
labels_file_valid = os.path.join(path_data_valid, data_config["labels_file_valid"])

# Target vocab size, add <SOS> and <EOS>
with open(dictionary, "r", encoding="utf8") as f:
    trg_vocab = np.shape(f.readlines())[0] + 2

# Define model
input_dim = model_config["input_dim"]
enc_hidden_dim = model_config["enc_hidden_dim"]
enc_bidirectional = model_config["enc_bidirectional"]

dec_hidden_dim = model_config["dec_hidden_dim"]
dec_bidirectional = model_config["dec_bidirectional"]

num_layer = model_config["num_layer"]
max_len = model_config["max_len"]

model = LSTMModel(input_dim, enc_hidden_dim, enc_bidirectional,
                  dec_hidden_dim, dec_bidirectional, num_layer,
                  trg_vocab)

# Define optimizer
optimizer = torch.optim.Adam(model.parameters(),
                             lr=trainer_config["lr"],
                             betas=(0.9, 0.98), eps=1e-9)

# Logger
train_logger = Logger("logs/katakana/lstm/train/")
valid_logger = Logger("logs/katakana/lstm/valid/")

# Trainer
trainer = LSTMTrainer(model=model, optimizer=optimizer, data_loader=JPDataloader_1, config=config,
                      resume_path=args.resume_path, train_logger=train_logger, valid_logger=valid_logger,
                      path_data_train=path_data_train, labels_train=labels_file_train,
                      path_data_valid=path_data_valid, labels_valid=labels_file_valid, path_dictionary=dictionary,
                      max_len=max_len, trg_vocab=trg_vocab)

# --------------------------------------------------------------------------------
#       Main
# --------------------------------------------------------------------------------
trainer.train()
