# Copyright (c) 2023 Amphion.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import collections
import json
import os
import sys
import time

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel
from torch.utils.data import ConcatDataset, DataLoader
from torch.utils.tensorboard import SummaryWriter

from models.base.base_sampler import BatchSampler
from utils.util import (
    Logger,
    remove_older_ckpt,
    save_config,
    set_all_random_seed,
    ValueWindow,
)


class BaseTrainer(object):
    def __init__(self, args, cfg):
        self.args = args
        self.log_dir = args.log_dir
        self.cfg = cfg

        self.checkpoint_dir = os.path.join(args.log_dir, "checkpoints")
        os.makedirs(self.checkpoint_dir, exist_ok=True)
        if not cfg.train.ddp or args.local_rank == 0:
            self.sw = SummaryWriter(os.path.join(args.log_dir, "events"))
            self.logger = self.build_logger()
        self.time_window = ValueWindow(50)

        self.step = 0
        self.epoch = -1
        self.max_epochs = self.cfg.train.epochs
        self.max_steps = self.cfg.train.max_steps

        # set random seed & init distributed training
        set_all_random_seed(self.cfg.train.random_seed)
        if cfg.train.ddp:
            dist.init_process_group(backend="nccl")

        if cfg.model_type not in ["AutoencoderKL", "AudioLDM"]:
            self.singers = self.build_singers_lut()

        # setup data_loader
        self.data_loader = self.build_data_loader()

        # setup model & enable distributed training
        self.model = self.build_model()
        print(self.model)

        if isinstance(self.model, dict):
            for key, value in self.model.items():
                value.cuda(self.args.local_rank)
                if key == "PQMF":
                    continue
                if cfg.train.ddp:
                    self.model[key] = DistributedDataParallel(
                        value, device_ids=[self.args.local_rank]
                    )
        else:
            self.model.cuda(self.args.local_rank)
            if cfg.train.ddp:
                self.model = DistributedDataParallel(
                    self.model, device_ids=[self.args.local_rank]
                )

        # create criterion
        self.criterion = self.build_criterion()
        if isinstance(self.criterion, dict):
            for key, value in self.criterion.items():
                self.criterion[key].cuda(args.local_rank)
        else:
            self.criterion.cuda(self.args.local_rank)

        # optimizer
        self.optimizer = self.build_optimizer()
        self.scheduler = self.build_scheduler()

        # save config file
        self.config_save_path = os.path.join(self.checkpoint_dir, "args.json")

    def build_logger(self):
        log_file = os.path.join(self.checkpoint_dir, "train.log")
        logger = Logger(log_file, level=self.args.log_level).logger

        return logger

    def build_dataset(self):
        raise NotImplementedError

    def build_data_loader(self):
        Dataset, Collator = self.build_dataset()
        # build dataset instance for each dataset and combine them by ConcatDataset
        datasets_list = []
        for dataset in self.cfg.dataset:
            subdataset = Dataset(self.cfg, dataset, is_valid=False)
            datasets_list.append(subdataset)
        train_dataset = ConcatDataset(datasets_list)

        train_collate = Collator(self.cfg)
        # TODO: multi-GPU training
        if self.cfg.train.ddp:
            raise NotImplementedError("DDP is not supported yet.")

        # sampler will provide indices to batch_sampler, which will perform batching and yield batch indices
        batch_sampler = BatchSampler(
            cfg=self.cfg, concat_dataset=train_dataset, dataset_list=datasets_list
        )

        # use batch_sampler argument instead of (sampler, shuffle, drop_last, batch_size)
        train_loader = DataLoader(
            train_dataset,
            collate_fn=train_collate,
            num_workers=self.args.num_workers,
            batch_sampler=batch_sampler,
            pin_memory=False,
        )
        if not self.cfg.train.ddp or self.args.local_rank == 0:
            datasets_list = []
            for dataset in self.cfg.dataset:
                subdataset = Dataset(self.cfg, dataset, is_valid=True)
                datasets_list.append(subdataset)
            valid_dataset = ConcatDataset(datasets_list)
            valid_collate = Collator(self.cfg)
            batch_sampler = BatchSampler(
                cfg=self.cfg, concat_dataset=valid_dataset, dataset_list=datasets_list
            )
            valid_loader = DataLoader(
                valid_dataset,
                collate_fn=valid_collate,
                num_workers=1,
                batch_sampler=batch_sampler,
            )
        else:
            raise NotImplementedError("DDP is not supported yet.")
            # valid_loader = None
        data_loader = {"train": train_loader, "valid": valid_loader}
        return data_loader

    def build_singers_lut(self):
        # combine singers
        if not os.path.exists(os.path.join(self.log_dir, self.cfg.preprocess.spk2id)):
            singers = collections.OrderedDict()
        else:
            with open(
                os.path.join(self.log_dir, self.cfg.preprocess.spk2id), "r"
            ) as singer_file:
                singers = json.load(singer_file)
        singer_count = len(singers)
        for dataset in self.cfg.dataset:
            singer_lut_path = os.path.join(
                self.cfg.preprocess.processed_dir, dataset, self.cfg.preprocess.spk2id
            )
            with open(singer_lut_path, "r") as singer_lut_path:
                singer_lut = json.load(singer_lut_path)
            for singer in singer_lut.keys():
                if singer not in singers:
                    singers[singer] = singer_count
                    singer_count += 1
        with open(
            os.path.join(self.log_dir, self.cfg.preprocess.spk2id), "w"
        ) as singer_file:
            json.dump(singers, singer_file, indent=4, ensure_ascii=False)
        print(
            "singers have been dumped to {}".format(
                os.path.join(self.log_dir, self.cfg.preprocess.spk2id)
            )
        )
        return singers

    def build_model(self):
        raise NotImplementedError()

    def build_optimizer(self):
        raise NotImplementedError

    def build_scheduler(self):
        raise NotImplementedError()

    def build_criterion(self):
        raise NotImplementedError

    def get_state_dict(self):
        raise NotImplementedError

    def save_config_file(self):
        save_config(self.config_save_path, self.cfg)

    # TODO, save without module.
    def save_checkpoint(self, state_dict, saved_model_path):
        torch.save(state_dict, saved_model_path)

    def load_checkpoint(self):
        checkpoint_path = os.path.join(self.checkpoint_dir, "checkpoint")
        assert os.path.exists(checkpoint_path)
        checkpoint_filename = open(checkpoint_path).readlines()[-1].strip()
        model_path = os.path.join(self.checkpoint_dir, checkpoint_filename)
        assert os.path.exists(model_path)
        if not self.cfg.train.ddp or self.args.local_rank == 0:
            self.logger.info(f"Re(store) from {model_path}")
        checkpoint = torch.load(model_path, map_location="cpu")
        return checkpoint

    def load_model(self, checkpoint):
        raise NotImplementedError

    def restore(self):
        checkpoint = self.load_checkpoint()
        self.load_model(checkpoint)

    def train_step(self, data):
        raise NotImplementedError(
            f"Need to implement function {sys._getframe().f_code.co_name} in "
            f"your sub-class of {self.__class__.__name__}. "
        )

    @torch.no_grad()
    def eval_step(self):
        raise NotImplementedError(
            f"Need to implement function {sys._getframe().f_code.co_name} in "
            f"your sub-class of {self.__class__.__name__}. "
        )

    def write_summary(self, losses, stats):
        raise NotImplementedError(
            f"Need to implement function {sys._getframe().f_code.co_name} in "
            f"your sub-class of {self.__class__.__name__}. "
        )

    def write_valid_summary(self, losses, stats):
        raise NotImplementedError(
            f"Need to implement function {sys._getframe().f_code.co_name} in "
            f"your sub-class of {self.__class__.__name__}. "
        )

    def echo_log(self, losses, mode="Training"):
        message = [
            "{} - Epoch {} Step {}: [{:.3f} s/step]".format(
                mode, self.epoch + 1, self.step, self.time_window.average
            )
        ]

        for key in sorted(losses.keys()):
            if isinstance(losses[key], dict):
                for k, v in losses[key].items():
                    message.append(
                        str(k).split("/")[-1] + "=" + str(round(float(v), 5))
                    )
            else:
                message.append(
                    str(key).split("/")[-1] + "=" + str(round(float(losses[key]), 5))
                )
        self.logger.info(", ".join(message))

    def eval_epoch(self):
        self.logger.info("Validation...")
        valid_losses = {}
        for i, batch_data in enumerate(self.data_loader["valid"]):
            for k, v in batch_data.items():
                if isinstance(v, torch.Tensor):
                    batch_data[k] = v.cuda()
            valid_loss, valid_stats, total_valid_loss = self.eval_step(batch_data, i)
            for key in valid_loss:
                if key not in valid_losses:
                    valid_losses[key] = 0
                valid_losses[key] += valid_loss[key]

        # Add mel and audio to the Tensorboard
        # Average loss
        for key in valid_losses:
            valid_losses[key] /= i + 1
        self.echo_log(valid_losses, "Valid")
        return valid_losses, valid_stats

    def train_epoch(self):
        for i, batch_data in enumerate(self.data_loader["train"]):
            start_time = time.time()
            # Put the data to cuda device
            for k, v in batch_data.items():
                if isinstance(v, torch.Tensor):
                    batch_data[k] = v.cuda(self.args.local_rank)

            # Training step
            train_losses, train_stats, total_loss = self.train_step(batch_data)
            self.time_window.append(time.time() - start_time)

            if self.args.local_rank == 0 or not self.cfg.train.ddp:
                if self.step % self.args.stdout_interval == 0:
                    self.echo_log(train_losses, "Training")

                if self.step % self.cfg.train.save_summary_steps == 0:
                    self.logger.info(f"Save summary as step {self.step}")
                    self.write_summary(train_losses, train_stats)

                if (
                    self.step % self.cfg.train.save_checkpoints_steps == 0
                    and self.step != 0
                ):
                    saved_model_name = "step-{:07d}_loss-{:.4f}.pt".format(
                        self.step, total_loss
                    )
                    saved_model_path = os.path.join(
                        self.checkpoint_dir, saved_model_name
                    )
                    saved_state_dict = self.get_state_dict()
                    self.save_checkpoint(saved_state_dict, saved_model_path)
                    self.save_config_file()
                    # keep max n models
                    remove_older_ckpt(
                        saved_model_name,
                        self.checkpoint_dir,
                        max_to_keep=self.cfg.train.keep_checkpoint_max,
                    )

                if self.step != 0 and self.step % self.cfg.train.valid_interval == 0:
                    if isinstance(self.model, dict):
                        for key in self.model.keys():
                            self.model[key].eval()
                    else:
                        self.model.eval()
                    # Evaluate one epoch and get average loss
                    valid_losses, valid_stats = self.eval_epoch()
                    if isinstance(self.model, dict):
                        for key in self.model.keys():
                            self.model[key].train()
                    else:
                        self.model.train()
                    # Write validation losses to summary.
                    self.write_valid_summary(valid_losses, valid_stats)
            self.step += 1

    def train(self):
        for epoch in range(max(0, self.epoch), self.max_epochs):
            self.train_epoch()
            self.epoch += 1
            if self.step > self.max_steps:
                self.logger.info("Training finished!")
                break