import logging
from datetime import datetime
from os.path import join
from pathlib import Path

import numpy as np
from tqdm import tqdm
import tensorflow as tf

from open3d.visualization.tensorboard_plugin import summary as summary3d
from .base_pipeline import BasePipeline
from ..modules.losses import SemSegLoss
from ..modules.metrics import SemSegMetric
from ..dataloaders import TFDataloader
from ...utils import make_dir, PIPELINE, get_runid, code2md

log = logging.getLogger(__name__)


class SemanticSegmentation(BasePipeline):
    """This class allows you to perform semantic segmentation for both training
    and inference using the TensorFlow framework. This pipeline has multiple
    stages: Pre-processing, loading dataset, testing, and inference or training.

    **Example:**
        This example loads the Semantic Segmentation and performs a training
        using the SemanticKITTI dataset.

        .. code::

            import tensorflow as tf
            from .base_pipeline import BasePipeline

            Mydataset = TFDataloader(dataset=tf.dataset.get_split('training')
            MyModel = SemanticSegmentation(self,model,dataset=Mydataset,
                name='SemanticSegmentation',
                batch_size=4,
                val_batch_size=4,
                test_batch_size=3,
                max_epoch=100,
                learning_rate=1e-2,
                lr_decays=0.95,
                save_ckpt_freq=20,
                adam_lr=1e-2,
                scheduler_gamma=0.95,
                momentum=0.98,
                main_log_dir='./logs/',
                device='gpu',
                split='train',
                train_sum_dir='train_log')

    **Args:**
            dataset: The 3D ML dataset class. You can use the base dataset,
                sample datasets, or a custom dataset.
            model: The model to be used for building the pipeline.
            name: The name of the current training.
            batch_size: The batch size to be used for training.
            val_batch_size: The batch size to be used for validation.
            test_batch_size: The batch size to be used for testing.
            max_epoch: The maximum size of the epoch to be used for training.
            leanring_rate: The hyperparameter that controls the weights during
                training. Also, known as step size.
            lr_decays: The learning rate decay for the training.
            save_ckpt_freq: The frequency in which the checkpoint should be
                saved.
            adam_lr: The leanring rate to be applied for Adam optimization.
            scheduler_gamma: The decaying factor associated with the scheduler.
            momentum: The momentum that accelerates the training rate schedule.
            main_log_dir: The directory where logs are stored.
            device: The device to be used for training.
            split: The dataset split to be used. In this example, we have used
                "train".
            train_sum_dir: The directory where the trainig summary is stored.

    **Returns:**
            class: The corresponding class.
    """

    def __init__(
            self,
            model,
            dataset=None,
            name='SemanticSegmentation',
            batch_size=4,
            val_batch_size=4,
            test_batch_size=3,
            max_epoch=100,  # maximum epoch during training
            learning_rate=1e-2,  # initial learning rate
            lr_decays=0.95,
            save_ckpt_freq=20,
            adam_lr=1e-2,
            scheduler_gamma=0.95,
            momentum=0.98,
            main_log_dir='./logs/',
            device='gpu',
            split='train',
            train_sum_dir='train_log',
            **kwargs):

        super().__init__(model=model,
                         dataset=dataset,
                         name=name,
                         batch_size=batch_size,
                         val_batch_size=val_batch_size,
                         test_batch_size=test_batch_size,
                         max_epoch=max_epoch,
                         learning_rate=learning_rate,
                         lr_decays=lr_decays,
                         save_ckpt_freq=save_ckpt_freq,
                         adam_lr=adam_lr,
                         scheduler_gamma=scheduler_gamma,
                         momentum=momentum,
                         main_log_dir=main_log_dir,
                         device=device,
                         split=split,
                         train_sum_dir=train_sum_dir,
                         **kwargs)
        self.summary = {}
        self.cfg.convert_to_tf_names('pipeline')

    def run_inference(self, data):
        """Run the inference using the data passed."""
        model = self.model
        log.info("running inference")

        model.inference_begin(data)

        while True:
            inputs = model.inference_preprocess()
            results = model(inputs, training=False)
            if model.inference_end(results):
                break

        metric = SemSegMetric()
        metric.update(
            tf.convert_to_tensor(model.inference_result['predict_scores']),
            tf.convert_to_tensor(data['label']))
        log.info(f"Accuracy : {metric.acc()}")
        log.info(f"IoU : {metric.iou()}")

        return model.inference_result

    def run_test(self):
        """Run the test using the data passed."""
        model = self.model
        dataset = self.dataset
        cfg = self.cfg

        self.load_ckpt(model.cfg.ckpt_path)
        timestamp = datetime.now().strftime('%Y-%m-%d_%H:%M:%S')

        log_file_path = join(cfg.logs_dir, 'log_test_' + timestamp + '.txt')
        log.info("Logging in file : {}".format(log_file_path))
        log.addHandler(logging.FileHandler(log_file_path))

        record_summary = cfg.get('summary').get('record_for', [])
        log.info("Started testing")

        metric = SemSegMetric()
        Loss = SemSegLoss(self, model, dataset)

        test_split = dataset.get_split('test')
        for idx in tqdm(range(len(test_split)), desc='test'):
            attr = test_split.get_attr(idx)
            data = test_split.get_data(idx)
            data['label'] = data['label'].astype(np.int32)
            results = self.run_inference(data)
            scores, labels = Loss.filter_valid_label(results['predict_scores'],
                                                     data['label'])
            metric.update(scores, labels)
            dataset.save_test_result(results, attr)
            # Save only for the first batch
            if 'test' in record_summary and 'test' not in self.summary:
                self.summary['test'] = self.get_3d_summary(tf.convert_to_tensor(
                    results['predict_scores']),
                                                           data,
                                                           0,
                                                           save_gt=False)

        accs = metric.acc()
        ious = metric.iou()

        log.info("Per class Accuracy : {}".format(accs[:-1]))
        log.info("Per class IOUs : {}".format(ious[:-1]))
        log.info("Overall Accuracy : {:.3f}".format(accs[-1]))
        log.info("Overall IOU : {:.3f}".format(ious[-1]))

    def run_train(self):
        """Run model training."""
        model = self.model
        dataset = self.dataset

        cfg = self.cfg

        log.info(model)
        timestamp = datetime.now().strftime('%Y-%m-%d_%H:%M:%S')
        log_file_path = join(cfg.logs_dir, 'log_train_' + timestamp + '.txt')
        log.info("Logging in file : {}".format(log_file_path))
        log.addHandler(logging.FileHandler(log_file_path))

        Loss = SemSegLoss(self, model, dataset)
        self.metric_train = SemSegMetric()
        self.metric_val = SemSegMetric()

        train_split = TFDataloader(dataset=dataset.get_split('training'),
                                   model=model,
                                   use_cache=dataset.cfg.use_cache,
                                   steps_per_epoch=dataset.cfg.get(
                                       'steps_per_epoch_train', None))
        train_loader, len_train = train_split.get_loader(cfg.batch_size)

        valid_split = TFDataloader(dataset=dataset.get_split('validation'),
                                   model=model,
                                   use_cache=dataset.cfg.use_cache,
                                   steps_per_epoch=dataset.cfg.get(
                                       'steps_per_epoch_valid', None))
        valid_loader, len_val = valid_split.get_loader(cfg.val_batch_size)

        dataset_name = dataset.name if dataset is not None else ''
        tensorboard_dir = join(
            self.cfg.train_sum_dir,
            model.__class__.__name__ + '_' + dataset_name + '_tf')
        runid = get_runid(tensorboard_dir)
        self.tensorboard_dir = join(self.cfg.train_sum_dir,
                                    runid + '_' + Path(tensorboard_dir).name)

        writer = tf.summary.create_file_writer(self.tensorboard_dir)
        self.save_config(writer)
        log.info("Writing summary in {}.".format(self.tensorboard_dir))
        record_summary = cfg.get('summary').get('record_for', [])
        self.optimizer = model.get_optimizer(cfg)

        is_resume = model.cfg.get('is_resume', True)
        self.load_ckpt(model.cfg.ckpt_path, is_resume=is_resume)

        for epoch in range(0, cfg.max_epoch + 1):
            log.info("=== EPOCH {}/{} ===".format(epoch, cfg.max_epoch))
            # --------------------- training
            self.metric_train.reset()
            self.metric_val.reset()
            self.losses = []
            step = 0

            for inputs in tqdm(train_loader, total=len_train, desc='training'):
                with tf.GradientTape(persistent=True) as tape:
                    results = model(inputs, training=True)

                    loss, gt_labels, predict_scores = model.get_loss(
                        Loss, results, inputs)

                if len(predict_scores.shape) < 2:
                    continue

                if predict_scores.shape[0] == 0:
                    continue
                scaled_params = []
                params = []
                for val in model.trainable_weights:
                    if 'deform' in val.name:
                        scaled_params.append(val)
                    else:
                        params.append(val)

                grads = tape.gradient(loss, params)
                scaled_grads = tape.gradient(loss, scaled_params)
                for i in range(len(scaled_grads)):
                    scaled_grads[i] *= 0.1

                norm = cfg.get('grad_clip_norm', 100.0)
                grads = [tf.clip_by_norm(g, norm) for g in grads]
                scaled_grads = [tf.clip_by_norm(g, norm) for g in scaled_grads]

                self.optimizer.apply_gradients(zip(grads, params))

                if len(scaled_grads) > 0:
                    self.optimizer.apply_gradients(
                        zip(scaled_grads, scaled_params))

                self.metric_train.update(predict_scores, gt_labels)
                self.losses.append(loss.numpy())

                if 'train' in record_summary and step == 0:
                    self.summary['train'] = self.get_3d_summary(
                        results, inputs, epoch)
                step = step + 1

            # --------------------- validation
            self.valid_losses = []
            step = 0

            for inputs in tqdm(valid_loader, total=len_val, desc='validation'):
                with tf.GradientTape() as tape:
                    results = model(inputs, training=False)
                    loss, gt_labels, predict_scores = model.get_loss(
                        Loss, results, inputs)

                if len(predict_scores.shape) < 2:
                    continue

                self.metric_val.update(predict_scores, gt_labels)
                self.valid_losses.append(loss.numpy())

                if 'valid' in record_summary and step == 0:
                    self.summary['valid'] = self.get_3d_summary(
                        results, inputs, epoch)
                step = step + 1

            self.save_logs(writer, epoch)

            if epoch % cfg.save_ckpt_freq == 0:
                self.save_ckpt(epoch)

    def get_3d_summary(self, results, input_data, epoch, save_gt=True):
        """
        Create visualization for network inputs and outputs.

        Args:
            results: Model output (see below).
            input_data: Model input (see below).
            epoch (int): step
            save_gt (bool): Save ground truth (for 'train' or 'valid' stages).

        RandLaNet:
            results (Tensor(B, N, C)): Prediction scores for all classes.
            input_data (Tuple): Batch of pointclouds and labels.
                input_data[0] (Tensor(B,N,3), float) : points
                input_data[-1] (Tensor(B,N), int) : labels

        SparseConvUNet:
            results (Tensor(SN, C)): Prediction scores for all classes. SN is
                total points in the batch.
            input_data (Dict): Batch of pointclouds and labels. Keys should be:
                'point' [Tensor(SN,3), float]: Concatenated points.
                'batch_lengths' [Tensor(B,), int]: Number of points in each
                    point cloud of the batch.
                'label' [Tensor(SN,) (optional)]: Concatenated labels.

        Returns:
            [Dict] visualizations of inputs and outputs suitable to save as an
                Open3D for TensorBoard summary.
        """
        if not hasattr(self, "_first_step"):
            self._first_step = epoch
        label_to_names = self.dataset.get_label_to_names()
        cfg = self.cfg.get('summary')
        max_pts = cfg.get('max_pts')
        if max_pts is None:
            max_pts = np.iinfo(np.int32).max
        use_reference = cfg.get('use_reference', False)
        max_outputs = cfg.get('max_outputs', 1)
        input_pcd = []
        gt_labels = []
        predict_labels = []

        def to_sum_fmt(tensor, add_dims=(0, 0), dtype=np.int32):
            np_tensor = tensor.numpy().astype(dtype)
            new_shape = (1,) * add_dims[0] + np_tensor.shape + (
                1,) * add_dims[1]
            return np_tensor.reshape(new_shape)

        # Variable size point clouds
        if self.model.cfg['name'] in ('KPFCNN', 'KPConv'):
            org_inputs = self.model.organise_inputs(input_data)
            predict_labels.append(to_sum_fmt(tf.argmax(results, 1), (0, 1)))
            if self._first_step == epoch or not use_reference:
                pointcloud = org_inputs['points'][0]
                input_pcd.append(
                    to_sum_fmt(pointcloud[:, :3], (0, 0), np.float32))
                if save_gt:
                    gtl = org_inputs['point_labels']
                    gt_labels.append(to_sum_fmt(gtl, (0, 1)))

        # Dict input, variable size point clouds
        elif self.model.cfg['name'] in ('SparseConvUnet', 'PointTransformer'):
            if self.model.cfg['name'] == 'SparseConvUnet':
                row_splits = np.hstack(
                    ((0,), np.cumsum(input_data['batch_lengths'])))
            else:
                row_splits = input_data.get('row_splits',
                                            [0, input_data['point'].shape[0]])
            max_outputs = min(max_outputs, len(row_splits) - 1)
            for k in range(max_outputs):
                blen_k = row_splits[k + 1] - row_splits[k]
                pcd_step = int(np.ceil(blen_k / min(max_pts, blen_k)))
                res_pcd = results[row_splits[k]:row_splits[k + 1]:pcd_step, :]
                predict_labels.append(to_sum_fmt(tf.argmax(res_pcd, 1), (0, 1)))
                if self._first_step != epoch and use_reference:
                    continue
                pointcloud = input_data['point'][
                    row_splits[k]:row_splits[k + 1]:pcd_step]
                pointcloud = tf.convert_to_tensor(pointcloud)
                input_pcd.append(
                    to_sum_fmt(pointcloud[:, :3], (0, 0), np.float32))
                if save_gt:
                    gtl = input_data['label'][
                        row_splits[k]:row_splits[k + 1]:pcd_step]
                    gt_labels.append(to_sum_fmt(gtl, (0, 1)))

        # Fixed size point clouds
        # Tuple input, same size point clouds
        elif self.model.cfg['name'] in ('RandLANet', 'PVCNN'):
            if self.model.cfg['name'] == 'RandLANet':
                if save_gt:
                    pointcloud = input_data[0]  # 0 => input to first layer
                else:  # Structured data during inference
                    pointcloud = tf.expand_dims(
                        tf.convert_to_tensor(input_data['point']), 0)
                    results = tf.expand_dims(results, 0)
            elif self.model.cfg['name'] == 'PVCNN':
                pointcloud = input_data['point']
            pcd_step = int(
                np.ceil(pointcloud.shape[1] /
                        min(max_pts, pointcloud.shape[1])))
            predict_labels = to_sum_fmt(
                tf.argmax(results[:max_outputs, ::pcd_step, :], -1), (0, 1))
            if self._first_step == epoch or not use_reference:
                input_pcd = to_sum_fmt(pointcloud[:max_outputs, ::pcd_step, :3],
                                       (0, 0), np.float32)
                if save_gt:
                    if self.model.cfg['name'] == 'RandLANet':
                        gtl = input_data[-1]
                        if 'int' not in repr(gtl.dtype).lower():
                            raise TypeError(
                                "Labels should be Int types. Received "
                                f"{gtl.dtype}")
                    elif self.model.cfg['name'] == 'PVCNN':
                        gtl = input_data['label']
                    gt_labels = to_sum_fmt(gtl[:max_outputs, ::pcd_step],
                                           (0, 1))
        else:
            raise NotImplementedError(
                "Saving 3D summary for the model "
                f"{self.model.cfg['name']} is not implemented.")

        def get_reference_or(data_tensor):
            if self._first_step == epoch or not use_reference:
                return data_tensor
            return self._first_step

        summary_dict = {
            'semantic_segmentation': {
                "vertex_positions": get_reference_or(input_pcd),
                "vertex_gt_labels": get_reference_or(gt_labels),
                "vertex_predict_labels": predict_labels,
                'label_to_names': label_to_names
            }
        }
        return summary_dict

    def save_logs(self, writer, epoch):
        """Save logs from the training and send results to TensorBoard."""
        train_accs = self.metric_train.acc()
        val_accs = self.metric_val.acc()

        train_ious = self.metric_train.iou()
        val_ious = self.metric_val.iou()

        loss_dict = {
            'Training loss': np.mean(self.losses),
            'Validation loss': np.mean(self.valid_losses)
        }
        acc_dicts = [{
            'Training accuracy': acc,
            'Validation accuracy': val_acc
        } for acc, val_acc in zip(train_accs, val_accs)]
        iou_dicts = [{
            'Training IoU': iou,
            'Validation IoU': val_iou
        } for iou, val_iou in zip(train_ious, val_ious)]

        log.info(f"loss train: {loss_dict['Training loss']:.3f} "
                 f" eval: {loss_dict['Validation loss']:.3f}")
        log.info(f"acc train: {acc_dicts[-1]['Training accuracy']:.3f} "
                 f" eval: {acc_dicts[-1]['Validation accuracy']:.3f}")
        log.info(f"iou train: {iou_dicts[-1]['Training IoU']:.3f} "
                 f" eval: {iou_dicts[-1]['Validation IoU']:.3f}")

        # send results to tensorboard
        with writer.as_default():
            for key, val in loss_dict.items():
                tf.summary.scalar(key, val, epoch)

            for key, val in acc_dicts[-1].items():
                tf.summary.scalar("{}/ Overall".format(key), val, epoch)
            for key, val in iou_dicts[-1].items():
                tf.summary.scalar("{}/ Overall".format(key), val, epoch)

            for stage in self.summary:
                for key, summary_dict in self.summary[stage].items():
                    label_to_names = summary_dict.pop('label_to_names', None)
                    summary3d.add_3d('/'.join((stage, key)),
                                     summary_dict,
                                     epoch,
                                     max_outputs=0,
                                     label_to_names=label_to_names,
                                     logdir=self.tensorboard_dir)

    def load_ckpt(self, ckpt_path=None, is_resume=True):
        """Load a checkpoint. You must pass the checkpoint and indicate if you want
        to resume.
        """
        train_ckpt_dir = join(self.cfg.logs_dir, 'checkpoint')
        make_dir(train_ckpt_dir)

        if hasattr(self, 'optimizer'):
            self.ckpt = tf.train.Checkpoint(step=tf.Variable(1),
                                            optimizer=self.optimizer,
                                            net=self.model)
        else:
            self.ckpt = tf.train.Checkpoint(step=tf.Variable(1), net=self.model)

        self.manager = tf.train.CheckpointManager(self.ckpt,
                                                  train_ckpt_dir,
                                                  max_to_keep=100)

        if ckpt_path is not None:
            self.ckpt.restore(ckpt_path).expect_partial()
            log.info("Restored from {}".format(ckpt_path))
        else:
            self.ckpt.restore(self.manager.latest_checkpoint)

            if self.manager.latest_checkpoint and is_resume:
                log.info("Restored from {}".format(
                    self.manager.latest_checkpoint))
            else:
                log.info("Initializing from scratch.")

    def save_ckpt(self, epoch):
        """Save a checkpoint at the passed epoch."""
        save_path = self.manager.save()
        log.info("Saved checkpoint at: {}".format(save_path))

    def save_config(self, writer):
        """Save experiment configuration with tensorboard summary."""
        if hasattr(self, 'cfg_tb'):
            with writer.as_default():
                with tf.name_scope("Description"):
                    tf.summary.text("Open3D-ML", self.cfg_tb['readme'], step=0)
                    tf.summary.text("Command line",
                                    self.cfg_tb['cmd_line'],
                                    step=0)
                with tf.name_scope("Configuration"):
                    tf.summary.text('Dataset',
                                    code2md(self.cfg_tb['dataset'],
                                            language='json'),
                                    step=0)
                    tf.summary.text('Model',
                                    code2md(self.cfg_tb['model'],
                                            language='json'),
                                    step=0)
                    tf.summary.text('Pipeline',
                                    code2md(self.cfg_tb['pipeline'],
                                            language='json'),
                                    step=0)


