import os
from copy import deepcopy
import json
import re
from _ctypes import PyObj_FromPtr

import mmcv
import mmdet.core


class NoIndent(object):
    """ Value wrapper. """
    def __init__(self, value):
        self.value = value


class MyEncoder(json.JSONEncoder):
    FORMAT_SPEC = '@@{}@@'
    regex = re.compile(FORMAT_SPEC.format(r'(\d+)'))

    def __init__(self, **kwargs):
        # Save copy of any keyword argument values needed for use here.
        self.__sort_keys = kwargs.get('sort_keys', None)
        super(MyEncoder, self).__init__(**kwargs)

    def default(self, obj):
        return (self.FORMAT_SPEC.format(id(obj)) if isinstance(obj, NoIndent)
                else super(MyEncoder, self).default(obj))

    def encode(self, obj):
        format_spec = self.FORMAT_SPEC  # Local var to expedite access.
        json_repr = super(MyEncoder, self).encode(obj)  # Default JSON.

        # Replace any marked-up object ids in the JSON repr with the
        # value returned from the json.dumps() of the corresponding
        # wrapped Python object.
        for match in self.regex.finditer(json_repr):
            # see https://stackoverflow.com/a/15012814/355230
            id = int(match.group(1))
            no_indent = PyObj_FromPtr(id)
            json_obj_repr = json.dumps(
                no_indent.value, sort_keys=self.__sort_keys
            )

            # Replace the matched id string with json formatted representation
            # of the corresponding Python object.
            json_repr = json_repr.replace(
                '"{}"'.format(format_spec.format(id)), json_obj_repr)

        return json_repr


DAG_template = {
    "comment": "DAG defined, use DAG generate data flow",
    "green_mode": True,
    "model_list": [
        {
            "model_runner": "",
            "model_id": 244,
            "model_init_params_file": ""
        }
    ],
    "DAG": {
        "DAG0": [
            NoIndent([["start_0"], ["0_0"]]),
            NoIndent([["start_0", "0_0"], [None]])
        ]
    },
    "report_params": {
        "0_0": {
            "dbkey": "bbox",
            "conf_thres":  0.3
        }
    }
}


init_params_template = {
    "model_path": "",
    "num_classes": 80,
    "remapping_type": "COCO",
    "input_shape": [640, 640],
    "conf_threshold": 0.3,
    "iou_threshold": 0.5,
    "top_k_num": -1,
    "post_process_type": "mm"
}


def main(cfg_path, onnx_path):
    out_init_params = deepcopy(init_params_template)
    cfg = mmcv.Config.fromfile(cfg_path)
    w, h = cfg.test_pipeline[1].img_scale
    dataset = cfg.train_dataset.dataset.type[:-7].upper()
    num_classes = len(mmdet.core.get_classes(dataset.lower()))
    out_init_params['num_classes'] = num_classes
    out_init_params['input_shape'] = NoIndent([h, w])
    out_init_params['iou_threshold'] = cfg.model.test_cfg.nms.iou_threshold
    out_init_params['model_path'] = os.path.abspath(onnx_path)
    out_init_params['remapping_type'] = dataset
    out_init_params_path = os.path.splitext(cfg_path)[0] + "_init_params.json"

    s = json.dumps(out_init_params, indent=4, cls=MyEncoder)
    with open(out_init_params_path, 'w') as f:
        f.write(s)

    out_DAG = deepcopy(DAG_template)
    out_DAG['model_list'][0]['model_runner'] = cfg.model.type.lower()
    out_DAG['model_list'][0][
        'model_init_params_file'
    ] = os.path.abspath(out_init_params_path)
    out_DAG_path = os.path.splitext(cfg_path)[0] + "_DAG.json"

    s = json.dumps(out_DAG, indent=4, cls=MyEncoder)
    with open(out_DAG_path, 'w') as f:
        f.write(s)


if __name__ == '__main__':
    import sys
    main(sys.argv[1], sys.argv[2])