# NeMo's "core" package
import nemo

# NeMo's ASR collection - this collections contains complete ASR models and
# building blocks (modules) for ASR
import nemo.collections.asr as nemo_asr
from omegaconf import OmegaConf, open_dict
from pytorch_lightning.loggers import WandbLogger
import pytorch_lightning as pl
import torch
torch.set_float32_matmul_precision('high')

wandb_logger = WandbLogger(
    log_model="all",
    project="hoot",
    name="test_subword_lyrics_en_lr1e3",
    save_dir="/home/tony/Data/checkpoints",
)

# params = OmegaConf.load("./configs/config_bpe.yaml")
# params = OmegaConf.load("./configs/config_faster_conformer_bpe.yaml")
params = OmegaConf.load("./configs/config_faster_conformer_bpe_a100.yaml")

# params.model.tokenizer.dir = "./tokenizers/exp/tokenizer_spe_unigram_v50/"  # note this is a directory, not a path to a vocabulary file
# params.model.tokenizer.dir = "./tokenizers/exp/tokenizer_spe_unigram_v1050/"  # note this is a directory, not a path to a vocabulary file
params.model.tokenizer.dir = "./tokenizers/en/tokenizer_spe_bpe_v1024/"  # note this is a directory, not a path to a vocabulary file
params.model.tokenizer.type = "bpe"

# trainer = pl.Trainer(devices=1, accelerator="gpu", max_epochs=50, logger=wandb_logger)
trainer = pl.Trainer(**params.trainer)
trainer.logger = wandb_logger

train_manifest = "/home/tony/Data/Hoot/en_train_manifest.json"
test_manifest = "/home/tony/Data/Hoot/en_test_manifest.json"

# Update paths to dataset
params.model.train_ds.manifest_filepath = train_manifest
params.model.validation_ds.manifest_filepath = test_manifest

# remove spec augment for this dataset
# params.model.spec_augment.rect_masks = 0

# first_asr_model = nemo_asr.models.EncDecCTCModelBPE(cfg=params.model, trainer=trainer)
first_asr_model = nemo_asr.models.EncDecCTCModelBPE(cfg=params.model, trainer=trainer)

# Start training!!!
trainer.fit(first_asr_model)
first_asr_model.save_to("first_model_lyrics_en_lr1e4.nemo")
# print(
#     first_asr_model.transcribe(
#         paths2audio_files=["/home/tony/Data/Hoot/exp/cmkTJWx-AGg.wav"], batch_size=4
#     )
# )
print("done!")