import random
import time
import subprocess
import traceback
import socket
import setproctitle

from torch.cuda.amp import GradScaler, autocast
import numpy as np
import torch.optim
import torch.utils.data
import copy
import logging
import os
import re
import sys
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import tqdm
import datetime

from utils.commons.ckpt_utils import get_last_checkpoint, get_all_ckpts
from utils.commons.ddp_utils import DDP
from utils.commons.hparams import hparams
from utils.commons.tensor_utils import move_to_cuda
from utils.commons.os_utils import remove_file
from utils.commons.meters import Timer


def check_port_is_occupied(host='localhost', port=10080):
    s = socket.socket()
    try:
        s.connect((host, port))
        print(f"{host}:{port} is occupied!")
        return True
    except:
        print(f"{host}:{port} is not occupied!")
        return False
    finally:
        s.close()


class Tee(object):
    def __init__(self, name, mode):
        self.file = open(name, mode)
        self.stdout = sys.stdout
        sys.stdout = self

    def __del__(self):
        sys.stdout = self.stdout
        self.file.close()

    def write(self, data):
        self.file.write(data)
        self.stdout.write(data)

    def flush(self):
        self.file.flush()


class Trainer:
    def __init__(
            self,
            work_dir,
            default_save_path=None,
            accumulate_grad_batches=1,
            max_updates=160000,
            print_nan_grads=False,
            val_check_interval=2000,
            num_sanity_val_steps=5,
            amp=False,
            # tb logger
            log_save_interval=100,
            tb_log_interval=10,
            # checkpoint
            monitor_key='val_loss',
            monitor_mode='min',
            num_ckpt_keep=5,
            save_best=True,
            resume_from_checkpoint=0,
            seed=1234,
            debug=False,
    ):
        os.makedirs(work_dir, exist_ok=True)
        self.work_dir = work_dir
        self.accumulate_grad_batches = accumulate_grad_batches
        self.max_updates = max_updates
        self.num_sanity_val_steps = num_sanity_val_steps
        self.print_nan_grads = print_nan_grads
        self.default_save_path = default_save_path
        self.resume_from_checkpoint = resume_from_checkpoint if resume_from_checkpoint > 0 else None
        self.seed = seed
        self.debug = debug
        # model and optm
        self.task = None
        self.optimizers = []

        # trainer state
        self.testing = False
        self.global_step = 0
        self.current_epoch = 0
        self.total_batches = 0

        # configure checkpoint
        self.monitor_key = monitor_key
        self.num_ckpt_keep = num_ckpt_keep
        self.save_best = save_best
        self.monitor_op = np.less if monitor_mode == 'min' else np.greater
        self.best_val_results = np.Inf if monitor_mode == 'min' else -np.Inf
        self.mode = 'min'

        # allow int, string and gpu list
        self.all_gpu_ids = [
            int(x) for x in os.environ.get("CUDA_VISIBLE_DEVICES", "").split(",") if x != '']
        
        self.use_multi_machine_ddp = hparams['world_size'] != -1
        self.num_local_gpus = len(self.all_gpu_ids) # if world_size is not -1, multi-machine setting
        self.num_total_gpus = len(self.all_gpu_ids) if not self.use_multi_machine_ddp else hparams['world_size'] # if world_size is not -1, multi-machine setting
        # self.num_gpus = len(self.all_gpu_ids) 
        self.on_gpu = self.num_local_gpus > 0
        self.root_gpu = 0
        logging.info(f'GPU available: {torch.cuda.is_available()}, GPU used: {self.all_gpu_ids}, world_size: {self.num_total_gpus}, multi-machine training: {self.use_multi_machine_ddp}')
        self.use_ddp = self.num_local_gpus > 1 or self.use_multi_machine_ddp
        self.proc_rank = 0
        # Tensorboard logging
        self.log_save_interval = log_save_interval
        self.val_check_interval = val_check_interval
        self.tb_log_interval = tb_log_interval
        self.amp = amp
        self.amp_scalar = GradScaler()

    def test(self, task_cls):
        self.testing = True
        self.fit(task_cls)

    def fit(self, task_cls):
        try:
            if self.use_ddp:
                # mp.spawn(self.ddp_run, nprocs=self.num_local_gpus, args=(task_cls, copy.deepcopy(hparams)))
                mp.start_processes(self.ddp_run,nprocs=self.num_local_gpus, args=(task_cls, copy.deepcopy(hparams)), start_method='spawn')
            else:
                # File "/mnt/bn/ailabrenyi/entries/yezhenhui/projects/GeneFace_private/venv_113/lib/python3.9/site-packages/torch/nn/modules/batchnorm.py", line 735, in forward
                # world_size = torch.distributed.get_world_size(process_group)
                # to address the error in batchNorm using venv_113
                # self.init_ddp_connection_tcp(0, 1) 
                self.task = task_cls()
                self.task.trainer = self
                setproctitle.setproctitle(f'GeneFace_worker ({hparams["work_dir"]})')
                self.run_single_process(self.task)
        except:
            traceback.print_exc()
            time.sleep(5)
            subprocess.check_call(f'pkill -f "GeneFace_worker \({hparams["work_dir"]}"', shell=True)
        return 1

        return 1

    def ddp_run(self, gpu_idx, task_cls, hparams_):
        hparams.update(hparams_)
        setproctitle.setproctitle(f'GeneFace_worker ({hparams_["work_dir"]}_{gpu_idx})')

        if hparams.get('use_file_system_mp'):
            torch.multiprocessing.set_sharing_strategy('file_system')
        if hparams.get('use_fork', True):
            torch.multiprocessing.set_start_method('fork', force=True)
        self.root_gpu = gpu_idx
        self.proc_rank = gpu_idx + hparams['start_rank'] if self.use_multi_machine_ddp else gpu_idx
        print("before init_tcp!")
        if hparams['init_method'] == 'file':
            self.init_ddp_connection_file(self.proc_rank, self.num_total_gpus)
        elif hparams['init_method'] == 'tcp':
            self.init_ddp_connection_tcp(self.proc_rank, self.num_total_gpus)
        else:
            raise NotImplementedError()

        if gpu_idx != 0 and not self.debug:
            sys.stdout = open(os.devnull, "w")
            sys.stderr = open(os.devnull, "w")
        dist.barrier()
        print("after init_tcp!")

        task = task_cls()
        task.trainer = self
        torch.cuda.set_device(gpu_idx)
        self.task = task
        self.run_single_process(task)

    def run_single_process(self, task):
        """Sanity check a few things before starting actual training.

        :param task:
        """
        # build model, optm and load checkpoint
        if self.proc_rank == 0:
            self.save_terminal_logs()
            if not self.testing:
                self.save_codes()

        model = task.build_model()
        if model is not None:
            task.model = model
        checkpoint, _ = get_last_checkpoint(self.work_dir, self.resume_from_checkpoint)
        if checkpoint is not None:
            self.restore_weights(checkpoint)
        elif self.on_gpu:
            task.cuda(self.root_gpu)
        if not self.testing:
            self.optimizers = task.configure_optimizers()
            self.fisrt_epoch = True
        if checkpoint is not None:
            self.restore_opt_state(checkpoint)
        del checkpoint
        # clear cache after restore
        if self.on_gpu:
            torch.cuda.empty_cache()

        if self.use_ddp:
            self.task = self.configure_ddp(self.task)
            dist.barrier()

        task_ref = self.get_task_ref()
        task_ref.trainer = self
        task_ref.testing = self.testing
        # link up experiment object
        if self.proc_rank == 0:
            task_ref.build_tensorboard(save_dir=self.work_dir, name='tb_logs')
        else:
            os.makedirs('tmp', exist_ok=True)
            task_ref.build_tensorboard(save_dir='tmp', name='tb_tmp')
        self.logger = task_ref.logger
        try:
            if self.testing:
                self.run_evaluation(test=True)
            else:
                self.train()
        except:
            traceback.print_exc()
            task_ref.on_keyboard_interrupt()
            time.sleep(5)
            if self.proc_rank == 0:
                subprocess.check_call(f'pkill -f "GeneFace_worker \({hparams["work_dir"]}"', shell=True)

    ####################
    # valid and test
    ####################
    def run_evaluation(self, test=False):
        eval_results = self.evaluate(self.task, test, tqdm_desc='Valid' if not test else 'test',
                                     max_batches=hparams['eval_max_batches'])
        if eval_results is not None and 'tb_log' in eval_results:
            tb_log_output = eval_results['tb_log']
            self.log_metrics_to_tb(tb_log_output)
        if self.proc_rank == 0 and not test:
            self.save_checkpoint(epoch=self.current_epoch, logs=eval_results)

    def evaluate(self, task, test=False, tqdm_desc='Valid', max_batches=None):
        if max_batches == -1:
            max_batches = None
        # enable eval mode
        task.zero_grad()
        task.eval()
        torch.set_grad_enabled(False)

        task_ref = self.get_task_ref()
        if test:
            ret = task_ref.test_start()
            if ret == 'EXIT':
                return
        else:
            task_ref.validation_start()
        outputs = []
        dataloader = task_ref.test_dataloader() if test else task_ref.val_dataloader()
        pbar = tqdm.tqdm(dataloader, desc=tqdm_desc, total=max_batches, dynamic_ncols=True, unit='step',
                         disable=self.root_gpu > 0)
        # give model a chance to do something with the outputs (and method defined)
        for batch_idx, batch in enumerate(pbar):
            if batch is None:  # pragma: no cover
                continue
            # stop short when on fast_dev_run (sets max_batch=1)
            if max_batches is not None and batch_idx >= max_batches:
                break

            # make dataloader_idx arg in validation_step optional
            if self.on_gpu:
                batch = move_to_cuda(batch, self.root_gpu)
            args = [batch, batch_idx]
            if self.use_ddp:
                output = task(*args)
            else:
                if test:
                    output = task_ref.test_step(*args)
                else:
                    output = task_ref.validation_step(*args)
            # track outputs for collation
            outputs.append(output)
        # give model a chance to do something with the outputs (and method defined)
        if test:
            eval_results = task_ref.test_end(outputs)
        else:
            eval_results = task_ref.validation_end(outputs)
        # enable train mode again
        task.train()
        torch.set_grad_enabled(True)
        return eval_results

    ####################
    # train
    ####################
    def train(self):
        task_ref = self.get_task_ref()
        task_ref.on_train_start()
        if self.num_sanity_val_steps > 0:
            # run tiny validation (if validation defined) to make sure program won't crash during val
            self.evaluate(self.task, False, 'Sanity Val', max_batches=self.num_sanity_val_steps)
        # clear cache before training
        if self.on_gpu:
            torch.cuda.empty_cache()
        dataloader = task_ref.train_dataloader()
        epoch = self.current_epoch
        # run all epochs
        while True:
            # set seed for distributed sampler (enables shuffling for each epoch)
            if self.use_ddp and hasattr(dataloader.sampler, 'set_epoch'):
                dataloader.sampler.set_epoch(epoch)
            # update training progress in trainer and model
            task_ref.current_epoch = epoch
            self.current_epoch = epoch
            # total batches includes multiple val checks
            self.batch_loss_value = 0  # accumulated grads
            # before epoch hook
            task_ref.on_epoch_start()

            # run epoch
            train_pbar = tqdm.tqdm(dataloader, initial=self.global_step, total=float('inf'),
                                   dynamic_ncols=True, unit='step', disable=self.root_gpu > 0)
            # for batch_idx, batch in enumerate(train_pbar):
            train_iterator = iter(enumerate(train_pbar))
            while True:
                with Timer("get_batch", enable=self.debug):
                    try:
                        batch_idx, batch = next(train_iterator)
                    except StopIteration:
                        train_iterator = iter(enumerate(train_pbar))
                        batch_idx, batch = next(train_iterator)

                if self.global_step % self.val_check_interval == 0 and not self.fisrt_epoch:
                    self.run_evaluation()
                pbar_metrics, tb_metrics = self.run_training_batch(batch_idx, batch)
                train_pbar.set_postfix(**pbar_metrics)
                self.fisrt_epoch = False
                # when metrics should be logged
                if (self.global_step + 1) % self.tb_log_interval == 0:
                    # logs user requested information to logger
                    self.log_metrics_to_tb(tb_metrics)

                self.global_step += 1
                task_ref.global_step = self.global_step
                if self.global_step > self.max_updates:
                    print("| Training end..")
                    break
            # epoch end hook
            epoch_loss_dict = task_ref.on_epoch_end()
            self.log_metrics_to_tb(epoch_loss_dict)
            epoch += 1
            if self.global_step > self.max_updates:
                break
        task_ref.on_train_end()

    def run_training_batch(self, batch_idx, batch):
        if batch is None:
            return {}
        all_progress_bar_metrics = []
        all_log_metrics = []
        task_ref = self.get_task_ref()
        for opt_idx, optimizer in enumerate(self.optimizers):
            if optimizer is None:
                continue
            # make sure only the gradients of the current optimizer's paramaters are calculated
            # in the training step to prevent dangling gradients in multiple-optimizer setup.
            if len(self.optimizers) > 1:
                for k, param in task_ref.named_parameters():
                    param.requires_grad = False
                for group in optimizer.param_groups:
                    for param in group['params']:
                        param.requires_grad = True

            # forward pass
            with Timer("forward_training_step", enable=self.debug):
                with autocast(enabled=self.amp):
                    if self.on_gpu:
                        batch = move_to_cuda(copy.copy(batch), self.root_gpu)
                    args = [batch, batch_idx, opt_idx]
                    if self.use_ddp:
                        output = self.task(*args)
                    else:
                        output = task_ref.training_step(*args)
                    loss = output['loss']
                    if loss is None:
                        continue
                    progress_bar_metrics = output['progress_bar']
                    log_metrics = output['tb_log']
                    # accumulate loss
                    loss = loss / self.accumulate_grad_batches
                
            # backward pass
            with Timer("backward_training_step", enable=self.debug):
                if loss.requires_grad:
                    if self.amp:
                        self.amp_scalar.scale(loss).backward()
                    else:
                        loss.backward()

                # track progress bar metrics
                all_log_metrics.append(log_metrics)
                all_progress_bar_metrics.append(progress_bar_metrics)

                if loss is None:
                    continue

            # nan grads
            with Timer("checkNan_training_step", enable=self.debug):
                has_nan_grad = False
                nan_params_names = []
                if self.print_nan_grads:
                    for name, param in task_ref.named_parameters():
                        if (param.grad is not None) and torch.isnan(param.grad.float()).any():
                            print("| NaN params: ", name, param, param.grad)
                            has_nan_grad = True
                            nan_params_names.append(name)
                    if has_nan_grad:
                        # exit(0)
                        print(f"| WARN: found nan in grad! first nan params: {nan_params_names[0]}; last nan params: {nan_params_names[-1]}.")
                        pass

            # gradient update with accumulated gradients
            with Timer("optimUpdate_training_step", enable=self.debug):
                if (self.global_step + 1) % self.accumulate_grad_batches == 0 and not has_nan_grad:
                # if (self.global_step + 1) % self.accumulate_grad_batches == 0:
                    # Unscales the gradients of optimizer's assigned params in-place
                    if self.amp:
                        self.amp_scalar.unscale_(optimizer)
                    grad_norm_dict = task_ref.on_before_optimization(opt_idx)
                    if grad_norm_dict is not None:
                        all_log_metrics[-1].update(grad_norm_dict)
                    if self.amp:
                        self.amp_scalar.step(optimizer)
                        self.amp_scalar.update()
                    else:
                        optimizer.step()
                    optimizer.zero_grad()
                    task_ref.on_after_optimization(self.current_epoch, batch_idx, optimizer, opt_idx)

        # collapse all metrics into one dict
        all_progress_bar_metrics = {k: v for d in all_progress_bar_metrics for k, v in d.items()}
        all_log_metrics = {k: v for d in all_log_metrics for k, v in d.items()}
        return all_progress_bar_metrics, all_log_metrics

    ####################
    # load and save checkpoint
    ####################
    def restore_weights(self, checkpoint):
        # load model state
        task_ref = self.get_task_ref()

        for k, v in checkpoint['state_dict'].items():
            if hasattr(task_ref, k):
                getattr(task_ref, k).load_state_dict(v, strict=True)
            else:
                print(f"| the checkpoint has unmatched keys {k}")

        if self.on_gpu:
            task_ref.cuda(self.root_gpu)
        # load training state (affects trainer only)
        self.best_val_results = checkpoint['checkpoint_callback_best']
        self.global_step = checkpoint['global_step']
        self.current_epoch = checkpoint['epoch']
        task_ref.global_step = self.global_step

        # wait for all models to restore weights
        if self.use_ddp:
            # wait for all processes to catch up
            dist.barrier()

    def restore_opt_state(self, checkpoint):
        if self.testing:
            return
        # restore the optimizers
        optimizer_states = checkpoint['optimizer_states']
        for optimizer, opt_state in zip(self.optimizers, optimizer_states):
            if optimizer is None:
                return
            try:
                optimizer.load_state_dict(opt_state)
                # move optimizer to GPU 1 weight at a time
                if self.on_gpu:
                    for state in optimizer.state.values():
                        for k, v in state.items():
                            if isinstance(v, torch.Tensor):
                                state[k] = v.cuda(self.root_gpu)
            except ValueError:
                print("| WARMING: optimizer parameters not match !!!")
        try:
            if dist.is_initialized() and dist.get_rank() > 0:
                return
        except Exception as e:
            print(e)
            return
        did_restore = True
        return did_restore

    def save_checkpoint(self, epoch, logs=None):
        monitor_op = np.less
        ckpt_path = f'{self.work_dir}/model_ckpt_steps_{self.global_step}.ckpt'
        logging.info(f'Epoch {epoch:05d}@{self.global_step}: saving model to {ckpt_path}')
        self._atomic_save(ckpt_path)

        get_ckpt_step_fn = lambda x: int(re.findall('.*steps\_(\d+)\.ckpt', x)[0])
        for old_ckpt in get_all_ckpts(self.work_dir)[self.num_ckpt_keep:]:
            # leave the milestone ckpts
            if hparams.get("ckpt_milestone_interval", 10_0000) != 0 and get_ckpt_step_fn(old_ckpt) % hparams.get("ckpt_milestone_interval", 10_0000) == 0:
                pass
            else:
                remove_file(old_ckpt)
                logging.info(f'Delete ckpt: {os.path.basename(old_ckpt)}')
        current = None
        if logs is not None and self.monitor_key in logs:
            current = logs[self.monitor_key]
        if current is not None and self.save_best:
            if monitor_op(current, self.best_val_results):
                best_filepath = f'{self.work_dir}/model_ckpt_best.pt'
                self.best_val_results = current
                logging.info(
                    f'Epoch {epoch:05d}@{self.global_step}: {self.monitor_key} reached {current:0.5f}. '
                    f'Saving model to {best_filepath}')
                self._atomic_save(best_filepath)

    def _atomic_save(self, filepath):
        checkpoint = self.dump_checkpoint()
        tmp_path = str(filepath) + ".part"
        torch.save(checkpoint, tmp_path, _use_new_zipfile_serialization=False)
        os.replace(tmp_path, filepath)
    
    def dump_checkpoint(self):
        checkpoint = {'epoch': self.current_epoch, 'global_step': self.global_step,
                      'checkpoint_callback_best': self.best_val_results}
        # save optimizers
        optimizer_states = []
        for i, optimizer in enumerate(self.optimizers):
            if optimizer is not None:
                state_dict = optimizer.state_dict()
                state_dict = {k.replace('_orig_mod.', ''): v for k, v in state_dict.items()}
                optimizer_states.append(state_dict)

        checkpoint['optimizer_states'] = optimizer_states
        task_ref = self.get_task_ref()
        state_dict = {
            k: v.state_dict() for k, v in task_ref.named_children()
            if len(list(v.parameters())) > 0 and (k not in hparams.get('not_save_modules', []))
        }
        for module_k, module_k_dict in list(state_dict.items()):
            for k, v in list(module_k_dict.items()):
                if '_orig_mod.' in k:
                    module_k_dict[k.replace('_orig_mod.', '')] = v
                    del module_k_dict[k]
            
        checkpoint['state_dict'] = state_dict
        return checkpoint
    ####################
    # DDP
    ####################
    def configure_ddp(self, task):
        task = torch.nn.SyncBatchNorm.convert_sync_batchnorm(task)
        task = DDP(task, device_ids=[self.root_gpu], find_unused_parameters=True)
        # task = DDP(task, device_ids=[self.root_gpu], find_unused_parameters=False)
        random.seed(self.seed)
        np.random.seed(self.seed)
        return task

    def init_ddp_connection_file(self, proc_rank, world_size):
        """
        use a shared file in the network file system to bind all process
        if you found num_worker is larger than world_size, remove the old shard_file_name
        """
        exp_name = hparams['exp_name']
        shared_file_name = f'file:///home/tiger/nfs/pytorch_ddp_sharedfile/{exp_name}'
        os.makedirs(os.path.dirname(shared_file_name).replace("file://",""), exist_ok=True)
        dist.init_process_group('nccl', init_method=shared_file_name,
                            world_size=world_size, rank=proc_rank)
        
    def init_ddp_connection_tcp(self, proc_rank, world_size):
        if self.use_multi_machine_ddp:
            root_node, port = os.environ['ARNOLD_WORKER_HOSTS'].split(",")[0].split(":")
            os.environ['MASTER_PORT'] = '6668'
        else:
            root_node = '127.0.0.1'
        root_node = self.resolve_root_node_address(root_node)
        os.environ['MASTER_ADDR'] = root_node

        dist.init_process_group('nccl', rank=proc_rank, world_size=world_size, timeout=datetime.timedelta(seconds=600))
        # dist.init_process_group('gloo', rank=proc_rank, world_size=world_size)

    def resolve_root_node_address(self, root_node):
        if '[' in root_node:
            name = root_node.split('[')[0]
            number = root_node.split(',')[0]
            if '-' in number:
                number = number.split('-')[0]
            number = re.sub('[^0-9]', '', number)
            root_node = name + number
        return root_node

    ####################
    # utils
    ####################
    def get_task_ref(self):
        from utils.commons.base_task import BaseTask
        task: BaseTask = self.task.module if isinstance(self.task, DDP) else self.task
        return task

    def log_metrics_to_tb(self, metrics, step=None):
        """Logs the metric dict passed in.

        :param metrics:
        """
        # turn all tensors to scalars
        scalar_metrics = self.metrics_to_scalars(metrics)

        step = step if step is not None else self.global_step
        # log actual metrics
        if self.proc_rank == 0:
            self.log_metrics(self.logger, scalar_metrics, step=step)

    @staticmethod
    def log_metrics(logger, metrics, step=None):
        for k, v in metrics.items():
            if isinstance(v, torch.Tensor):
                v = v.item()
            logger.add_scalar(k, v, step)

    def metrics_to_scalars(self, metrics):
        new_metrics = {}
        for k, v in metrics.items():
            if isinstance(v, torch.Tensor):
                v = v.item()

            if type(v) is dict:
                v = self.metrics_to_scalars(v)

            new_metrics[k] = v

        return new_metrics

    def save_terminal_logs(self):
        t = datetime.datetime.now().strftime('%Y%m%d%H%M%S')
        os.makedirs(f'{self.work_dir}/terminal_logs', exist_ok=True)
        Tee(f'{self.work_dir}/terminal_logs/log_{t}.txt', 'w')

    def save_codes(self):
        if len(hparams['save_codes']) > 0:
            t = datetime.datetime.now().strftime('%Y%m%d%H%M%S')
            code_dir = f'{self.work_dir}/codes/{t}'
            subprocess.check_call(f'mkdir -p "{code_dir}"', shell=True)
            for c in hparams['save_codes']:
                if os.path.exists(c):
                    subprocess.check_call(
                        f'rsync -aR '
                        f'--include="*.py" '
                        f'--include="*.yaml" '
                        f'--exclude="__pycache__" '
                        f'--include="*/" '
                        f'--exclude="*" '
                        f'"./{c}" "{code_dir}/"',
                        shell=True)
            print(f"| Copied codes to {code_dir}.")