import numpy as np
import torch
import sys, os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from models.pipeline import Pipeline
from datasets.dataloader import Collection_Loader
from models.util_learner import CriterionPoint, Optimizer
from tqdm import tqdm
torch.manual_seed(0)

class Trainer():
    def __init__(self, args, cfg):
        self.args = args
        print(f"[INFOR] Model: {cfg.regressor.name}")
        self.log_name = str(args.dataset) + "_" + str(args.scene) + "_" + str(cfg.regressor.name)
        self.pipeline = Pipeline(cfg)
        self.criterion = CriterionPoint(cfg.train.loss.reprojection, cfg.train.num_iters)
        self.device = torch.device(f'cuda:{args.cudaid}' if torch.cuda.is_available() else 'cpu')
        
        # to device
        self.pipeline.to(self.device)
        self.criterion.to(self.device)

        # dataloader
        train_collection = Collection_Loader(args, cfg, mode="train")
        print("[INFOR] Loaded data collection")
        self.train_loader = torch.utils.data.DataLoader(train_collection, batch_size=cfg.train.batch_size,
                                                        shuffle=cfg.train.loader_shuffle, num_workers=cfg.train.loader_num_workers, 
                                                        pin_memory=True)
        
        self.length_train_loader = len(self.train_loader)
        self.epochs = int(cfg.train.num_iters / self.length_train_loader)
        print(f"[INFOR] Total epochs: {self.epochs}")
        self.optimizer = Optimizer(self.pipeline.regressor.parameters(), self.epochs, self.length_train_loader, **cfg.optimizer)
        
        if self.args.checkpoint:
            # load checkpoint and resume training
            self.start_epoch = self.pipeline.load_checkpoint(self.args.outputs, self.log_name)
            # self.start_epoch = 2024
            self.lr = self.optimizer.adjust_lr(self.start_epoch)
        else:
            self.start_epoch = 0
            self.lr = self.optimizer.lr
        self.train_log = Train_Log(args, cfg, self.length_train_loader, self.start_epoch, self.epochs)
        

    def train(self):
        print("[INFOR] Start training")
        for epoch in range(self.start_epoch, self.epochs):
            if self.train_log.is_save_checkpoint():
                self.pipeline.save_checkpoint(self.args.outputs, self.log_name, epoch) # overwrite(save) checkpoint per epoch
            for batch_idx, (data, target) in enumerate(tqdm(self.train_loader)):
                iters = epoch*self.length_train_loader + batch_idx
                loss,_,_ = step_fwd(self.pipeline, self.device, data, target, iters, epoch,
                                   self.criterion, self.optimizer, train=True)
                self.train_log.update(epoch, batch_idx, loss, self.lr)
            self.lr = self.optimizer.adjust_lr(epoch) # adjust learning rate
            self.train_log.show(epoch) # show loss per epoch
        # self.pipeline.save_checkpoint(self.args.outputs, self.log_name, epoch, True)    


def step_fwd(model, device, data, target=None, iteration=2500000, epoch=0,
             criterion=None, optim=None, train=False):
    """
    A training/validation step."""
    if train: 
        assert criterion is not None
        assert target is not None
    for k,v in data.items():
        if isinstance(v,list):
            continue
        data[k] = data[k].to(device)
    if target is not None:
        for k,_ in target.items():
                target[k] = target[k].to(device)
    output = model(data)
    loss = None
    lr = None
    if train:
        loss = criterion(output, target, iteration)
        if optim is not None:
            optim.learner.zero_grad()
            loss[0].backward()
            optim.learner.step()
            lr = optim.lr_scheduler(epoch)
    return loss, output, lr

