## @package net_printer
# Module caffe2.python.net_printer


from caffe2.proto.caffe2_pb2 import OperatorDef, NetDef
from caffe2.python.checkpoint import Job
from caffe2.python.core import Net, ExecutionStep, Plan
from caffe2.python.task import Task, TaskGroup, WorkspaceType, TaskOutput
from collections import defaultdict
from contextlib import contextmanager
from copy import copy
from itertools import chain


class Visitor:
    @classmethod
    def register(cls, Type):
        if not(hasattr(cls, 'visitors')):
            cls.visitors = {}
        else:
            assert Type not in cls.visitors, \
                '{} already registered!'.format(Type)

        def _register(func):
            cls.visitors[Type] = func
            return func

        return _register

    def __call__(self, obj, *args, **kwargs):
        if obj is None:
            return

        Type = type(obj)
        if Type not in self.__class__.visitors:
            raise TypeError('%s: unsupported object type: %s' % (
                self.__class__.__name__, Type))

        func = self.__class__.visitors[Type]
        return func(self, obj, *args, **kwargs)


class Analyzer(Visitor):
    PREFIXES_TO_IGNORE = {'distributed_ctx_init'}

    def __init__(self):
        self.workspaces = defaultdict(lambda: defaultdict(lambda: 0))
        self.workspace_ctx = []

    @property
    def workspace(self):
        return self.workspace_ctx[-1]

    @contextmanager
    def set_workspace(self, node=None, ws=None, do_copy=False):
        if ws is not None:
            ws = ws
        elif node is not None:
            ws = self.workspaces[str(node)]
        else:
            ws = self.workspace
        if do_copy:
            ws = copy(ws)
        self.workspace_ctx.append(ws)
        try:
            yield ws
        finally:
            del self.workspace_ctx[-1]

    def define_blob(self, blob):
        self.workspace[blob] += 1

    def need_blob(self, blob):
        if any(blob.startswith(p) for p in Analyzer.PREFIXES_TO_IGNORE):
            return
        assert blob in self.workspace, 'Blob undefined: %s' % blob


@Analyzer.register(OperatorDef)
def analyze_op(analyzer, op):
    for x in op.input:
        analyzer.need_blob(x)
    for x in op.output:
        analyzer.define_blob(x)


@Analyzer.register(Net)
def analyze_net(analyzer, net):
    for x in net.Proto().op:
        analyzer(x)


@Analyzer.register(ExecutionStep)
def analyze_step(analyzer, step):
    proto = step.Proto()
    with analyzer.set_workspace(do_copy=proto.create_workspace):
        if proto.report_net:
            with analyzer.set_workspace(do_copy=True):
                analyzer(step.get_net(proto.report_net))
        all_new_blobs = set()
        substeps = step.Substeps() + [step.get_net(n) for n in proto.network]
        for substep in substeps:
            with analyzer.set_workspace(
                    do_copy=proto.concurrent_substeps) as ws_in:
                analyzer(substep)
                if proto.should_stop_blob:
                    analyzer.need_blob(proto.should_stop_blob)
            if proto.concurrent_substeps:
                new_blobs = set(ws_in.keys()) - set(analyzer.workspace.keys())
                assert len(all_new_blobs & new_blobs) == 0, (
                    'Error: Blobs created by multiple parallel steps: %s' % (
                        ', '.join(all_new_blobs & new_blobs)))
                all_new_blobs |= new_blobs
    for x in all_new_blobs:
        analyzer.define_blob(x)


@Analyzer.register(Task)
def analyze_task(analyzer, task):
    # check that our plan protobuf is not too large (limit of 64Mb)
    step = task.get_step()
    plan = Plan(task.node)
    plan.AddStep(step)
    proto_len = len(plan.Proto().SerializeToString())
    assert proto_len < 2 ** 26, (
        'Due to a protobuf limitation, serialized tasks must be smaller '
        'than 64Mb, but this task has {} bytes.' % proto_len)

    is_private = task.workspace_type() != WorkspaceType.GLOBAL
    with analyzer.set_workspace(do_copy=is_private):
        analyzer(step)


