import argparse
import os
import warnings
from pathlib import Path

import mmcv
import numpy as np
from mmcv import Config

from mmseg.datasets.builder import build_dataset


def parse_args():
    parser = argparse.ArgumentParser(description='Browse a dataset')
    parser.add_argument('config', help='train config file path')
    parser.add_argument(
        '--show-origin',
        default=False,
        action='store_true',
        help='if True, omit all augmentation in pipeline,'
        ' show origin image and seg map')
    parser.add_argument(
        '--skip-type',
        type=str,
        nargs='+',
        default=['DefaultFormatBundle', 'Normalize', 'Collect'],
        help='skip some useless pipeline，if `show-origin` is true, '
        'all pipeline except `Load` will be skipped')
    parser.add_argument(
        '--output-dir',
        default='./output',
        type=str,
        help='If there is no display interface, you can save it')
    parser.add_argument('--show', default=False, action='store_true')
    parser.add_argument(
        '--show-interval',
        type=int,
        default=999,
        help='the interval of show (ms)')
    parser.add_argument(
        '--opacity',
        type=float,
        default=0.5,
        help='the opacity of semantic map')
    args = parser.parse_args()
    return args


def imshow_semantic(img,
                    seg,
                    class_names,
                    palette=None,
                    win_name='',
                    show=False,
                    wait_time=0,
                    out_file=None,
                    opacity=0.5):
    """Draw `result` over `img`.

    Args:
        img (str or Tensor): The image to be displayed.
        seg (Tensor): The semantic segmentation results to draw over
            `img`.
        class_names (list[str]): Names of each classes.
        palette (list[list[int]]] | np.ndarray | None): The palette of
            segmentation map. If None is given, random palette will be
            generated. Default: None
        win_name (str): The window name.
        wait_time (int): Value of waitKey param.
            Default: 0.
        show (bool): Whether to show the image.
            Default: False.
        out_file (str or None): The filename to write the image.
            Default: None.
        opacity(float): Opacity of painted segmentation map.
            Default 0.5.
            Must be in (0, 1] range.
    Returns:
        img (Tensor): Only if not `show` or `out_file`
    """
    img = mmcv.imread(img)
    img = img.copy()
    if palette is None:
        palette = np.random.randint(0, 255, size=(len(class_names), 3))
    palette = np.array(palette)
    assert palette.shape[0] == len(class_names)
    assert palette.shape[1] == 3
    assert len(palette.shape) == 2
    assert 0 < opacity <= 1.0
    color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8)
    for label, color in enumerate(palette):
        color_seg[seg == label, :] = color
    # convert to BGR
    color_seg = color_seg[..., ::-1]

    img = img * (1 - opacity) + color_seg * opacity
    img = img.astype(np.uint8)
    # if out_file specified, do not show image in window
    if out_file is not None:
        show = False

    if show:
        mmcv.imshow(img, win_name, wait_time)
    if out_file is not None:
        mmcv.imwrite(img, out_file)

    if not (show or out_file):
        warnings.warn('show==False and out_file is not specified, only '
                      'result image will be returned')
        return img


def _retrieve_data_cfg(_data_cfg, skip_type, show_origin):
    if show_origin is True:
        # only keep pipeline of Loading data and ann
        _data_cfg['pipeline'] = [
            x for x in _data_cfg.pipeline if 'Load' in x['type']
        ]
    else:
        _data_cfg['pipeline'] = [
            x for x in _data_cfg.pipeline if x['type'] not in skip_type
        ]


def retrieve_data_cfg(config_path, skip_type, show_origin=False):
    cfg = Config.fromfile(config_path)
    train_data_cfg = cfg.data.train
    if isinstance(train_data_cfg, list):
        for _data_cfg in train_data_cfg:
            if 'pipeline' in _data_cfg:
                _retrieve_data_cfg(_data_cfg, skip_type, show_origin)
            elif 'dataset' in _data_cfg:
                _retrieve_data_cfg(_data_cfg['dataset'], skip_type,
                                   show_origin)
            else:
                raise ValueError
    elif 'dataset' in train_data_cfg:
        _retrieve_data_cfg(train_data_cfg['dataset'], skip_type, show_origin)
    else:
        _retrieve_data_cfg(train_data_cfg, skip_type, show_origin)
    return cfg


def main():
    args = parse_args()
    cfg = retrieve_data_cfg(args.config, args.skip_type, args.show_origin)
    dataset = build_dataset(cfg.data.train)
    progress_bar = mmcv.ProgressBar(len(dataset))
    for item in dataset:
        filename = os.path.join(args.output_dir,
                                Path(item['filename']).name
                                ) if args.output_dir is not None else None
        imshow_semantic(
            item['img'],
            item['gt_semantic_seg'],
            dataset.CLASSES,
            dataset.PALETTE,
            show=args.show,
            wait_time=args.show_interval,
            out_file=filename,
            opacity=args.opacity,
        )
        progress_bar.update()


if __name__ == '__main__':
    main()