PIPELINE._register_module(SemanticSegmentation, "tf")

import logging
from datetime import datetime
from os.path import join
from pathlib import Path

import numpy as np
from tqdm import tqdm
import tensorflow as tf

from open3d.visualization.tensorboard_plugin import summary as summary3d
from .base_pipeline import BasePipeline
from ..modules.losses import SemSegLoss
from ..modules.metrics import SemSegMetric
from ..dataloaders import TFDataloader
from ...utils import make_dir, PIPELINE, get_runid, code2md

log = logging.getLogger(__name__)


class SemanticSegmentation(BasePipeline):
    """This class allows you to perform semantic segmentation for both training
    and inference using the TensorFlow framework. This pipeline has multiple
    stages: Pre-processing, loading dataset, testing, and inference or training.

    **Example:**
        This example loads the Semantic Segmentation and performs a training
        using the SemanticKITTI dataset.

        .. code::

            import tensorflow as tf
            from .base_pipeline import BasePipeline

            Mydataset = TFDataloader(dataset=tf.dataset.get_split('training')
            MyModel = SemanticSegmentation(self,model,dataset=Mydataset,
                name='SemanticSegmentation',
                batch_size=4,
                val_batch_size=4,
                test_batch_size=3,
                max_epoch=100,
                learning_rate=1e-2,
                lr_decays=0.95,
                save_ckpt_freq=20,
                adam_lr=1e-2,
                scheduler_gamma=0.95,
                momentum=0.98,
                main_log_dir='./logs/',
                device='gpu',
                split='train',
                train_sum_dir='train_log')

    **Args:**
            dataset: The 3D ML dataset class. You can use the base dataset,
                sample datasets, or a custom dataset.
            model: The model to be used for building the pipeline.
            name: The name of the current training.
            batch_size: The batch size to be used for training.
            val_batch_size: The batch size to be used for validation.
            test_batch_size: The batch size to be used for testing.
            max_epoch: The maximum size of the epoch to be used for training.
            leanring_rate: The hyperparameter that controls the weights during
                training. Also, known as step size.
            lr_decays: The learning rate decay for the training.
            save_ckpt_freq: The frequency in which the checkpoint should be
                saved.
            adam_lr: The leanring rate to be applied for Adam optimization.
            scheduler_gamma: The decaying factor associated with the scheduler.
            momentum: The momentum that accelerates the training rate schedule.
            main_log_dir: The directory where logs are stored.
            device: The device to be used for training.
            split: The dataset split to be used. In this example, we have used
                "train".
            train_sum_dir: The directory where the trainig summary is stored.

    **Returns:**
            class: The corresponding class.
    """

    def __init__(
            self,
            model,
            dataset=None,
            name='SemanticSegmentation',
            batch_size=4,
            val_batch_size=4,
            test_batch_size=3,
            max_epoch=100,  # maximum epoch during training
            learning_rate=1e-2,  # initial learning rate
            lr_decays=0.95,
            save_ckpt_freq=20,
            adam_lr=1e-2,
            scheduler_gamma=0.95,
            momentum=0.98,
            main_log_dir='./logs/',
            device='gpu',
            split='train',
            train_sum_dir='train_log',
            **kwargs):

        super().__init__(model=model,
                         dataset=dataset,
                         name=name,
                         batch_size=batch_size,
                         val_batch_size=val_batch_size,
                         test_batch_size=test_batch_size,
                         max_epoch=max_epoch,
                         learning_rate=learning_rate,
                         lr_decays=lr_decays,
                         save_ckpt_freq=save_ckpt_freq,
                         adam_lr=adam_lr,
                         scheduler_gamma=scheduler_gamma,
                         momentum=momentum,
                         main_log_dir=main_log_dir,
                         device=device,
                         split=split,
                         train_sum_dir=train_sum_dir,
                         **kwargs)
        self.summary = {}
        self.cfg.convert_to_tf_names('pipeline')

    def run_inference(self, data):
        """Run the inference using the data passed."""
        model = self.model
        log.info("running inference")

        model.inference_begin(data)

        while True:
            inputs = model.inference_preprocess()
            results = model(inputs, training=False)
            if model.inference_end(results):
                break

        metric = SemSegMetric()
        metric.update(
            tf.convert_to_tensor(model.inference_result['predict_scores']),
            tf.convert_to_tensor(data['label']))
        log.info(f"Accuracy : {metric.acc()}")
        log.info(f"IoU : {metric.iou()}")

        return model.inference_result

    def run_test(self):
        """Run the test using the data passed."""
        model = self.model
        dataset = self.dataset
        cfg = self.cfg

        self.load_ckpt(model.cfg.ckpt_path)
        timestamp = datetime.now().strftime('%Y-%m-%d_%H:%M:%S')

        log_file_path = join(cfg.logs_dir, 'log_test_' + timestamp + '.txt')
        log.info("Logging in file : {}".format(log_file_path))
        log.addHandler(logging.FileHandler(log_file_path))

        record_summary = cfg.get('summary').get('record_for', [])
        log.info("Started testing")

        metric = SemSegMetric()
        Loss = SemSegLoss(self, model, dataset)

        test_split = dataset.get_split('test')
        for idx in tqdm(range(len(test_split)), desc='test'):
            attr = test_split.get_attr(idx)
            data = test_split.get_data(idx)
            data['label'] = data['label'].astype(np.int32)
            results = self.run_inference(data)
            scores, labels = Loss.filter_valid_label(results['predict_scores'],
                                                     data['label'])
            metric.update(scores, labels)
            dataset.save_test_result(results, attr)
            # Save only for the first batch
            if 'test' in record_summary and 'test' not in self.summary:
                self.summary['test'] = self.get_3d_summary(tf.convert_to_tensor(
                    results['predict_scores']),
                                                           data,
                                                           0,
                                                           save_gt=False)

        accs = metric.acc()
        ious = metric.iou()

        log.info("Per class Accuracy : {}".format(accs[:-1]))
        log.info("Per class IOUs : {}".format(ious[:-1]))
        log.info("Overall Accuracy : {:.3f}".format(accs[-1]))
        log.info("Overall IOU : {:.3f}".format(ious[-1]))

    def run_train(self):
        """Run model training."""
        model = self.model
        dataset = self.dataset

        cfg = self.cfg

        log.info(model)
        timestamp = datetime.now().strftime('%Y-%m-%d_%H:%M:%S')
        log_file_path = join(cfg.logs_dir, 'log_train_' + timestamp + '.txt')
        log.info("Logging in file : {}".format(log_file_path))
        log.addHandler(logging.FileHandler(log_file_path))

        Loss = SemSegLoss(self, model, dataset)
        self.metric_train = SemSegMetric()
        self.metric_val = SemSegMetric()

        train_split = TFDataloader(dataset=dataset.get_split('training'),
                                   model=model,
                                   use_cache=dataset.cfg.use_cache,
                                   steps_per_epoch=dataset.cfg.get(
                                       'steps_per_epoch_train', None))
        train_loader, len_train = train_split.get_loader(cfg.batch_size)

        valid_split = TFDataloader(dataset=dataset.get_split('validation'),
                                   model=model,
                                   use_cache=dataset.cfg.use_cache,
                                   steps_per_epoch=dataset.cfg.get(
                                       'steps_per_epoch_valid', None))
        valid_loader, len_val = valid_split.get_loader(cfg.val_batch_size)

        dataset_name = dataset.name if dataset is not None else ''
        tensorboard_dir = join(
            self.cfg.train_sum_dir,
            model.__class__.__name__ + '_' + dataset_name + '_tf')
        runid = get_runid(tensorboard_dir)
        self.tensorboard_dir = join(self.cfg.train_sum_dir,
                                    runid + '_' + Path(tensorboard_dir).name)

        writer = tf.summary.create_file_writer(self.tensorboard_dir)
        self.save_config(writer)
        log.info("Writing summary in {}.".format(self.tensorboard_dir))
        record_summary = cfg.get('summary').get('record_for', [])
        self.optimizer = model.get_optimizer(cfg)

        is_resume = model.cfg.get('is_resume', True)
        self.load_ckpt(model.cfg.ckpt_path, is_resume=is_resume)

        for epoch in range(0, cfg.max_epoch + 1):
            log.info("=== EPOCH {}/{} ===".format(epoch, cfg.max_epoch))
            # --------------------- training
            self.metric_train.reset()
            self.metric_val.reset()
            self.losses = []
            step = 0

            for inputs in tqdm(train_loader, total=len_train, desc='training'):
                with tf.GradientTape(persistent=True) as tape:
                    results = model(inputs, training=True)

                    loss, gt_labels, predict_scores = model.get_loss(
                        Loss, results, inputs)

                if len(predict_scores.shape) < 2:
                    continue

                if predict_scores.shape[0] == 0:
                    continue
                scaled_params = []
                params = []
                for val in model.trainable_weights:
                    if 'deform' in val.name:
                        scaled_params.append(val)
                    else:
                        params.append(val)

                grads = tape.gradient(loss, params)
                scaled_grads = tape.gradient(loss, scaled_params)
                for i in range(len(scaled_grads)):
                    scaled_grads[i] *= 0.1

                norm = cfg.get('grad_clip_norm', 100.0)
                grads = [tf.clip_by_norm(g, norm) for g in grads]
                scaled_grads = [tf.clip_by_norm(g, norm) for g in scaled_grads]

                self.optimizer.apply_gradients(zip(grads, params))

                if len(scaled_grads) > 0:
                    self.optimizer.apply_gradients(
                        zip(scaled_grads, scaled_params))

                self.metric_train.update(predict_scores, gt_labels)
                self.losses.append(loss.numpy())

                if 'train' in record_summary and step == 0:
                    self.summary['train'] = self.get_3d_summary(
                        results, inputs, epoch)
                step = step + 1

            # --------------------- validation
            self.valid_losses = []
            step = 0

            for inputs in tqdm(valid_loader, total=len_val, desc='validation'):
                with tf.GradientTape() as tape:
                    results = model(inputs, training=False)
                    loss, gt_labels, predict_scores = model.get_loss(
                        Loss, results, inputs)

                if len(predict_scores.shape) < 2:
                    continue

                self.metric_val.update(predict_scores, gt_labels)
                self.valid_losses.append(loss.numpy())

                if 'valid' in record_summary and step == 0:
                    self.summary['valid'] = self.get_3d_summary(
                        results, inputs, epoch)
                step = step + 1

            self.save_logs(writer, epoch)

            if epoch % cfg.save_ckpt_freq == 0:
                self.save_ckpt(epoch)

    def get_3d_summary(self, results, input_data, epoch, save_gt=True):
        """
        Create visualization for network inputs and outputs.

        Args:
            results: Model output (see below).
            input_data: Model input (see below).
            epoch (int): step
            save_gt (bool): Save ground truth (for 'train' or 'valid' stages).

        RandLaNet:
            results (Tensor(B, N, C)): Prediction scores for all classes.
            input_data (Tuple): Batch of pointclouds and labels.
                input_data[0] (Tensor(B,N,3), float) : points
                input_data[-1] (Tensor(B,N), int) : labels

        SparseConvUNet:
            results (Tensor(SN, C)): Prediction scores for all classes. SN is
                total points in the batch.
            input_data (Dict): Batch of pointclouds and labels. Keys should be:
                'point' [Tensor(SN,3), float]: Concatenated points.
                'batch_lengths' [Tensor(B,), int]: Number of points in each
                    point cloud of the batch.
                'label' [Tensor(SN,) (optional)]: Concatenated labels.

        Returns:
            [Dict] visualizations of inputs and outputs suitable to save as an
                Open3D for TensorBoard summary.
        """
        if not hasattr(self, "_first_step"):
            self._first_step = epoch
        label_to_names = self.dataset.get_label_to_names()
        cfg = self.cfg.get('summary')
        max_pts = cfg.get('max_pts')
        if max_pts is None:
            max_pts = np.iinfo(np.int32).max
        use_reference = cfg.get('use_reference', False)
        max_outputs = cfg.get('max_outputs', 1)
        input_pcd = []
        gt_labels = []
        predict_labels = []

        def to_sum_fmt(tensor, add_dims=(0, 0), dtype=np.int32):
            np_tensor = tensor.numpy().astype(dtype)
            new_shape = (1,) * add_dims[0] + np_tensor.shape + (
                1,) * add_dims[1]
            return np_tensor.reshape(new_shape)

        # Variable size point clouds
        if self.model.cfg['name'] in ('KPFCNN', 'KPConv'):
            org_inputs = self.model.organise_inputs(input_data)
            predict_labels.append(to_sum_fmt(tf.argmax(results, 1), (0, 1)))
            if self._first_step == epoch or not use_reference:
                pointcloud = org_inputs['points'][0]
                input_pcd.append(
                    to_sum_fmt(pointcloud[:, :3], (0, 0), np.float32))
                if save_gt:
                    gtl = org_inputs['point_labels']
                    gt_labels.append(to_sum_fmt(gtl, (0, 1)))

        # Dict input, variable size point clouds
        elif self.model.cfg['name'] in ('SparseConvUnet', 'PointTransformer'):
            if self.model.cfg['name'] == 'SparseConvUnet':
                row_splits = np.hstack(
                    ((0,), np.cumsum(input_data['batch_lengths'])))
            else:
                row_splits = input_data.get('row_splits',
                                            [0, input_data['point'].shape[0]])
            max_outputs = min(max_outputs, len(row_splits) - 1)
            for k in range(max_outputs):
                blen_k = row_splits[k + 1] - row_splits[k]
                pcd_step = int(np.ceil(blen_k / min(max_pts, blen_k)))
                res_pcd = results[row_splits[k]:row_splits[k + 1]:pcd_step, :]
                predict_labels.append(to_sum_fmt(tf.argmax(res_pcd, 1), (0, 1)))
                if self._first_step != epoch and use_reference:
                    continue
                pointcloud = input_data['point'][
                    row_splits[k]:row_splits[k + 1]:pcd_step]
                pointcloud = tf.convert_to_tensor(pointcloud)
                input_pcd.append(
                    to_sum_fmt(pointcloud[:, :3], (0, 0), np.float32))
                if save_gt:
                    gtl = input_data['label'][
                        row_splits[k]:row_splits[k + 1]:pcd_step]
                    gt_labels.append(to_sum_fmt(gtl, (0, 1)))

        # Fixed size point clouds
        # Tuple input, same size point clouds
        elif self.model.cfg['name'] in ('RandLANet', 'PVCNN'):
            if self.model.cfg['name'] == 'RandLANet':
                if save_gt:
                    pointcloud = input_data[0]  # 0 => input to first layer
                else:  # Structured data during inference
                    pointcloud = tf.expand_dims(
                        tf.convert_to_tensor(input_data['point']), 0)
                    results = tf.expand_dims(results, 0)
            elif self.model.cfg['name'] == 'PVCNN':
                pointcloud = input_data['point']
            pcd_step = int(
                np.ceil(pointcloud.shape[1] /
                        min(max_pts, pointcloud.shape[1])))
            predict_labels = to_sum_fmt(
                tf.argmax(results[:max_outputs, ::pcd_step, :], -1), (0, 1))
            if self._first_step == epoch or not use_reference:
                input_pcd = to_sum_fmt(pointcloud[:max_outputs, ::pcd_step, :3],
                                       (0, 0), np.float32)
                if save_gt:
                    if self.model.cfg['name'] == 'RandLANet':
                        gtl = input_data[-1]
                        if 'int' not in repr(gtl.dtype).lower():
                            raise TypeError(
                                "Labels should be Int types. Received "
                                f"{gtl.dtype}")
                    elif self.model.cfg['name'] == 'PVCNN':
                        gtl = input_data['label']
                    gt_labels = to_sum_fmt(gtl[:max_outputs, ::pcd_step],
                                           (0, 1))
        else:
            raise NotImplementedError(
                "Saving 3D summary for the model "
                f"{self.model.cfg['name']} is not implemented.")

        def get_reference_or(data_tensor):
            if self._first_step == epoch or not use_reference:
                return data_tensor
            return self._first_step

        summary_dict = {
            'semantic_segmentation': {
                "vertex_positions": get_reference_or(input_pcd),
                "vertex_gt_labels": get_reference_or(gt_labels),
                "vertex_predict_labels": predict_labels,
                'label_to_names': label_to_names
            }
        }
        return summary_dict

    def save_logs(self, writer, epoch):
        """Save logs from the training and send results to TensorBoard."""
        train_accs = self.metric_train.acc()
        val_accs = self.metric_val.acc()

        train_ious = self.metric_train.iou()
        val_ious = self.metric_val.iou()

        loss_dict = {
            'Training loss': np.mean(self.losses),
            'Validation loss': np.mean(self.valid_losses)
        }
        acc_dicts = [{
            'Training accuracy': acc,
            'Validation accuracy': val_acc
        } for acc, val_acc in zip(train_accs, val_accs)]
        iou_dicts = [{
            'Training IoU': iou,
            'Validation IoU': val_iou
        } for iou, val_iou in zip(train_ious, val_ious)]

        log.info(f"loss train: {loss_dict['Training loss']:.3f} "
                 f" eval: {loss_dict['Validation loss']:.3f}")
        log.info(f"acc train: {acc_dicts[-1]['Training accuracy']:.3f} "
                 f" eval: {acc_dicts[-1]['Validation accuracy']:.3f}")
        log.info(f"iou train: {iou_dicts[-1]['Training IoU']:.3f} "
                 f" eval: {iou_dicts[-1]['Validation IoU']:.3f}")

        # send results to tensorboard
        with writer.as_default():
            for key, val in loss_dict.items():
                tf.summary.scalar(key, val, epoch)

            for key, val in acc_dicts[-1].items():
                tf.summary.scalar("{}/ Overall".format(key), val, epoch)
            for key, val in iou_dicts[-1].items():
                tf.summary.scalar("{}/ Overall".format(key), val, epoch)

            for stage in self.summary:
                for key, summary_dict in self.summary[stage].items():
                    label_to_names = summary_dict.pop('label_to_names', None)
                    summary3d.add_3d('/'.join((stage, key)),
                                     summary_dict,
                                     epoch,
                                     max_outputs=0,
                                     label_to_names=label_to_names,
                                     logdir=self.tensorboard_dir)

    def load_ckpt(self, ckpt_path=None, is_resume=True):
        """Load a checkpoint. You must pass the checkpoint and indicate if you want
        to resume.
        """
        train_ckpt_dir = join(self.cfg.logs_dir, 'checkpoint')
        make_dir(train_ckpt_dir)

        if hasattr(self, 'optimizer'):
            self.ckpt = tf.train.Checkpoint(step=tf.Variable(1),
                                            optimizer=self.optimizer,
                                            net=self.model)
        else:
            self.ckpt = tf.train.Checkpoint(step=tf.Variable(1), net=self.model)

        self.manager = tf.train.CheckpointManager(self.ckpt,
                                                  train_ckpt_dir,
                                                  max_to_keep=100)

        if ckpt_path is not None:
            self.ckpt.restore(ckpt_path).expect_partial()
            log.info("Restored from {}".format(ckpt_path))
        else:
            self.ckpt.restore(self.manager.latest_checkpoint)

            if self.manager.latest_checkpoint and is_resume:
                log.info("Restored from {}".format(
                    self.manager.latest_checkpoint))
            else:
                log.info("Initializing from scratch.")

    def save_ckpt(self, epoch):
        """Save a checkpoint at the passed epoch."""
        save_path = self.manager.save()
        log.info("Saved checkpoint at: {}".format(save_path))

    def save_config(self, writer):
        """Save experiment configuration with tensorboard summary."""
        if hasattr(self, 'cfg_tb'):
            with writer.as_default():
                with tf.name_scope("Description"):
                    tf.summary.text("Open3D-ML", self.cfg_tb['readme'], step=0)
                    tf.summary.text("Command line",
                                    self.cfg_tb['cmd_line'],
                                    step=0)
                with tf.name_scope("Configuration"):
                    tf.summary.text('Dataset',
                                    code2md(self.cfg_tb['dataset'],
                                            language='json'),
                                    step=0)
                    tf.summary.text('Model',
                                    code2md(self.cfg_tb['model'],
                                            language='json'),
                                    step=0)
                    tf.summary.text('Pipeline',
                                    code2md(self.cfg_tb['pipeline'],
                                            language='json'),
                                    step=0)