@Analyzer.register(TaskGroup)
def analyze_task_group(analyzer, tg):
    for task in tg.tasks_by_node().tasks():
        with analyzer.set_workspace(node=task.node):
            analyzer(task)


@Analyzer.register(Job)
def analyze_job(analyzer, job):
    analyzer(job.init_group)
    analyzer(job.epoch_group)


def analyze(obj):
    """
    Given a Job, visits all the execution steps making sure that:
      - no undefined blobs will be found during execution
      - no blob with same name is defined in concurrent steps
    """
    Analyzer()(obj)


class Text:
    def __init__(self):
        self._indent = 0
        self._lines_in_context = [0]
        self.lines = []

    @contextmanager
    def context(self, text):
        if text is not None:
            self.add('with %s:' % text)
            self._indent += 4
            self._lines_in_context.append(0)
        try:
            yield
        finally:
            if text is not None:
                if self._lines_in_context[-1] == 0:
                    self.add('pass')
                self._indent -= 4
                del self._lines_in_context[-1]

    def add(self, text):
        self._lines_in_context[-1] += 1
        self.lines.append((' ' * self._indent) + text)

    def __str__(self):
        return '\n'.join(self.lines)


class Printer(Visitor, Text):
    def __init__(self, factor_prefixes=False, c2_syntax=True):
        super(Visitor, self).__init__()
        super(Text, self).__init__()
        self.factor_prefixes = factor_prefixes
        self.c2_syntax = c2_syntax
        self.c2_net_name = None


def _sanitize_str(s):
    if isinstance(s, str):
        sanitized = s
    elif isinstance(s, bytes):
        sanitized = s.decode('ascii', errors='ignore')
    else:
        sanitized = str(s)
    if len(sanitized) < 64:
        return "'%s'" % sanitized
    else:
        return "'%s'" % sanitized[:64] + '...<+len=%d>' % (len(sanitized) - 64)


def _arg_val(arg):
    if arg.HasField('f'):
        return str(arg.f)
    if arg.HasField('i'):
        return str(arg.i)
    if arg.HasField('s'):
        return _sanitize_str(arg.s)
    if arg.floats:
        return str(list(arg.floats))
    if arg.ints:
        return str(list(arg.ints))
    if arg.strings:
        return str([_sanitize_str(s) for s in arg.strings])
    return '[]'


def commonprefix(m):
    "Given a list of strings, returns the longest common prefix"
    if not m:
        return ''
    s1 = min(m)
    s2 = max(m)
    for i, c in enumerate(s1):
        if c != s2[i]:
            return s1[:i]
    return s1


def format_value(val):
    if isinstance(val, list):
        return '[%s]' % ', '.join("'%s'" % str(v) for v in val)
    else:
        return str(val)


def factor_prefix(vals, do_it):
    vals = [format_value(v) for v in vals]
    prefix = commonprefix(vals) if len(vals) > 1 and do_it else ''
    joined = ', '.join(v[len(prefix):] for v in vals)
    return '%s[%s]' % (prefix, joined) if prefix else joined


def call(op, inputs=None, outputs=None, factor_prefixes=False):
    if not inputs:
        inputs = ''
    else:
        inputs_v = [a for a in inputs if not isinstance(a, tuple)]
        inputs_kv = [a for a in inputs if isinstance(a, tuple)]
        inputs = ', '.join(
            x
            for x in chain(
                [factor_prefix(inputs_v, factor_prefixes)],
                ('%s=%s' % kv for kv in inputs_kv),
            )
            if x
        )
    call = '%s(%s)' % (op, inputs)
    return call if not outputs else '%s = %s' % (
        factor_prefix(outputs, factor_prefixes), call)


def format_device_option(dev_opt):
    if not dev_opt or not (
            dev_opt.device_type or dev_opt.device_id or dev_opt.node_name):
        return None
    return call(
        'DeviceOption',
        [dev_opt.device_type, dev_opt.device_id, "'%s'" % dev_opt.node_name])