class Train_Log():
    def __init__(self, args, cfg, length_loader, start_epoch, total_epoch=0) -> None:
        self.args = args
        self.cfg = cfg
        self.total_epoch = total_epoch
        self.log_interval = cfg.train.log_interval
        self.length_train_loader = length_loader
        self.vis_env = str(args.dataset) + "_" + str(args.scene) + \
              "_" + str(cfg.regressor.name) +"_"+ str(args.experiment_version)
        self.showloss = ShowLosses(total_epoch=self.total_epoch)
        self.list_fignames = ['total_loss', 'point_loss', 'point_uncer_loss', 
                              'line_loss', 'line_uncer_loss', 'points_prj_loss',
                                'lines_prj_loss', 'learning_rate']
        if self.args.visdom:
            from visdom import Visdom
            print("[INFOR] Visdom is used for log visualization")
            self.vis = Visdom()
            for name in self.list_fignames:
                self.add_fig(name, start_epoch)

    def add_fig(self, name, start_epoch):
        self.vis.line(X=np.asarray([start_epoch]), Y=np.zeros(1), win=name,
                            opts={'legend': [name], 'xlabel': 'epochs',
                                'ylabel': name}, env=self.vis_env)
    def update_fig(self, idx, epoch_count, value):
        name = self.list_fignames[idx]
        self.vis.line(X=np.asarray([epoch_count]), Y=np.asarray([value]), win=name,
                        update='append', env=self.vis_env)
        
    def update(self, epoch, batch_idx, loss, lr):
        self.showloss.update(loss)
        self.lr = lr
        if self.args.visdom:
            if batch_idx % self.log_interval == 0:
                n_iter = epoch*self.length_train_loader + batch_idx
                epoch_count = float(n_iter)/self.length_train_loader
                l = len(self.list_fignames)
                for idx in range(l-1):
                    self.update_fig(idx, epoch_count, loss[idx].item())
                self.update_fig(l-1, epoch_count, lr)

    def show(self, epoch):
        self.showloss.show(epoch, self.lr)
    def is_save_checkpoint(self):
        his_epoch_loss = self.showloss.dict_losses[0].his_epoch_loss
        if len(his_epoch_loss) == 0:
            return False
        if min(his_epoch_loss) >= his_epoch_loss[-1]:
            return True
        else:
            return False

class His_Loss():
    def __init__(self)->None:
        self.his_epoch_loss = []
        self.temp_batch_loss = []
    def update_loss(self, loss):
        self.temp_batch_loss.append(loss)
    def show(self):
        avg_loss = np.mean(self.temp_batch_loss)
        self.his_epoch_loss.append(avg_loss)
        self.temp_batch_loss = [] # reset
        return avg_loss

class ShowLosses():
    # for debugging, showing all losses if needed
    def __init__(self, list_display=[True, True, False, True, False, True, True], total_epoch=0):
        '''
        corresponding to show following losses:
        ['total_loss', 'point_loss', 'point_uncer_loss', 
                              'line_loss', 'line_uncer_loss', 'points_prj_loss',
                                'lines_prj_loss']
        '''
        self.list_display = list_display
        self.length = len(self.list_display)
        self.names = ['Avg total loss', 'A.P.L', 'A.P.U.L', 'A.L.L', 'A.L.U.L', 'A.P.P.L', 'A.P.L.L']
        # A.P.L means average point loss, A.P.P.L means average point projection loss, etc.
        self.create_dict_losses()
        self.total_epoch = total_epoch

    def create_dict_losses(self):
        self.dict_losses = {}
        for i in range(self.length):
            if self.list_display[i]:
                self.dict_losses[i] = His_Loss()
    
    def update(self, loss):
        for k,_ in self.dict_losses.items():
            self.dict_losses[k].update_loss(loss[k].item())


    def show(self, epoch, lr=0.0):
        content = f"Epoch {epoch}/{self.total_epoch} | "
        for k,_ in self.dict_losses.items():
            avg_loss = self.dict_losses[k].show()
            content += self.names[k] + f": {avg_loss:.5f} | "
        content = content + f"lr: {lr:.6f}"
        print(content)

    