# Copyright (c) OpenMMLab. All rights reserved.
import argparse
import json
import os.path as osp

import cv2
import numpy as np
from mmcv.image import imfrombytes, imwrite
from mmengine.fileio import get
from mmengine.structures import InstanceData
from mmengine.utils import mkdir_or_exist

from mmdet.structures import DetDataSample
from mmdet.visualization import DetLocalVisualizer
from mmdet.visualization.palette import _get_adaptive_scales

# backend_args = dict(
#     backend='petrel',
#     path_mapping=dict({
#         './data/': 's3://openmmlab/datasets/detection/',
#         'data/': 's3://openmmlab/datasets/detection/'
#     }))
backend_args = None


def parse_args():
    parser = argparse.ArgumentParser(description='Browse a dataset')
    parser.add_argument('data_root')
    parser.add_argument('ann_file')
    parser.add_argument('img_prefix')
    parser.add_argument('--label-map-file', '-m', default=None)
    parser.add_argument(
        '--output-dir',
        '-o',
        default=None,
        type=str,
        help='If there is no display interface, you can save it')
    parser.add_argument('--not-show', default=False, action='store_true')
    parser.add_argument('--show-num', '-n', type=int, default=30)
    parser.add_argument('--shuffle', default=False, action='store_true')
    parser.add_argument(
        '--show-interval',
        type=float,
        default=0,
        help='the interval of show (s)')
    args = parser.parse_args()
    return args


def draw_all_character(visualizer, characters, w):
    start_index = 2
    y_index = 5
    for char in characters:
        if isinstance(char, str):
            visualizer.draw_texts(
                str(char),
                positions=np.array([start_index, y_index]),
                colors=(0, 0, 0),
                font_families='monospace')
            start_index += len(char) * 8
        else:
            visualizer.draw_texts(
                str(char[0]),
                positions=np.array([start_index, y_index]),
                colors=char[1],
                font_families='monospace')
            start_index += len(char[0]) * 8

        if start_index > w - 10:
            start_index = 2
            y_index += 15

    drawn_text = visualizer.get_image()
    return drawn_text


