#!/usr/bin/env python
import argparse
import oneflow as flow
from data import AudioDataLoader, AudioDataset
from decoder import Decoder
from encoder import Encoder
from transformer import Transformer
from solver import Solver
from utils import process_dict
from optimizer import TransformerOptimizer


parser = argparse.ArgumentParser(
    "End-to-End Automatic Speech Recognition Training " "(Transformer framework)."
)
# General config
# Task related
parser.add_argument(
    "--train-json",
    type=str,
    default="../../egs/aishell/dump/train/deltafalse/data.json",
    help="Filename of train label data (json)",
)
parser.add_argument(
    "--valid-json",
    type=str,
    default="../../egs/aishell/dump/dev/deltafalse/data.json",
    help="Filename of validation label data (json)",
)
parser.add_argument(
    "--dict",
    type=str,
    default="../../egs/aishell/data/lang_1char/train_chars.txt",
    help="Dictionary which should include <unk> <sos> <eos>",
)
# Low Frame Rate (stacking and skipping frames)
parser.add_argument(
    "--LFR_m", default=4, type=int, help="Low Frame Rate: number of frames to stack"
)
parser.add_argument(
    "--LFR_n", default=3, type=int, help="Low Frame Rate: number of frames to skip"
)
# Network architecture
# encoder
parser.add_argument(
    "--d_input", default=80, type=int, help="Dim of encoder input (before LFR)"
)
parser.add_argument(
    "--n_layers_enc", default=6, type=int, help="Number of encoder stacks"
)
parser.add_argument(
    "--n_head", default=8, type=int, help="Number of Multi Head Attention (MHA)"
)
parser.add_argument("--d_k", default=64, type=int, help="Dimension of key")
parser.add_argument("--d_v", default=64, type=int, help="Dimension of value")
parser.add_argument("--d_model", default=512, type=int, help="Dimension of model")
parser.add_argument("--d_inner", default=2048, type=int, help="Dimension of inner")
parser.add_argument("--dropout", default=0.1, type=float, help="Dropout rate")
parser.add_argument(
    "--pe_maxlen", default=5000, type=int, help="Positional Encoding max len"
)
# decoder
parser.add_argument(
    "--d_word_vec", default=512, type=int, help="Dim of decoder embedding"
)
parser.add_argument(
    "--n_layers_dec", default=6, type=int, help="Number of decoder stacks"
)
parser.add_argument(
    "--tgt_emb_prj_weight_sharing",
    default=1,
    type=int,
    help="share decoder embedding with decoder projection",
)
# Loss
parser.add_argument(
    "--label_smoothing", default=0.1, type=float, help="label smoothing"
)
# Training config
parser.add_argument("--epochs", default=150, type=int, help="Number of maximum epochs")
# minibatch
parser.add_argument(
    "--shuffle", default=1, type=int, help="reshuffle the data at every epoch"
)
parser.add_argument("--batch-size", default=16, type=int, help="Batch size")
parser.add_argument(
    "--batch_frames",
    default=15000,
    type=int,
    help="Batch frames. If this is not 0, batch size will make no sense",
)
parser.add_argument(
    "--maxlen-in",
    default=800,
    type=int,
    metavar="ML",
    help="Batch size is reduced if the input sequence length > ML",
)
parser.add_argument(
    "--maxlen-out",
    default=150,
    type=int,
    metavar="ML",
    help="Batch size is reduced if the output sequence length > ML",
)
parser.add_argument(
    "--num-workers", default=0, type=int, help="Number of workers to generate minibatch"
)
# optimizer
parser.add_argument(
    "--k", default=0.2, type=float, help="tunable scalar multiply to learning rate"
)
parser.add_argument("--warmup_steps", default=4000, type=int, help="warmup steps")
parser.add_argument("--step_num", default=0, type=int)
# save and load model
parser.add_argument(
    "--save-folder", default="exp/temp", help="Location to save epoch models"
)
parser.add_argument(
    "--checkpoint",
    dest="checkpoint",
    default=0,
    type=int,
    help="Enables checkpoint saving of model",
)
parser.add_argument(
    "--continue-from", default="", help="Continue from checkpoint model"
)
parser.add_argument(
    "--model-path",
    default="final.pth.tar",
    help="Location to save best validation model",
)
# logging
parser.add_argument(
    "--print-freq",
    default=10,
    type=int,
    help="Frequency of printing training infomation",
)


def main(args):
    # Construct Solver
    # data
    tr_dataset = AudioDataset(
        args.train_json,
        args.batch_size,
        args.maxlen_in,
        args.maxlen_out,
        batch_frames=args.batch_frames,
    )
    cv_dataset = AudioDataset(
        args.valid_json,
        args.batch_size,
        args.maxlen_in,
        args.maxlen_out,
        batch_frames=args.batch_frames,
    )
    tr_loader = AudioDataLoader(
        tr_dataset,
        batch_size=1,
        num_workers=args.num_workers,
        shuffle=args.shuffle,
        LFR_m=args.LFR_m,
        LFR_n=args.LFR_n,
    )
    cv_loader = AudioDataLoader(
        cv_dataset,
        batch_size=1,
        num_workers=args.num_workers,
        LFR_m=args.LFR_m,
        LFR_n=args.LFR_n,
    )

    # load dictionary and generate char_list, sos_id, eos_id
    char_list, sos_id, eos_id = process_dict(args.dict)
    vocab_size = len(char_list)
    data = {"tr_loader": tr_loader, "cv_loader": cv_loader}

    # model
    encoder = Encoder(
        args.d_input * args.LFR_m,
        args.n_layers_enc,
        args.n_head,
        args.d_k,
        args.d_v,
        args.d_model,
        args.d_inner,
        dropout=args.dropout,
        pe_maxlen=args.pe_maxlen,
    )
    decoder = Decoder(
        sos_id,
        eos_id,
        vocab_size,
        args.d_word_vec,
        args.n_layers_dec,
        args.n_head,
        args.d_k,
        args.d_v,
        args.d_model,
        args.d_inner,
        dropout=args.dropout,
        tgt_emb_prj_weight_sharing=args.tgt_emb_prj_weight_sharing,
        pe_maxlen=args.pe_maxlen,
    )
    model = Transformer(encoder, decoder)

    device = flow.device("cuda")
    model.to(device)

    # optimizer
    optimizier = TransformerOptimizer(
        flow.optim.Adam(model.parameters(), betas=(0.9, 0.98), eps=1e-09),
        args.k,
        args.d_model,
        args.warmup_steps,
        args.step_num,
    )

    # solver
    solver = Solver(data, model, optimizier, device, args)
    solver.train()


args = parser.parse_args()
print(args)
main(args)