@Printer.register(OperatorDef)
def print_op(text, op):
    args = [(a.name, _arg_val(a)) for a in op.arg]
    dev_opt_txt = format_device_option(op.device_option)
    if dev_opt_txt:
        args.append(('device_option', dev_opt_txt))

    if text.c2_net_name:
        text.add(call(
            text.c2_net_name + '.' + op.type,
            [list(op.input), list(op.output)] + args))
    else:
        text.add(call(
            op.type,
            list(op.input) + args,
            op.output,
            factor_prefixes=text.factor_prefixes))
    for arg in op.arg:
        if arg.HasField('n'):
            with text.context('arg: %s' % arg.name):
                text(arg.n)


@Printer.register(NetDef)
def print_net_def(text, net_def):
    if text.c2_syntax:
        text.add(call('core.Net', ["'%s'" % net_def.name], [net_def.name]))
        text.c2_net_name = net_def.name
    else:
        text.add('# net: %s' % net_def.name)
    for op in net_def.op:
        text(op)
    if text.c2_syntax:
        text.c2_net_name = None


@Printer.register(Net)
def print_net(text, net):
    text(net.Proto())


def _get_step_context(step):
    proto = step.Proto()
    if proto.should_stop_blob:
        return call('loop'), False
    if proto.num_iter and proto.num_iter != 1:
        return call('loop', [proto.num_iter]), False
    if proto.num_concurrent_instances > 1:
        return (
            call('parallel',
                 [('num_instances', proto.num_concurrent_instances)]),
            len(step.Substeps()) > 1)
    concurrent = proto.concurrent_substeps and len(step.Substeps()) > 1
    if concurrent:
        return call('parallel'), True
    if proto.report_net:
        return call('run_once'), False
    return None, False


@Printer.register(ExecutionStep)
def print_step(text, step):
    proto = step.Proto()
    step_ctx, do_substep = _get_step_context(step)
    with text.context(step_ctx):
        if proto.report_net:
            with text.context(call('report_net', [proto.report_interval])):
                text(step.get_net(proto.report_net))
        substeps = step.Substeps() + [step.get_net(n) for n in proto.network]
        for substep in substeps:
            sub_proto = (
                substep.Proto() if isinstance(substep, ExecutionStep) else None)
            if sub_proto is not None and sub_proto.run_every_ms:
                substep_ctx = call(
                    'reporter',
                    [str(substep), ('interval_ms', sub_proto.run_every_ms)])
            elif do_substep:
                title = (
                    'workspace'
                    if sub_proto is not None and sub_proto.create_workspace else
                    'step')
                substep_ctx = call(title, [str(substep)])
            else:
                substep_ctx = None
            with text.context(substep_ctx):
                text(substep)
                if proto.should_stop_blob:
                    text.add(call('yield stop_if', [proto.should_stop_blob]))


def _print_task_output(x):
    assert isinstance(x, TaskOutput)
    return 'Output[' + ', '.join(str(x) for x in x.names) + ']'


@Printer.register(Task)
def print_task(text, task):
    outs = ', '.join(_print_task_output(o) for o in task.outputs())
    context = [('node', task.node), ('name', task.name), ('outputs', outs)]
    with text.context(call('Task', context)):
        text(task.get_step())


@Printer.register(TaskGroup)
def print_task_group(text, tg, header=None):
    with text.context(header or call('TaskGroup')):
        for task in tg.tasks_by_node().tasks():
            text(task)


@Printer.register(Job)
def print_job(text, job):
    text(job.init_group, 'Job.current().init_group')
    text(job.epoch_group, 'Job.current().epoch_group')
    with text.context('Job.current().stop_conditions'):
        for out in job.stop_conditions:
            text.add(_print_task_output(out))
    text(job.download_group, 'Job.current().download_group')
    text(job.exit_group, 'Job.current().exit_group')


def to_string(obj, **kwargs):
    """
    Given a Net, ExecutionStep, Task, TaskGroup or Job, produces a string
    with detailed description of the execution steps.
    """
    printer = Printer(**kwargs)
    printer(obj)
    return str(printer)


def debug_net(net):
    """
    Given a Net, produce another net that logs info about the operator call
    before each operator execution. Use for debugging purposes.
    """
    assert isinstance(net, Net)
    debug_net = Net(str(net))
    assert isinstance(net, Net)
    for op in net.Proto().op:
        text = Text()
        print_op(op, text)
        debug_net.LogInfo(str(text))
        debug_net.Proto().op.extend([op])
    return debug_net