PIPELINE._register_module(SemanticSegmentation, "tf")

This class allows you to perform semantic segmentation for both training     and inference using the TensorFlow framework. This pipeline has multiple     stages: Pre-processing, loading dataset, testing, and inference or training.      **Example:**         This example loads the Semantic Segmentation and performs a training         using the SemanticKITTI dataset.          .. code::              import tensorflow as tf             from .base_pipeline import BasePipeline              Mydataset = TFDataloader(dataset=tf.dataset.get_split('training')             MyModel = SemanticSegmentation(self,model,dataset=Mydataset,                 name='SemanticSegmentation',                 batch_size=4,                 val_batch_size=4,                 test_batch_size=3,                 max_epoch=100,                 learning_rate=1e-2,                 lr_decays=0.95,                 save_ckpt_freq=20,                 adam_lr=1e-2,                 scheduler_gamma=0.95,                 momentum=0.98,                 main_log_dir='./logs/',                 device='gpu',                 split='train',                 train_sum_dir='train_log')      **Args:**             dataset: The 3D ML dataset class. You can use the base dataset,                 sample datasets, or a custom dataset.             model: The model to be used for building the pipeline.             name: The name of the current training.             batch_size: The batch size to be used for training.             val_batch_size: The batch size to be used for validation.             test_batch_size: The batch size to be used for testing.             max_epoch: The maximum size of the epoch to be used for training.             leanring_rate: The hyperparameter that controls the weights during                 training. Also, known as step size.             lr_decays: The learning rate decay for the training.             save_ckpt_freq: The frequency in which the checkpoint should be                 saved.             adam_lr: The leanring rate to be applied for Adam optimization.             scheduler_gamma: The decaying factor associated with the scheduler.             momentum: The momentum that accelerates the training rate schedule.             main_log_dir: The directory where logs are stored.             device: The device to be used for training.             split: The dataset split to be used. In this example, we have used                 "train".             train_sum_dir: The directory where the trainig summary is stored.      **Returns:**             class: The corresponding class. # maximum epoch during training # initial learning rate Run the inference using the data passed. Run the test using the data passed. # Save only for the first batch Run model training. # --------------------- training # --------------------- validation Create visualization for network inputs and outputs.          Args:             results: Model output (see below).             input_data: Model input (see below).             epoch (int): step             save_gt (bool): Save ground truth (for 'train' or 'valid' stages).          RandLaNet:             results (Tensor(B, N, C)): Prediction scores for all classes.             input_data (Tuple): Batch of pointclouds and labels.                 input_data[0] (Tensor(B,N,3), float) : points                 input_data[-1] (Tensor(B,N), int) : labels          SparseConvUNet:             results (Tensor(SN, C)): Prediction scores for all classes. SN is                 total points in the batch.             input_data (Dict): Batch of pointclouds and labels. Keys should be:                 'point' [Tensor(SN,3), float]: Concatenated points.                 'batch_lengths' [Tensor(B,), int]: Number of points in each                     point cloud of the batch.                 'label' [Tensor(SN,) (optional)]: Concatenated labels.          Returns:             [Dict] visualizations of inputs and outputs suitable to save as an                 Open3D for TensorBoard summary. # Variable size point clouds # Dict input, variable size point clouds # Fixed size point clouds # Tuple input, same size point clouds # 0 => input to first layer # Structured data during inference Save logs from the training and send results to TensorBoard. # send results to tensorboard Load a checkpoint. You must pass the checkpoint and indicate if you want         to resume. Save a checkpoint at the passed epoch. Save experiment configuration with tensorboard summary.