def main():
    args = parse_args()
    assert args.show_num > 0

    local_path = osp.join(args.data_root, args.ann_file)
    with open(local_path, 'r') as f:
        data_list = [json.loads(line) for line in f]

    dataset_index = list(range(len(data_list)))
    if args.shuffle:
        import random
        random.shuffle(dataset_index)

    if args.label_map_file is not None:
        label_map_file = osp.join(args.data_root, args.label_map_file)
        with open(label_map_file, 'r') as file:
            label_map = json.load(file)

    visualizer = DetLocalVisualizer()

    for i in dataset_index[:args.show_num]:
        item = data_list[i]

        img_path = osp.join(args.data_root, args.img_prefix, item['filename'])
        if backend_args is not None:
            img_bytes = get(img_path, backend_args)
            img = imfrombytes(img_bytes, flag='color')
        else:
            img = cv2.imread(img_path)
        img = img[..., [2, 1, 0]]  # bgr to rgb

        base_name, extension = osp.splitext(item['filename'])

        out_file = osp.join(args.output_dir, base_name + '_' + str(i) +
                            extension) if args.output_dir is not None else None

        if args.output_dir is not None:
            mkdir_or_exist(args.output_dir)

        if 'detection' in item:
            anno = item['detection']

            instances = [obj for obj in anno['instances']]
            bboxes = [obj['bbox'] for obj in instances]
            bbox_labels = [int(obj['label']) for obj in instances]
            label_names = [label_map[str(label)] for label in bbox_labels]

            data_sample = DetDataSample()
            gt_instances = InstanceData()
            if len(instances) > 0 and 'score' in instances[0]:
                score = [obj['score'] for obj in instances]
                gt_instances['scores'] = np.array(score)

            gt_instances['bboxes'] = np.array(bboxes).reshape(-1, 4)
            gt_instances['labels'] = np.array(bbox_labels)
            gt_instances['label_names'] = label_names
            data_sample.gt_instances = gt_instances

            visualizer.add_datasample(
                osp.basename(img_path),
                img,
                data_sample,
                draw_pred=False,
                show=not args.not_show,
                wait_time=args.show_interval,
                out_file=out_file)
        elif 'grounding' in item:
            anno = item['grounding']
            text = anno['caption']
            regions = anno['regions']

            max_label = len(regions) if len(regions) > 0 else 0
            palette = np.random.randint(0, 256, size=(max_label + 1, 3))
            bbox_palette = [tuple(c) for c in palette]
            # bbox_palette = get_palette('random', max_label + 1)
            colors = [bbox_palette[label] for label in range(max_label)]

            visualizer.set_image(img)

            gt_tokens_positive = []
            for i, region in enumerate(regions):
                bbox = region['bbox']
                bbox = np.array(bbox).reshape(-1, 4)
                tokens_positive = region['tokens_positive']
                gt_tokens_positive.append(tokens_positive)
                visualizer.draw_bboxes(
                    bbox,
                    edge_colors=colors[i],
                    face_colors=colors[i],
                    alpha=0.3)
                visualizer.draw_bboxes(bbox, edge_colors=colors[i], alpha=1)

                if 'score' in region:
                    areas = (bbox[:, 3] - bbox[:, 1]) * (
                        bbox[:, 2] - bbox[:, 0])
                    scales = _get_adaptive_scales(areas)
                    score = region['score'][0]
                    score = [str(s) for s in score]
                    font_sizes = [
                        int(13 * scales[i]) for i in range(len(scales))
                    ]
                    visualizer.draw_texts(
                        score,
                        bbox[:, :2].astype(np.int32),
                        colors=(255, 255, 255),
                        font_sizes=font_sizes,
                        bboxes=[{
                            'facecolor': 'black',
                            'alpha': 0.8,
                            'pad': 0.7,
                            'edgecolor': 'none'
                        }] * len(bbox))

            drawn_img = visualizer.get_image()
            new_image = np.ones((100, img.shape[1], 3), dtype=np.uint8) * 255
            visualizer.set_image(new_image)

            split_by_character = [char for char in text]
            characters = []
            start_index = 0
            end_index = 0
            for w in split_by_character:
                end_index += len(w)
                is_find = False
                for i, positive in enumerate(gt_tokens_positive):
                    for p in positive:
                        if start_index >= p[0] and end_index <= p[1]:
                            characters.append([w, colors[i]])
                            is_find = True
                            break
                    if is_find:
                        break
                if not is_find:
                    characters.append([w, (0, 0, 0)])
                start_index = end_index

            drawn_text = draw_all_character(visualizer, characters,
                                            img.shape[1])
            drawn_img = np.concatenate((drawn_img, drawn_text), axis=0)

            if not args.not_show:
                visualizer.show(
                    drawn_img,
                    win_name=base_name,
                    wait_time=args.show_interval)

            if out_file is not None:
                imwrite(drawn_img[..., ::-1], out_file)

        elif 'referring' in item:
            referring = item['referring']

            max_label = len(referring) if len(referring) > 0 else 0
            palette = np.random.randint(0, 256, size=(max_label + 1, 3))
            bbox_palette = [tuple(c) for c in palette]
            # bbox_palette = get_palette('random', max_label + 1)
            colors = [bbox_palette[label] for label in range(max_label)]

            visualizer.set_image(img)
            phrases = []
            for i, ref in enumerate(referring):
                bbox = ref['bbox']
                phrase = ref['phrase']
                phrases.append(' // '.join(phrase))
                bbox = np.array(bbox).reshape(-1, 4)

                visualizer.draw_bboxes(
                    bbox,
                    edge_colors=colors[i],
                    face_colors=colors[i],
                    alpha=0.3)
                visualizer.draw_bboxes(bbox, edge_colors=colors[i], alpha=1)
            drawn_img = visualizer.get_image()

            new_image = np.ones((100, img.shape[1], 3), dtype=np.uint8) * 255
            visualizer.set_image(new_image)

            start_index = 2
            y_index = 5

            chunk_size = max(min(img.shape[1] - 400, 70), 50)
            for i, p in enumerate(phrases):
                chunk_p = [
                    p[i:i + chunk_size] for i in range(0, len(p), chunk_size)
                ]
                for cp in chunk_p:
                    visualizer.draw_texts(
                        cp,
                        positions=np.array([start_index, y_index]),
                        colors=colors[i],
                        font_families='monospace')
                    y_index += 15

            drawn_text = visualizer.get_image()
            drawn_img = np.concatenate((drawn_img, drawn_text), axis=0)

            if not args.not_show:
                visualizer.show(
                    drawn_img,
                    win_name=base_name,
                    wait_time=args.show_interval)

            if out_file is not None:
                imwrite(drawn_img[..., ::-1], out_file)


if __name__ == '__main__':
    main()