import argparse
from tqdm import tqdm
import torch
from torch.nn.parallel import DistributedDataParallel
from optimizer.optimizers import Eve, ScaledAdam
from schedulers.scheduler import NoamScheduler, Eden
from models.tts.valle.valle_dataset import VALLEDataset, VALLECollator
from models.tts.base import TTSTrainer
from models.tts.valle.valle import VALLE

class VALLETrainer(TTSTrainer):
    def __init__(self, args, cfg):
        TTSTrainer.__init__(self, args, cfg)

    def _build_model(self):
        
        model = VALLE(self.cfg.model)
        
        return model
    
    def _build_dataset(self):
        return VALLEDataset, VALLECollator

    def _build_optimizer(self):
        if self.args.train_stage:
            if isinstance(self.model, DistributedDataParallel):
                model = self.model.module
            else:
                model = self.model
            model_parameters = model.stage_parameters(self.args.train_stage)
        else:
            model_parameters = self.model.parameters()   

                                        
        if self.cfg.train.optimizer == "ScaledAdam":
            parameters_names = []
            if self.args.train_stage != 0:
                parameters_names.append(
                    [
                        name_param_pair[0]
                        for name_param_pair in model.stage_named_parameters(
                            self.args.train_stage
                        )
                    ]
                )
            else:
                parameters_names.append(
                    [
                        name_param_pair[0]
                        for name_param_pair in model.named_parameters()
                    ]
                )

            optimizer = ScaledAdam(
                model_parameters,
                lr=self.cfg.train.base_lr,
                betas=(0.9, 0.95),
                clipping_scale=2.0,
                parameters_names=parameters_names,
                show_dominant_parameters=False,
                clipping_update_period=1000,
            )
        elif self.cfg.train.optimizer == "Eve":
            optimizer = Eve(
                model_parameters,
                lr=self.cfg.train.base_lr,
                betas=(0.9, 0.98),
                target_rms=0.1,
            )
        elif self.cfg.train.optimizer == "AdamW":
            optimizer = torch.optim.AdamW(
                model_parameters,
                lr=self.cfg.train.base_lr,
                betas=(0.9, 0.95),
                weight_decay=1e-2,
                eps=1e-8,
            )
        elif self.cfg.train.optimizer == "Adam":
            optimizer = torch.optim.Adam(
                model_parameters,
                lr=self.cfg.train.base_lr,
                betas=(0.9, 0.95),
                eps=1e-8,
            )
        else:
            raise NotImplementedError()
        
        return optimizer
        
    def _build_scheduler(self):
        if self.cfg.train.scheduler.lower() == "eden":
            scheduler = Eden(self.optimizer, 5000, 4, warmup_batches=self.cfg.train.warmup_steps)
        elif self.cfg.train.scheduler.lower() == "noam":
            scheduler = NoamScheduler(
                self.cfg.train.base_lr,
                self.optimizer,
                self.cfg.model.decoder_dim,
                warmup_steps=self.cfg.train.warmup_steps,
            )
        elif self.cfg.train.scheduler.lower() == "cosine":
            scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
                self.cfg.train.warmup_steps,
                self.optimizer,
                eta_min=self.cfg.train.base_lr,
            )
        else:
            raise NotImplementedError(f"{self.cfg.train.scheduler}")

        return scheduler      

    def _train_epoch(self):
        r"""Training epoch. Should return average loss of a batch (sample) over
        one epoch. See ``train_loop`` for usage.
        """
        if isinstance(self.model, dict):
            for key in self.model.keys():
                self.model[key].train()
        else:
            self.model.train()

        epoch_sum_loss: float = 0.0
        epoch_losses: dict = {}
        epoch_step: int = 0
        for batch in tqdm(
            self.train_dataloader,
            desc=f"Training Epoch {self.epoch}",
            unit="batch",
            colour="GREEN",
            leave=False,
            dynamic_ncols=True,
            smoothing=0.04,
            disable=not self.accelerator.is_main_process,
        ):

            # Do training step and BP
            with self.accelerator.accumulate(self.model):
                total_loss, train_losses = self._train_step(batch)
                self.accelerator.backward(total_loss)  
                self.optimizer.step()
                self.optimizer.zero_grad()   
            self.batch_count += 1  

            if self.batch_count % self.cfg.train.gradient_accumulation_step == 0:
                if self.cfg.train.optimizer not in ["ScaledAdam", "Eve"]:
                    torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)

                for k in range(self.cfg.train.gradient_accumulation_step):
                    if isinstance(self.scheduler, Eden):
                        self.scheduler.step_batch(self.step)
                    else:
                        self.scheduler.step()
                                        
                epoch_sum_loss += total_loss.detach().cpu().item()
                
                if isinstance(train_losses, dict):
                    for key, value in train_losses.items():
                        if key not in epoch_losses.keys():
                            epoch_losses[key] = value
                        else:
                            epoch_losses[key] += value

                if isinstance(train_losses, dict):
                    for key, loss in train_losses.items():
                        self.accelerator.log(
                            {"Step/Train {}".format(key): "{:.6f}".format(loss)},
                            step=self.step,
                        )
                else:
                    self.accelerator.log(
                        {"Step/Train Loss": loss},
                        step=self.step,
                    )
                                        
                self.accelerator.log(
                    {"Step/lr": self.scheduler.get_last_lr()[0]},
                    step=self.step,
                )
                
                # print loss every log_epoch_step steps
                if epoch_step % self.cfg.train.log_epoch_step == 0:
                    for key, loss in train_losses.items():
                        self.logger.info("Step/Train {}: {:.6f}".format(key, loss))
                        print("Step/Train {}: {:.6f}".format(key, loss))
                    

                self.step += 1
                epoch_step += 1
              
                
        self.accelerator.wait_for_everyone()

        epoch_sum_loss = (
            epoch_sum_loss
            / len(self.train_dataloader)
            * self.cfg.train.gradient_accumulation_step
        )

        for key in epoch_losses.keys():
            epoch_losses[key] = (
                epoch_losses[key]
                / len(self.train_dataloader)
                * self.cfg.train.gradient_accumulation_step
            )

        return epoch_sum_loss, epoch_losses
               
    def _get_state_dict(self):
        state_dict = {
            "model": self.model.state_dict(),
            "optimizer": self.optimizer.state_dict(),
            "scheduler": self.scheduler.state_dict(),
            "step": self.step,
            "epoch": self.epoch,
            "batch_size": self.cfg.train.batch_size,
            "trian_stage": self.args.train_stage
        }        
        return state_dict
    
    def _train_step(self, batch, is_training=True):
        text_tokens = batch["phone_seq"].to(self.device)
        text_tokens_lens = batch["phone_len"].to(self.device)
        assert text_tokens.ndim == 2

        audio_features = batch["acoustic_token"].to(self.device)
        audio_features_lens = batch["target_len"].to(self.device)
        assert audio_features.ndim == 3
  
                
        with torch.set_grad_enabled(is_training):
            loss, losses = self.model(
                x=text_tokens,
                x_lens=text_tokens_lens,
                y=audio_features,
                y_lens=audio_features_lens,
                train_stage=self.args.train_stage
            )
            
        assert loss.requires_grad == is_training

        loss_dict = {}
        frames_sum = (audio_features_lens).sum()
        
        avg_loss = loss / frames_sum
        
        loss_dict['loss'] = avg_loss.detach().cpu().item()
        for l in losses:
            loss_dict[l] = losses[l].detach().cpu().item() / frames_sum.item()
        
        return avg_loss, loss_dict
    
    def _valid_step(self, batch):
        valid_losses = {}
        total_loss = 0
        valid_stats = {}
        
        total_loss, valid_losses = self._train_step(
            batch=batch,
            is_training=False,
        )
        assert total_loss.requires_grad is False
        
        total_loss = total_loss.detach().cpu().item()
        
        return total_loss, valid_losses, valid_stats

    def add_arguments(parser: argparse.ArgumentParser):
            parser.add_argument(
                "--train_stage",
                type=int,
                default="1",
                help="0: train all modules, 1: AR Decoder, 2: NAR Decoder",
            )
