# Copyright 2019 Vitaliy Zakaznikov (TestFlows Test Framework http://testflows.com)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
import sys
import time
import inspect
import functools
import tempfile
import threading
import textwrap

from contextlib import ExitStack, contextmanager

import testflows.settings as settings

from .exceptions import DummyTestException, ArgumentError, ResultException, RepeatTestException
from .flags import Flags, SKIP, TE, FAIL_NOT_COUNTED, ERROR_NOT_COUNTED, NULL_NOT_COUNTED
from .flags import CFLAGS, PAUSE_BEFORE, PAUSE_AFTER
from .testtype import TestType, TestSubType
from .objects import get, Null, OK, Fail, Skip, Error, Argument, Attribute, ExamplesTable
from .constants import name_sep, id_sep
from .io import TestIO, LogWriter
from .name import join, depth, match, absname
from .funcs import current, top, previous, main, skip, ok, fail, error, exception, pause, load
from .init import init
from .cli.arg.parser import ArgumentParser
from .cli.arg.exit import ExitWithError, ExitException
from .cli.arg.type import key_value as key_value_type
from .cli.text import danger, warning
from .exceptions import exception as get_exception
from .filters import the

class xfails(dict):
    """xfails container.

    xfails = {
        "pattern": [("result", "reason")],
        ...
        }
    """
    def add(self, pattern, *results):
        """Add an entry to the xfails.

        :param pattern: test name pattern to match
        :param *results: one or more results to cross out
            where each result is a two-tuple of (result, reason)
        """
        self[pattern] = results
        return self

class xflags(dict):
    """xflags container.

    xflags = {
        "filter": (set_flags, clear_flags),
        ...
    }
    """
    def add(self, pattern, set_flags=0, clear_flags=0):
        """Add an entry to the xflags.

        :param pattern: test name pattern to match
        :param set_flags: flags to set
        :param clear_flags: flags to clear, default: None
        """
        self[pattern] = [Flags(set_flags), Flags(clear_flags)]
        return self

class DummyTest(object):
    """Base class for dummy tests.
    """
    def __init__(self, *args, **kwargs):
        pass

    def __enter__(self):
        def dummy(*args, **kwargs):
            pass

        self.trace = sys.gettrace()
        sys.settrace(dummy)
        sys._getframe(1).f_trace = self.__skip__

    def __skip__(self, *args):
        sys.settrace(self.trace)
        raise DummyTestException()

    def __exit__(self, exception_type, exception_value, exception_traceback):
        if isinstance(exception_value, DummyTestException):
            return True

class TestBase(object):
    """Base class for all the tests.

    :param name: name
    :param flags: flags
    :param parent: parent name
    :param only: tests to run
    :param start: name of the starting test
    :param end: name of the last test
    """
    uid = None
    tags = set()
    attributes = []
    requirements = []
    users = []
    tickets = []
    examples = None
    name = None
    description = None
    node = None
    map = None
    flags = Flags()
    name_sep = "."
    type = TestType.Test
    subtype = TestSubType.Empty

    @classmethod
    def argparser(cls):
        """Command line argument parser.

        :return: argument parser
        """
        definitions_description = """

        Option values

        pattern
          used to match test names using a unix-like file path pattern that supports wildcards
            '/' path level separator
            '*' matches everything
            '?' matches any single character
            '[seq]' matches any character in seq
            '[!seq]' matches any character not in seq
            ':' matches anything at the current path level
          for example: "suiteA/*" selects all the tests in 'suiteA'
        """
        parser = ArgumentParser(
                prog=sys.argv[0],
                description=((cls.description or "") + definitions_description),
                description_prog="Test - Framework"
            )
        parser.add_argument("--name", dest="_name", metavar="name",
            help="test run name", type=str, required=False)
        parser.add_argument("--tag", dest="_tags", metavar="value", nargs="+",
            help="test run tags", type=str, required=False)
        parser.add_argument("--attr", dest="_attrs", metavar="name=value", nargs="+",
            help="test run attributes", type=key_value_type, required=False)
        parser.add_argument("--only", dest="_only", metavar="pattern", nargs="+",
            help="run only selected tests", type=str, required=False)
        parser.add_argument("--skip", dest="_skip", metavar="pattern", nargs="+",
            help="skip selected tests", type=str, required=False)
        parser.add_argument("--start", dest="_start", metavar="pattern", nargs=1,
            help="start at the selected test", type=str, required=False)
        parser.add_argument("--end", dest="_end", metavar="pattern", nargs=1,
            help="end at the selected test", type=str, required=False)
        parser.add_argument("--pause-before", dest="_pause_before", metavar="pattern", nargs="+",
            help="pause before executing selected tests", type=str, required=False)
        parser.add_argument("--pause-after", dest="_pause_after", metavar="pattern", nargs="+",
            help="pause after executing selected tests", type=str, required=False)
        parser.add_argument("--debug", dest="_debug", action="store_true",
            help="enable debugging mode", default=False)
        parser.add_argument("--no-colors", dest="_no_colors", action="store_true",
            help="disable terminal color highlighting", default=False)
        parser.add_argument("--id", metavar="id", dest="_id", type=str, help="custom test id")
        parser.add_argument("-o", "--output", dest="_output", metavar="format", type=str,
            choices=["slick", "nice", "quiet", "short", "dots", "raw", "silent"], default="nice",
            help="""stdout output format, choices are: ['slick','nice','short','dots','quiet','raw','silent'],
                default: 'nice'""")
        parser.add_argument("-l", "--log", dest="_log", metavar="file", type=str,
            help="path to the log file where test output will be stored, default: uses temporary log file")
        parser.add_argument("--show-skipped", dest="_show_skipped", action="store_true",
            help="show skipped tests, default: False", default=False)
        return parser

    def parse_cli_args(self, xflags=None, only=None, skip=None, start=None, end=None,
            name=None, tags=None, attributes=None):
        """Parse command line arguments.

        :return: parsed known arguments
        """
        try:
            parser = self.argparser()
            args, unknown = parser.parse_known_args()
            args = vars(args)

            if args.get("_name"):
                name = args.pop("_name")

            if name is None:
                raise TypeError("name must be specified")

            if args.get("_debug"):
                settings.debug = True
                args.pop("_debug")

            if args.get("_no_colors"):
                settings.no_colors = True
                args.pop("_no_colors")

            if unknown:
                raise ExitWithError(f"unknown argument {unknown}")

            if args.get("_id"):
                settings.test_id = args.get("_id")
                args.pop("_id")

            if args.get("_log"):
                logfile = os.path.abspath(args.get("_log"))
                settings.write_logfile = logfile
                args.pop("_log")
            else:
                settings.write_logfile = os.path.join(tempfile.gettempdir(), f"testflows.{os.getpid()}.log")

            settings.read_logfile = settings.write_logfile
            if os.path.exists(settings.write_logfile):
                os.remove(settings.write_logfile)

            settings.output_format = args.pop("_output")

            if args.get("_show_skipped"):
                settings.show_skipped = True
                args.pop("_show_skipped")

            if args.get("_pause_before"):
                if not xflags:
                    xflags = globals()["xflags"]()
                for pattern in args.get("_pause_before"):
                    pattern = absname(pattern, name)
                    xflags[pattern] = xflags.get(pattern, [0, 0])
                    xflags[pattern][0] |= PAUSE_BEFORE
                args.pop("_pause_before")

            if args.get("_pause_after"):
                if not xflags:
                    xflags = globals()["xflags"]()
                for pattern in args.get("_pause_after"):
                    pattern = absname(pattern, name)
                    xflags[pattern] = xflags.get(pattern, [0, 0])
                    xflags[pattern][0] |= PAUSE_AFTER
                args.pop("_pause_after")

            if args.get("_only"):
                only = [] # clear whatever was passed
                for pattern in args.get("_only"):
                    only.append(the(pattern).at(name))
                args.pop("_only")

            if args.get("_skip"):
                skip = [] # clear whatever was passed
                for pattern in args.get("_skip"):
                    only.append(the(pattern).at(name))
                args.pop("_skip")

            if args.get("_start"):
                start = the(args.get("_start")[0]).at(name)
                args.pop("_start")

            if args.get("_end"):
                end = the(args.get("_end")[0]).at(name)
                args.pop("_end")

            if args.get("_tags"):
                tags = [value for value in args.pop("_tags")]

            if args.get("_attrs"):
                attributes = [Attribute(item.key, item.value) for item in args.pop("_attrs")]

        except (ExitException, KeyboardInterrupt, Exception) as exc:
            #if settings.debug:
            sys.stderr.write(warning(get_exception(), eol='\n'))
            sys.stderr.write(danger("error: " + str(exc).strip()))
            if isinstance(exc, ExitException):
                sys.exit(exc.exitcode)
            else:
                sys.exit(1)

        return args, xflags, only, skip, start, end, name, tags, attributes

    def __init__(self, name=None, flags=None, cflags=None, type=None, subtype=None,
                 uid=None, tags=None, attributes=None, requirements=None,
                 users=None, tickets=None, examples=None, description=None, parent=None,
                 xfails=None, xflags=None, only=None, skip=None,
                 start=None, end=None, args=None, id=None, node=None, map=None,
                 _frame=None, _run=True):

        self.lock = threading.Lock()
        self._run = _run

        cli_args = {}
        if current() is None:
            if top() is not None:
                raise RuntimeError("only one top level test is allowed")
            top(self)
            # flag to indicate if main test called init
            self._init= False
            frame = get(_frame, inspect.currentframe().f_back.f_back.f_back)
            if main(frame):
                cli_args, xflags, only, skip, start, end, name, tags, attributes = self.parse_cli_args(
                    xflags, only, skip, start, end, name, tags, attributes)

        self.name = name
        if self.name is None:
            raise TypeError("name must be specified")
        self.child_count = 0
        self.start_time = time.time()
        self.parent = parent
        self.id = get(id, [settings.test_id])
        self.node = get(node, self.node)
        self.map = get(map, self.map)
        self.tags = tags
        self.requirements = get(requirements, self.requirements)
        self.attributes = get(attributes, self.attributes)
        self.users = get(users, self.users)
        self.tickets = get(tickets, self.tickets)
        self.description = get(description, self.description)
        self.examples = get(examples, get(self.examples, ExamplesTable()))
        self.args = get(args, {})
        self.args.update({k:v for k, v in cli_args.items() if not k.startswith("_")})
        self._process_args()
        self.result = Null(self.name)
        if flags is not None:
            self.flags = Flags(flags)
        self.type = get(type, self.type)
        self.subtype = get(subtype, self.subtype)
        self.cflags = Flags(cflags) | (self.flags & CFLAGS)
        self.uid = get(uid, self.uid)
        self.xfails = get(xfails, {})
        self.xflags = get(xflags, {})
        self.only = get(only, [])
        self.skip = get(skip, [])
        self.start = get(start, None)
        self.end = get(end, None)
        self.caller_test = None

    @classmethod
    def make_name(cls, name, parent=None):
        """Make full name.

        :param name: name
        :param parent: parent name
        """
        name = name % {"name": cls.name} if name is not None else cls.name
        if name is None:
            raise TypeError("name must be specified")
        # '/' is not allowed just like in Unix file names
        # so convert any '/' to U+2215 division slash
        name = name.replace(name_sep, "\u2215")
        return join(get(parent, name_sep), name)

    @classmethod
    def make_tags(cls, tags):
        return set(get(tags, cls.tags))

    def _process_args(self):
        """Process arguments by converting
        them into a dictionary of
        "name:Argument" pairs
        """
        args = []
        try:
            for name in dict(self.args):
                value = self.args.get(name)
                if not isinstance(value, Argument):
                    value = Argument(name=name, value=value)
                args.append(value)

            self.args = {}

            for arg in args:
                if not isinstance(arg, Argument):
                    raise ValueError(f"not an argument {arg}")
                self.args[arg.name] = arg
        except ArgumentError:
            raise
        except Exception as e:
            raise ArgumentError(str(e))

    def __enter__(self):
        self.io = TestIO(self)
        if top() is self:
            self.io.output.protocol()
            self.io.output.version()
        self.io.output.test_message()

        if self.flags & PAUSE_BEFORE:
            pause()

        self.caller_test = current()
        current(self)

        if self.flags & SKIP:
            raise ResultException(Skip(self.name, "skip flag set"))
        else:
            if top() is self:
                self._init = init()
            if self._run:
                self.run(**{name: arg.value for name, arg in self.args.items()})
            return self

    def __exit__(self, exception_type, exception_value, exception_traceback):
        if top() is self and not self._init:
            return False

        current(self.caller_test)
        previous(self)

        try:
            if exception_value:
                if isinstance(exception_value, ResultException):
                    self.result = exception_value.result
                elif isinstance(exception_value, AssertionError):
                    exception(test=self)
                    self.result = Fail(self.name, str(exception_value))
                else:
                    exception(test=self)
                    self.result = Error(self.name,
                        "unexpected %s: %s" % (exception_type.__name__, str(exception_value)))
                    if isinstance(exception_value, KeyboardInterrupt):
                        raise ResultException(self.result)
            else:
                if isinstance(self.result, Null):
                    self.result = OK(self.name)
        finally:
            self._apply_xfails()
            self.io.output.result(self.result)
            self.io.close()

            if self.flags & PAUSE_AFTER:
                pause()

        return True

    def _apply_xfails(self):
        """Apply xfails to self.result.
        """
        for pattern, xouts in self.xfails.items():
            if match(self.name, pattern):
                for xout in xouts:
                    result, reason = xout
                    if isinstance(self.result, result):
                        self.result = self.result.xout(reason)

    def run(self, **args):
        pass

    def bind(self, func):
        """Bind function to the current test.

        :param func: function that takes an instance of test
            as the argument named 'test'
        :return: partial function with the 'test' argument set to self
        """
        return functools.partial(func, test=self)

    def message_io(self, name=""):
        """Return an instance of test's message IO
        used to write messages using print() method
        or other methods that takes a file-like
        object.

        Note: only write() and flush() methods
        are supported.

        :param name: name of the stream, default: None
        """
        return self.io.message_io(name=name)

class _test(object):
    """Test definition.

    :param name: name of the test
    :param test: test class (optional), default: Test
    :param **kwargs: test class arguments
    """
    def __init__(self, name, **kwargs):
        current_thread = threading.current_thread()
        if getattr(current_thread, "_parent", None):
            parent_current = current(thread=current_thread._parent)
            parent_top = top(thread=current_thread._parent)
            parent_previous = previous(thread=current_thread._parent)
            if parent_current and not current():
                current(value=parent_current)
                top(value=parent_top)
                previous(value=parent_previous)

        parent = kwargs.pop("parent", None) or current()
        test = kwargs.pop("test", None)
        keep_type = kwargs.pop("keep_type", None)

        test = test if test is not None else TestBase

        if parent:
            name = test.make_name(name, parent.name)
            kwargs["parent"] = parent.name
            kwargs["id"] = parent.id + [parent.child_count]
            kwargs["cflags"] = parent.cflags
            # propagate xfails, xflags, only and skip that prefix match the name of the test
            kwargs["xfails"] = xfails({
                    k: v for k, v in parent.xfails.items() if match(name, k, prefix=True)
                }) or kwargs.get("xfails")
            kwargs["xflags"] = xflags({
                    k: v for k, v in parent.xflags.items() if match(name, k, prefix=True)
                }) or kwargs.get("xflags")
            kwargs["only"] = parent.only or kwargs.get("only")
            kwargs["skip"] = parent.skip or kwargs.get("skip")
            kwargs["start"] = parent.start or kwargs.get("start")
            kwargs["end"] = parent.end or kwargs.get("end")
            # handle parent test type propagation
            if keep_type is None:
                self._parent_type_propagation(parent, kwargs)
            with parent.lock:
                parent.child_count += 1
        else:
            name = test.make_name(name)
        tags = test.make_tags(kwargs.pop("tags", None))

        # anchor all patterns
        kwargs["xfails"] = xfails({
                absname(k, name if name else name_sep): v for k, v in (kwargs.get("xfails") or {}).items()
            }) or None
        kwargs["xflags"] = xflags({
                absname(k, name if name else name_sep): v for k, v in (kwargs.get("xflags") or {}).items()
            }) or None
        kwargs["only"] = [f.at(name if name else name_sep) for f in kwargs.get("only") or []] or None
        kwargs["skip"] = [f.at(name if name else name_sep) for f in kwargs.get("skip") or []] or None
        kwargs["start"] = kwargs.get("start").at(name if name else name_sep) if kwargs.get("start") else None
        kwargs["end"] = kwargs.get("end").at(name if name else name_sep) if kwargs.get("end") else None

        self.parent = parent
        self._apply_xflags(name, kwargs)
        self._apply_only(name, tags, kwargs)
        self._apply_skip(name, tags, kwargs)
        self._apply_start(name, tags, parent, kwargs)
        self._apply_end(name, tags, parent, kwargs)
        self._repeat = kwargs.pop("repeat", None)
        self._tags = tags
        self._kwargs = dict(kwargs)
        self.test = test(name, tags=tags, **kwargs)
        if getattr(self, "parent_type", None):
            self.test.parent_type = self.parent_type

    def _apply_end(self, name, tags, parent, kwargs):
        end = kwargs.get("end")
        if not end:
            return

        if end.match(name, tags):
            if parent:
                with parent.lock:
                    parent.end = None
                    parent.skip = [the("/*")]

    def _apply_start(self, name, tags, parent, kwargs):
        start = kwargs.get("start")
        if not start:
            return

        if not start.match(name, tags):
            kwargs["flags"] = Flags(kwargs.get("flags")) | SKIP
        else:
            kwargs["flags"] = Flags(kwargs.get("flags")) & ~SKIP
            if parent:
                with parent.lock:
                    parent.start = None

    def _apply_only(self, name, tags, kwargs):
        only = kwargs.get("only")
        if not only:
            return

        # only should not skip Given and Finally steps
        if kwargs.get("subtype") in (TestSubType.Given, TestSubType.Finally):
            only.append(the(join(name, "*")))

        found = False
        for item in only:
            if item.match(name, tags):
                found = True
                break

        if not found:
            kwargs["flags"] = Flags(kwargs.get("flags")) | SKIP
        else:
            kwargs["flags"] = Flags(kwargs.get("flags")) & ~SKIP

    def _apply_skip(self, name, tags, kwargs):
        skip = kwargs.get("skip")
        if not skip:
            return

        for item in skip:
            if item.match(name, tags, prefix=False):
                kwargs["flags"] = Flags(kwargs.get("flags")) | SKIP
                break

    def _apply_xflags(self, name, kwargs):
        xflags = kwargs.get("xflags")
        if not xflags:
            return
        for pattern, item in xflags.items():
            if match(name, pattern):
                set_flags, clear_flags = item
                kwargs["flags"] = (Flags(kwargs.get("flags")) & ~Flags(clear_flags)) | Flags(set_flags)

    def _parent_type_propagation(self, parent, kwargs):
        """Propagate parent test type if lower.

        :param parent: parent
        :param kwargs: test's kwargs
        """
        type = kwargs.pop("type", TestType.Test)
        subtype = kwargs.pop("subtype", TestSubType.Empty)

        parent_type = parent.type

        if parent_type == TestType.Iteration:
            parent_type = parent.parent_type

        if int(parent_type) < int(type):
            type = parent.type
            subtype = parent.subtype

        kwargs["subtype"] = subtype
        kwargs["type"] = type

    def __enter__(self):
        def dummy(*args, **kwargs):
            pass
        try:
            if self._repeat is not None:
                self._with_frame = inspect.currentframe().f_back
                self._with_block_start_lineno = self._with_frame.f_lineno
                self._with_source = inspect.getsourcelines(self._with_frame)
                self.trace = sys.gettrace()
                sys.settrace(dummy)
                sys._getframe(1).f_trace = functools.partial(self.__repeat__, None, None, None)
            else:
                return self.test.__enter__()
        except (KeyboardInterrupt, Exception):
            self.trace = sys.gettrace()
            sys.settrace(dummy)
            sys._getframe(1).f_trace = functools.partial(self.__nop__, *sys.exc_info())

    def __repeat__(self, *args):
        sys.settrace(self.trace)
        raise RepeatTestException()

    def __nop__(self, exc_type, exc_value, exc_tb, *args):
        sys.settrace(self.trace)
        raise exc_value.with_traceback(exc_tb)

    def __exit__(self, exception_type, exception_value, exception_traceback):
        frame = inspect.currentframe().f_back

        if isinstance(exception_value, RepeatTestException) and not frame.f_locals.get("_repeat") and top() != self.test:
            try:
                self.test.__enter__()
                self._with_block_end = frame.f_lasti
                self._with_block_end_lineno = self._with_frame.f_lineno
                lines, index = self._with_source
                index = max(1, index)
                source = textwrap.dedent("".join(
                    lines[(self._with_block_start_lineno - index)+1:(self._with_block_end_lineno - index) + 1]))
                code = compile(source, self._with_frame.f_code.co_filename, mode="exec")
                frame.f_locals["_repeat"] = True
                __kwargs = dict(self._kwargs)
                __kwargs.pop("name", None)
                __kwargs.pop("parent", None)
                __kwargs["type"] = TestType.Iteration
                __kwargs["subtype"] = TestSubType.Empty
                for i in range(self._repeat):
                    with Iteration(name=f"{i}", tags=self._tags, **__kwargs, parent_type=self.test.type):
                        exec(code, frame.f_globals, frame.f_locals)
                frame.f_locals.pop("_repeat")
            except:
                try:
                    test__exit__ = self.test.__exit__(*sys.exc_info())
                except(KeyboardInterrupt, Exception):
                    raise
            else:
                try:
                    test__exit__ = self.test.__exit__(None, None, None)
                except(KeyboardInterrupt, Exception):
                    raise
        else:
            try:
                test__exit__ = self.test.__exit__(exception_type, exception_value, exception_traceback)
            except (KeyboardInterrupt, Exception):
                raise

        # if test did not handle the exception in __exit__ then re-raise it
        if exception_value and not test__exit__:
            raise exception_value.with_traceback(exception_traceback)

        if not self.test.result:
            if not self.parent:
                sys.exit(1)

            if isinstance(self.test.result, Fail):
                result = Fail(self.parent.name, self.test.result.message)
            else:
                # convert Null into an Error
                result = Error(self.parent.name, self.test.result.message)

            if TE not in self.test.flags:
                raise ResultException(result)
            else:
                with self.parent.lock:
                    if isinstance(self.parent.result, Error):
                        pass
                    elif isinstance(self.test.result, Error) and ERROR_NOT_COUNTED not in self.test.flags:
                        self.parent.result = result
                    elif isinstance(self.test.result, Null) and NULL_NOT_COUNTED not in self.test.flags:
                        self.parent.result = result
                    elif isinstance(self.parent.result, Fail):
                        pass
                    elif isinstance(self.test.result, Fail) and FAIL_NOT_COUNTED not in self.test.flags:
                        self.parent.result = result
                    else:
                        pass
        return True

class Module(_test):
    """Module definition."""
    def __init__(self, name, **kwargs):
        kwargs["type"] = TestType.Module
        return super(Module, self).__init__(name, **kwargs)

class Suite(_test):
    """Suite definition."""
    def __init__(self, name, **kwargs):
        kwargs["type"] = TestType.Suite
        return super(Suite, self).__init__(name, **kwargs)

class Test(_test):
    """Test definition."""
    def __init__(self, name, **kwargs):
        kwargs["type"] = TestType.Test
        return super(Test, self).__init__(name, **kwargs)

class Iteration(_test):
    """Test iteration definition."""
    def __init__(self, name, **kwargs):
        kwargs["type"] = TestType.Iteration
        self.parent_type = kwargs.pop("parent_type", TestType.Test)
        return super(Iteration, self).__init__(name, **kwargs)

class Step(_test):
    """Step definition."""
    def __init__(self, name, **kwargs):
        kwargs["type"] = TestType.Step
        return super(Step, self).__init__(name, **kwargs)

# support for BDD
class Feature(Suite):
    def __init__(self, name, **kwargs):
        kwargs["subtype"] = TestSubType.Feature
        kwargs["_frame"] = kwargs.pop("_frame", inspect.currentframe().f_back )
        return super(Feature, self).__init__(name, **kwargs)

class Scenario(Test):
    def __init__(self, name, **kwargs):
        kwargs["subtype"] = TestSubType.Scenario
        kwargs["_frame"] = kwargs.pop("_frame", inspect.currentframe().f_back )
        return super(Scenario, self).__init__(name, **kwargs)

class _background(Test):
    def __init__(self, name, **kwargs):
        kwargs["subtype"] = TestSubType.Background
        kwargs["_frame"] = kwargs.pop("_frame", inspect.currentframe().f_back.f_back)
        return super(_background, self).__init__(name, **kwargs)

class Steps(ExitStack):
    def __init__(self, *args, **kwargs):
        self.values = []
        super(Steps, self).__init__(*args, **kwargs)

    def __call__(self, ctx):
        step = self.enter_context(ctx)
        self.values.append(step)
        return step

@contextmanager
def Background(name, **kwargs):
    with _background(name, **kwargs) as bg:
        with ExitStack() as stack:
            bg.stack = stack
            yield bg

def Given(name, **kwargs):
    kwargs["subtype"] = TestSubType.Given
    kwargs["_frame"] = kwargs.pop("_frame", inspect.currentframe().f_back )
    if isinstance(name, TestStep):
        return name(**kwargs)
    else:
        return Step(name, **kwargs)

def When(step, **kwargs):
    kwargs["subtype"] = TestSubType.When
    kwargs["_frame"] = kwargs.pop("_frame", inspect.currentframe().f_back )
    if isinstance(step, TestStep):
        return step(**kwargs)
    else:
        return Step(step, **kwargs)

def Then(name, **kwargs):
    kwargs["subtype"] = TestSubType.Then
    kwargs["_frame"] = kwargs.pop("_frame", inspect.currentframe().f_back )
    if isinstance(name, TestStep):
        return name(**kwargs)
    else:
        return Step(name, **kwargs)

def And(name, **kwargs):
    kwargs["subtype"] = TestSubType.And
    kwargs["_frame"] = kwargs.pop("_frame", inspect.currentframe().f_back )
    if isinstance(name, TestStep):
        return name(**kwargs)
    else:
        return Step(name, **kwargs)

def But(name, **kwargs):
    kwargs["subtype"] = TestSubType.But
    kwargs["_frame"] = kwargs.pop("_frame", inspect.currentframe().f_back )
    if isinstance(name, TestStep):
        return name(**kwargs)
    else:
        return Step(name, **kwargs)

def By(name, **kwargs):
    kwargs["subtype"] = TestSubType.By
    kwargs["_frame"] = kwargs.pop("_frame", inspect.currentframe().f_back )
    if isinstance(name, TestStep):
        return name(**kwargs)
    else:
        return Step(name, **kwargs)

def Finally(name, **kwargs):
    kwargs["subtype"] = TestSubType.Finally
    kwargs["_frame"] = kwargs.pop("_frame", inspect.currentframe().f_back )
    if isinstance(name, TestStep):
        return name(**kwargs)
    else:
        return Step(name, **kwargs)

class NullStep():
    def __enter__(self):
        return None

    def __exit__(self, *args, **kwargs):
        return False

# decorators
class _testdecorator(object):
    type = Test
    def __init__(self, func):
        func.name = getattr(func, "name", func.__name__.replace("_", " "))
        func.description = getattr(func, "description", func.__doc__)
        self.func = func
        functools.update_wrapper(self, func)

    def __call__(self, **kwargs):
        frame = kwargs.pop("_frame", inspect.currentframe().f_back)
        _kwargs = dict(vars(self.func))
        _name = kwargs.pop("name", None)
        _repeat = kwargs.pop("repeat", None)
        if kwargs.get("args") is None:
            # use function signature to get default argument values
            # and set them as args
            signature = inspect.signature(self.func)
            kwargs["args"] = { p.name:p.default for p in signature.parameters.values() if p.default != inspect.Parameter.empty }
        if _name is not None:
            kwargs["name"] = _name % (_kwargs)
        _kwargs.update(kwargs)
        # handle repeats
        if _repeat is not None:
            with self.type(**_kwargs, _frame=frame, _run=False) as parent_test:
                __kwargs = dict(_kwargs)
                __kwargs.pop("name")
                __kwargs["type"] = TestType.Iteration
                __kwargs["subtype"] = TestSubType.Empty
                for i in range(_repeat):
                    with Iteration(name=f"{i}", parent_type=parent_test.type, **__kwargs) as test:
                        self.func(**{name: arg.value for name, arg in test.args.items()})
            return parent_test
        else:
            with self.type(**_kwargs, _frame=frame) as test:
                self.func(**{name: arg.value for name, arg in test.args.items()})
            return test

class TestStep(_testdecorator):
    type = Step

class TestCase(_testdecorator):
    type = Test

class TestScenario(TestCase):
    type = Scenario

class TestSuite(_testdecorator):
    type = Suite

class TestFeature(TestSuite):
    type = Feature

class TestModule(_testdecorator):
    type = Module

class TestClass(object):
    def __init__(self, test):
        self.test = test

    def __call__(self, func):
        func.test = self.test
        return func

class Name(object):
    def __init__(self, name):
        self.name = name

    def __call__(self, func):
        func.name = self.name
        return func

class Description(object):
    def __init__(self, description):
        self.description = description

    def __call__(self, func):
        func.description = self.description
        return func

class Examples(object):
    def __init__(self, header, rows, row_format=None):
        self.examples = ExamplesTable(header, rows=rows, row_format=row_format)

    def __call__(self, func):
        func.examples = self.examples
        return func

class Attributes(object):
    def __init__(self, *attributes):
        self.attributes = attributes

    def __call__(self, func):
        func.attributes = self.attributes
        return func

class Requirements(object):
    def __init__(self, *requirements):
        self.requirements = requirements

    def __call__(self, func):
        func.requirements = self.requirements
        return func

class Tags(object):
    def __init__(self, *tags):
        self.tags = tags

    def __call__(self, func):
        func.tags = self.tags
        return func

class Uid(object):
    def __init__(self, uid):
        self.uid = uid

    def __call__(self, func):
        func.uid = self.uid
        return func

class Users(object):
    def __init__(self, *users):
        self.users = users

    def __call__(self, func):
        func.users = self.users
        return func

class Tickets(object):
    def __init__(self, *tickets):
        self.tickets = tickets

    def __call__(self, func):
        func.tickets = self.tickets
        return func

def run(comment=None, test=None, **kwargs):
    """Run a test.

    :param comment: comment
    :param test: test
    :param **kwargs: other test arguments
    """
    if test is None:
        raise TypeError("run() test argument must be specified")
    if inspect.isclass(test) and issubclass(test, TestBase):
        test = test
    elif issubclass(type(test), _testdecorator):
        return test(**kwargs, _frame=inspect.currentframe().f_back)
    elif inspect.isfunction(test):
        return test(**kwargs)
    elif inspect.ismethod(test):
        return test(**kwargs)
    else:
        raise TypeError(f"invalid test type '{type(test)}'")

    _frame = inspect.currentframe().f_back
    _repeat = kwargs.pop("repeat", None)
    if _repeat is not None:
        with globals()["Test"](test=test, name=kwargs.pop("name", None), **kwargs,
                               _frame=_frame, _run=False) as parent_test:
            _kwargs = dict(kwargs)
            _kwargs["type"] = TestType.Iteration
            _kwargs["subtype"] = TestSubType.Empty
            for i in range(_repeat):
                with Iteration(test=test, name=f"{i}", **kwargs, _frame=_frame, parent_type=parent_test.type):
                    pass
        return parent_test.result
    else:
        with globals()["Test"](test=test, name=kwargs.pop("name", None), **kwargs, _frame=frame) as test:
            pass
        return test.result


'''https://pythonprogramming.net/basics-data-analysis-python-pandas-tutorial/
'''
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib import style
style.use('ggplot')
style.use('fivethirtyeight')

web_stats = {'Day':[1,2,3,4,5,6],
             'Visitors':[43,34,65,56,29,76],
             'Bounce Rate':[65,67,78,65,45,52]}

df = pd.DataFrame(web_stats)

#print(df.head())
#print(df.tail())
#print(df.tail(2))

#df.set_index('Day', inplace=True)
#df = df.set_index('Day')

#print(df['Visitors'])
#print(df.Visitors)

#df['Visitors'].plot()
#plt.show()

df.plot()
plt.show()

print(df[['Visitors','Bounce Rate']])


import html5lib

html5parser = html5lib.HTMLParser(strict=True)

try:
    p = html5parser.parse(""" <!DOCTYPE html>

<title>Page Title</title>

This is a Heading

This is a paragraph.

</htm>
""")
except:
    p = False

if p:
    print("OK")
else:
    print("No")
print(p)


# Copyright (c) 2017 The Khronos Group Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import division, print_function

from collections import OrderedDict

import nnef
import numpy as np

from .. import tf_compat
from ...common import dog
from ...common import utils

class Converter(object):
    def __init__(self, nnefdog):
        # TODO maybe we would need tfdn_by_nnefdn_name
        self.tfdn_by_name = OrderedDict()  # doesn't have to be ordered
        self.nnefop_by_tfop = OrderedDict()
        self.nnefdn_by_tfdn = OrderedDict()
        self.nnefdog = nnefdog

    def make_tfdn(self, nnefdn=None, name=None, discriminator=None):
        assert (nnefdn is None) != (name is None)
        if name is None:
            name = nnefdn.name
        if discriminator is not None:
            counter = 1
            name_candidate = name + "_{}{}".format(discriminator, counter)
            while name_candidate in self.tfdn_by_name:  # TODO optimize
                counter += 1
                name_candidate = name + "_{}{}".format(discriminator, counter)
            name = name_candidate
        if name in self.tfdn_by_name:
            utils.print_error("Duplicate dn definition: {}".format(name))
            return self.tfdn_by_name[name]
        tfdn = dog.DataNode(name)
        self.tfdn_by_name[name] = tfdn
        self.nnefdn_by_tfdn[tfdn] = nnefdn
        return tfdn

    def get_tfdn_safe(self, nnefdn=None, name=None):
        if isinstance(nnefdn, (bool, int, float)):
            return nnefdn
        assert (nnefdn is None) != (name is None)
        orig_name = name
        if name is None:
            name = nnefdn.name
        if name not in self.tfdn_by_name:
            utils.print_error("Dn not defined: {}".format(name))
            return self.make_tfdn(nnefdn=nnefdn, name=orig_name)
        return self.tfdn_by_name[name]

    def add_tfop(self, tfop, nnefop=None):
        assert tfop not in self.nnefop_by_tfop
        self.nnefop_by_tfop[tfop] = nnefop

    @staticmethod
    def get_shape_safe(tfdn_or_other):
        return dog.get_shape_safe(tfdn_or_other)

    @staticmethod
    def get_rank_safe(tfdn_or_other):
        return len(dog.get_shape_safe(tfdn_or_other))

    @staticmethod
    def is_binary_nnefop_name(nnefop_name):
        return DefaultConverters[nnefop_name] == generic_convert_binary

# HELPERS:

# TODO revise these methods, is negative pad_total possible?

def calculate_same_padding_element_ex(input_size, filter_size, stride, dilation):
    dilated_filter_size = (filter_size - 1) * dilation + 1
    output_size = int(np.ceil(float(input_size) / float(stride)))
    pad_total = (output_size - 1) * stride + dilated_filter_size - input_size
    if pad_total >= 0:
        pad_front = pad_total // 2
        pad_back = pad_total - pad_front
        return pad_front, pad_back
    else:
        return 0, pad_total

def calculate_same_padding_ex(input_size, filter_size, stride, dilation):
    return [calculate_same_padding_element_ex(i, f, s, d) for i, f, s, d in
            zip(input_size, filter_size, stride, dilation)]

def calculate_valid_padding_element_ex(input_size, filter_size, stride, dilation):
    dilated_filter_size = (filter_size - 1) * dilation + 1
    output_size = int(np.ceil(float(input_size - dilated_filter_size + 1) / float(stride)))
    pad_total = (output_size - 1) * stride + dilated_filter_size - input_size
    return 0, pad_total

def calculate_valid_padding_ex(input_size, filter_size, stride, dilation):
    return [calculate_valid_padding_element_ex(i, f, s, d) for i, f, s, d in
            zip(input_size, filter_size, stride, dilation)]

# TODO unify this with trafos.get_paddings

def calculate_tfpadding_for_deconv(nnefop):
    nnefborder = nnefop.args["border"].lower()
    nnefpadding = nnefop.args["padding"]

    if "input" in nnefop.args:
        input_shape = dog.get_shape_safe(nnefop.args["input"])
    elif "orig_input" in nnefop.args:
        input_shape = dog.get_shape_safe(nnefop.args["orig_input"])
    else:
        input_shape = nnefop.args["orig_input_shape"]

    if "filter" in nnefop.args:
        filter_shape = dog.get_shape_safe(nnefop.args["filter"])
    elif "orig_filter" in nnefop.args:
        filter_shape = dog.get_shape_safe(nnefop.args["orig_filter"])
    else:
        filter_shape = nnefop.args["orig_filter_shape"]
    # TODO Check this
    same_padding = calculate_same_padding_ex(input_size=input_shape[1:-1],  # tf order
                                             filter_size=filter_shape[:-2],
                                             stride=nnefop.args["stride"],
                                             dilation=nnefop.args["dilation"])
    valid_padding = calculate_valid_padding_ex(input_size=input_shape[1:-1],
                                               filter_size=filter_shape[:-2],
                                               stride=nnefop.args["stride"],
                                               dilation=nnefop.args["dilation"])
    if nnefborder == "constant" and (nnefpadding == [] or nnefpadding == same_padding):
        return "SAME"
    elif nnefborder == "constant" and (not utils.has_greater_than_0(nnefpadding) or nnefpadding == valid_padding):
        return "VALID"
    else:
        utils.print_error("Special paddings or borders are not supported for deconv.")
        return "SAME"

def calculate_nnef_reshape_shape(old_shape, new_shape):
    new_shape = list(new_shape)
    minus_one_idx = None
    prod = 1
    for i in range(len(new_shape)):
        if new_shape[i] == 0:
            if i >= len(old_shape):
                utils.print_error("New shape has 0 at index {}, but old shape is only {} long, setting dimension to 1"
                                  .format(i, len(old_shape)))
                new_shape[i] = 1
            else:
                new_shape[i] = old_shape[i]
        if new_shape[i] == -1:
            if minus_one_idx is None:
                minus_one_idx = i
            else:
                utils.print_error("Shape had more than one -1's, setting it to 1")
                new_shape[i] = 1
        if new_shape[i] != -1:
            prod *= new_shape[i]

    if minus_one_idx is not None:
        if int(np.prod(old_shape)) % prod != 0:
            utils.print_error("-1 is used in shape but the product of other dimensions does not divide the old size")
        new_shape[minus_one_idx] = int(np.prod(old_shape)) // prod

    return new_shape

def nnefarray_to_tf(nnefarray, nnefshape):
    if not isinstance(nnefarray, np.ndarray):
        nnefarray = np.array(nnefarray)

    if nnefarray.size == 1:
        return nnefarray

    return np.reshape(nnefarray, nnefshape)

def nnefstride_to_tf_nhwc(stride, in_spatial, out_spatial, empty_value=None):
    if not stride:
        return empty_value

    if in_spatial:
        stride = [1, 1] + stride

    nhwc_stride = utils.shape_nchw_to_nhwc(stride)

    if out_spatial:
        nhwc_stride = nhwc_stride[1:-1]

    return nhwc_stride

def spatial_nhwc(stride_like, in_spatial, out_spatial, empty_value, unit_elem=1):
    if not stride_like:
        return empty_value

    if bool(in_spatial) == bool(out_spatial):
        return list(stride_like)

    if in_spatial:
        stride_like = [unit_elem] + stride_like + [unit_elem]

    if out_spatial:
        stride_like = stride_like[1:-1]

    return stride_like

def nnefborder_to_tf(border):
    to_tf = {
        'ignore': None,
        'constant': 'CONSTANT',
        'replicate': None,
        'reflect': 'REFLECT',
        'reflect-even': 'SYMMETRIC'
    }

    if to_tf[border] is None:
        utils.print_error("Border mode {} is not supported in TensorFlow.".format(border))
        return 'CONSTANT'
    else:
        return to_tf[border]

# CONVERTERS:

def convert_external(nnefop, converter):
    tfop = dog.OperationNode("tf.placeholder")
    tfop.add_arg("dtype", utils.nnef_dtype_to_tf(nnefop.result.dtype))
    tfop.add_arg("shape", nnefop.args["shape"])
    tfop.add_arg("name", nnefop.result.name)
    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop, nnefop)

def convert_variable(nnefop, converter):
    tfop = dog.OperationNode("tf.get_variable")
    tfop.add_arg("name", nnefop.args["label"])
    tfop.add_arg("shape", nnefop.args["shape"])
    tfop.add_arg("dtype", utils.nnef_dtype_to_tf(nnefop.result.dtype))
    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop, nnefop)

def convert_constant(nnefop, converter):
    array = nnefarray_to_tf(nnefop.args["value"], nnefop.args["shape"])

    shape = nnefop.args["shape"]
    if list(array.shape) != list(shape) and list(array.shape) != [1]:
        utils.print_error("Internal error, shape mismatch: {} {} {} {}"
                          .format(nnefop.result.name, list(array.shape), list(shape), nnefop.result.extra))

    tfop = dog.OperationNode("tf.constant")
    tfop.add_arg("value", array)
    tfop.add_arg("dtype", utils.nnef_dtype_to_tf(nnefop.result.dtype))
    tfop.add_arg("shape", shape)
    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop, nnefop)

def convert_update(nnefop, converter):
    tfop = dog.OperationNode("tf.assign")
    tfop.add_arg("ref", converter.get_tfdn_safe(nnefop.args["variable"]))
    tfop.add_arg("value", converter.get_tfdn_safe(nnefop.args["value"]))
    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop, nnefop)

def convert_reshape(nnefop, converter):
    if nnefop.args["shape"] == [0, -1]:
        tfop = dog.OperationNode("tf.contrib.layers.flatten")
        tfop.add_arg("inputs", converter.get_tfdn_safe(nnefop.args["input"]))
        tfop.add_result("result", converter.make_tfdn(nnefop.result))
        converter.add_tfop(tfop, nnefop)
    else:
        nnefdn_input = nnefop.args["input"]
        tfop = dog.OperationNode("tf.reshape")
        tfop.add_arg("tensor", converter.get_tfdn_safe(nnefdn_input))

        nnef_shape = calculate_nnef_reshape_shape(dog.get_shape_safe(nnefdn_input),
                                                  nnefop.args["shape"])  # TODO needed?
        tfop.add_arg("shape", nnef_shape)

        tfop.add_result("result", converter.make_tfdn(nnefop.result))
        converter.add_tfop(tfop, nnefop)

def convert_squeeze(nnefop, converter):
    nnef_axes = nnefop.args["axes"]
    nnefdn_input = nnefop.args["input"]

    tfop = dog.OperationNode("tf.squeeze")
    tfop.add_arg("input", converter.get_tfdn_safe(nnefdn_input))
    tfop.add_arg("axis", nnef_axes)
    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop, nnefop)

def partial_convert_unsqueeze_to_reshape(nnefop, converter):
    nnef_axes = nnefop.args["axes"]
    nnefdn_input = nnefop.args["input"]
    nnef_input_shape = dog.get_shape_safe(nnefdn_input)

    shape = utils.apply_unsqueeze_shape(list(nnef_input_shape), nnef_axes)

    tfop = dog.OperationNode("tf.reshape")
    tfop.add_arg("tensor", converter.get_tfdn_safe(nnefdn_input))
    tfop.add_arg("shape", shape)
    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop, nnefop)

def partial_convert_unsqueeze_to_expand_dims(nnefop, converter):
    assert len(nnefop.args["axes"]) == 1

    axis = nnefop.args["axes"][0]
    nnefdn_input = nnefop.args["input"]

    tfop = dog.OperationNode("tf.expand_dims")
    tfop.add_arg("input", converter.get_tfdn_safe(nnefdn_input))
    tfop.add_arg("axis", axis)
    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop, nnefop)

def convert_unsqueeze(nnefop, converter):
    if len(nnefop.args["axes"]) == 1:
        partial_convert_unsqueeze_to_expand_dims(nnefop, converter)
    else:
        partial_convert_unsqueeze_to_reshape(nnefop, converter)

def convert_transpose(nnefop, converter):
    tfop = dog.OperationNode("tf.transpose")
    tfop.add_arg("a", nnefop.args["input"])
    # TODO something might be needed like adding axes which are not written out
    tfop.add_arg("perm", nnefop.args["axes"])
    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop, nnefop)

def convert_concat(nnefop, converter):
    nnefdns = nnefop.args["values"]
    tfop = dog.OperationNode("tf.concat")
    tfop.add_arg("values", [converter.get_tfdn_safe(nnefdn) for nnefdn in nnefdns])
    tfop.add_arg("axis", nnefop.args['axis'])
    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop, nnefop)

def convert_split(nnefop, converter):
    nnefaxis = nnefop.args["axis"]
    nnefdn_value = nnefop.args["value"]
    nnefratios = nnefop.args["ratios"]
    nnefshape = dog.get_shape_safe(nnefdn_value)
    if nnefshape[nnefaxis] == 0:
        if len(np.unique(nnefratios)) == 1:
            tf_num_or_size_splits = len(nnefratios)
        else:
            utils.print_error("Split with different ratios on an axis ({}) with unspecified size.".format(nnefaxis))
            tf_num_or_size_splits = len(nnefratios)
    else:
        part_size = nnefshape[nnefaxis] // int(np.sum(nnefratios))
        tf_num_or_size_splits = [ratio * part_size for ratio in nnefratios]

    tfop = dog.OperationNode("tf.split")
    tfop.add_arg("value", converter.get_tfdn_safe(nnefdn_value))
    tfop.add_arg("num_or_size_splits", tf_num_or_size_splits)
    tfop.add_arg("axis", nnefaxis)
    for i in range(len(nnefratios)):
        tfop.add_result("result{}".format(i), converter.make_tfdn(nnefop.result[i]))
    converter.add_tfop(tfop, nnefop)

def generic_convert_binary(nnefop, converter):
    tfop = dog.OperationNode(nnefopname_to_tf(nnefop.name))
    tfop.add_arg("x", converter.get_tfdn_safe(nnefop.args["x"]))
    tfop.add_arg("y", converter.get_tfdn_safe(nnefop.args["y"]))
    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop, nnefop)

def generic_convert_unary(nnefop, converter):
    tfop = dog.OperationNode(nnefopname_to_tf(nnefop.name))
    tfop.add_arg("x", converter.get_tfdn_safe(nnefop.args["x"]))
    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop, nnefop)

def convert_rsqr(nnefop, converter):
    tfop = dog.OperationNode("tf.pow")
    tfop.add_arg("x", converter.get_tfdn_safe(nnefop.args["x"]))
    tfop.add_arg("y", -2.0)
    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop, nnefop)

def convert_select(nnefop, converter):
    x = converter.get_tfdn_safe(nnefop.args["true_value"])
    y = converter.get_tfdn_safe(nnefop.args["false_value"])
    cond_shape = dog.get_shape_safe(nnefop.args["condition"])

    if isinstance(x, (float, int, bool)):
        tfop_constant = dog.OperationNode("tf.constant")
        tfop_constant.add_arg("value", x)
        tfop_constant.add_arg("dtype", utils.tf_type_of_python_scalar(x))
        tfop_constant.add_arg("shape", cond_shape)
        tfdn_constant = dog.DataNode(nnefop.result.name + "_x__")
        tfdn_constant.shape = cond_shape
        tfop_constant.add_result("result", converter.make_tfdn(tfdn_constant))
        converter.add_tfop(tfop_constant, nnefop)
        x = tfdn_constant

    if isinstance(y, (float, int, bool)):
        tfop_constant = dog.OperationNode("tf.constant")
        tfop_constant.add_arg("value", y)
        tfop_constant.add_arg("dtype", utils.tf_type_of_python_scalar(y))
        tfop_constant.add_arg("shape", cond_shape)
        tfdn_constant = dog.DataNode(nnefop.result.name + "_y__")
        tfdn_constant.shape = cond_shape
        tfop_constant.add_result("result", converter.make_tfdn(tfdn_constant))
        converter.add_tfop(tfop_constant, nnefop)
        y = tfdn_constant

    tfop = dog.OperationNode("tf.where")
    tfop.add_arg("condition", converter.get_tfdn_safe(nnefop.args["condition"]))
    tfop.add_arg("x", x)
    tfop.add_arg("y", y)
    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop, nnefop)

def convert_clamp(nnefop, converter):
    tfop = dog.OperationNode("tf.clip_by_value")
    tfop.add_arg("t", converter.get_tfdn_safe(nnefop.args["x"]))
    tfop.add_arg("clip_value_min", converter.get_tfdn_safe(nnefop.args["a"]))
    tfop.add_arg("clip_value_max", converter.get_tfdn_safe(nnefop.args["b"]))
    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop, nnefop)

def convert_matmul(nnefop, converter):
    tfop = dog.OperationNode("tf.matmul")
    tfop.add_arg("a", converter.get_tfdn_safe(nnefop.args["A"]))
    tfop.add_arg("b", converter.get_tfdn_safe(nnefop.args["B"]))
    tfop.add_arg("transpose_a", nnefop.args["transposeA"])
    tfop.add_arg("transpose_b", nnefop.args["transposeB"])
    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop, nnefop)

def partial_convert_conv_to_conv2d_or_conv3d(nnefop, converter):
    d = len(dog.get_shape_safe(nnefop.args["input"])) - 2

    tfop = dog.OperationNode("tf.nn.conv2d" if d == 2 else "tf.nn.conv3d")

    tfop.add_arg("input", converter.get_tfdn_safe(nnefop.args["input"]))
    tfop.add_arg("filter", converter.get_tfdn_safe(nnefop.args["filter"]))
    tfop.add_arg("strides", spatial_nhwc(nnefop.args["stride"], in_spatial=True,
                                         out_spatial=False, empty_value=[1] * (d + 2)))
    tfop.add_arg("padding", nnefop.args["_tf_padding"])

    if utils.has_greater_than_1(nnefop.args["dilation"]):
        tfop.add_arg("dilations", spatial_nhwc(nnefop.args["dilation"], in_spatial=True,
                                               out_spatial=False, empty_value=[1] * (d + 2)))

    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop)

def partial_convert_conv_to_convolution(nnefop, converter):
    d = len(dog.get_shape_safe(nnefop.args["input"])) - 2

    tfop = dog.OperationNode("tf.nn.convolution")
    tfop.add_arg("input", converter.get_tfdn_safe(nnefop.args["input"]))
    tfop.add_arg("filter", converter.get_tfdn_safe(nnefop.args["filter"]))
    tfop.add_arg("padding", nnefop.args["_tf_padding"])
    if utils.has_greater_than_1(nnefop.args["stride"]):
        tfop.add_arg("strides", nnefop.args["stride"])
    if utils.has_greater_than_1(nnefop.args["dilation"]):
        tfop.add_arg("dilation_rate", nnefop.args["dilation"])
    if utils.has_greater_than_1(nnefop.args["stride"]) and utils.has_greater_than_1(nnefop.args["dilation"]):
        utils.print_error("Custom stride AND dilation is not supported by tf.nn.convolution!")
    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop)

def partial_convert_conv_to_deptwise_conv2d(nnefop, converter):
    tfop = dog.OperationNode("tf.nn.depthwise_conv2d")
    tfop.add_arg("input", converter.get_tfdn_safe(nnefop.args["input"]))
    tfop.add_arg("filter", converter.get_tfdn_safe(nnefop.args["filter"]))
    tfop.add_arg("strides", spatial_nhwc(nnefop.args["stride"],
                                         in_spatial=True, out_spatial=False, empty_value=[1] * 4))
    tfop.add_arg("padding", nnefop.args["_tf_padding"])
    if utils.has_greater_than_1(nnefop.args["dilation"]):
        tfop.add_arg("rate", spatial_nhwc(nnefop.args["dilation"], in_spatial=True, out_spatial=True, empty_value=None))
    if utils.has_greater_than_1(nnefop.args["stride"]) and utils.has_greater_than_1(nnefop.args["dilation"]):
        utils.print_error("Custom stride AND dilation is not supported by tf.nn.depthwise_conv2d!")
    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop)

def convert_conv(nnefop, converter):
    groups = nnefop.args["groups"]
    nnef_input_shape = dog.get_shape_safe(nnefop.args["input"])
    d = len(nnef_input_shape) - 2

    if not nnefop.args["stride"]:
        nnefop.args["stride"] = [1 for _ in range(d)]
    if not nnefop.args["dilation"]:
        nnefop.args["dilation"] = [1 for _ in range(d)]

    dilated = utils.has_greater_than_1(nnefop.args["dilation"])

    if groups == 1:
        if d in [2, 3] and not dilated:
            partial_convert_conv_to_conv2d_or_conv3d(nnefop, converter)
        else:
            partial_convert_conv_to_convolution(nnefop, converter)
    else:
        if groups in [0, nnef_input_shape[-1]] and d == 2:
            partial_convert_conv_to_deptwise_conv2d(nnefop, converter)
        else:
            # TODO consider implementing this with split and conv
            utils.print_error(
                "Grouped convolutions are only supported if they can be converted to tf.nn.depthwise_conv2d.")
            return

def partial_convert_conv_grad_input_normal_2d(nnefop, converter):
    tfop = dog.OperationNode("tf.nn.conv2d_backprop_input")
    tfop.add_arg("input_sizes", nnefop.args["orig_input_shape"])
    tfop.add_arg("filter", converter.get_tfdn_safe(nnefop.args["orig_filter"]))
    tfop.add_arg("out_backprop", converter.get_tfdn_safe(nnefop.args["output_grad"]))
    tfop.add_arg("strides", spatial_nhwc(nnefop.args["stride"],
                                         in_spatial=True, out_spatial=False, empty_value=[1] * 4))
    tfop.add_arg("padding", calculate_tfpadding_for_deconv(nnefop))
    if utils.has_greater_than_1(nnefop.args["dilation"]):
        tfop.add_arg("dilations", spatial_nhwc(nnefop.args["dilation"],
                                               in_spatial=True, out_spatial=False, empty_value=[1] * 4))
    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop)

def partial_convert_conv_grad_input_normal_3d(nnefop, converter):
    tfop = dog.OperationNode("tf.nn.conv3d_transpose")
    tfop.add_arg("value", converter.get_tfdn_safe(nnefop.args["output_grad"]))
    tfop.add_arg("filter", converter.get_tfdn_safe(nnefop.args["orig_filter"]))
    tfop.add_arg("output_shape", nnefop.args["orig_input_shape"])
    tfop.add_arg("strides", spatial_nhwc(nnefop.args["stride"],
                                         in_spatial=True, out_spatial=False, empty_value=[1] * 5))
    tfop.add_arg("padding", calculate_tfpadding_for_deconv(nnefop))
    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop)

def partial_convert_conv_grad_input_planewise(nnefop, converter):
    tfop = dog.OperationNode("tf.nn.depthwise_conv2d_native_backprop_input")
    tfop.add_arg("input_sizes", nnefop.args["orig_input_shape"])
    tfop.add_arg("filter", converter.get_tfdn_safe(nnefop.args["orig_filter"]))
    tfop.add_arg("out_backprop", converter.get_tfdn_safe(nnefop.args["output_grad"]))
    tfop.add_arg("strides", spatial_nhwc(nnefop.args["stride"],
                                         in_spatial=True, out_spatial=False, empty_value=[1] * 4))
    tfop.add_arg("padding", calculate_tfpadding_for_deconv(nnefop))
    if utils.has_greater_than_1(nnefop.args["dilation"]):
        tfop.add_arg("rate",
                     spatial_nhwc(nnefop.args["dilation"],
                                  in_spatial=True, out_spatial=True, empty_value=None))
    if utils.has_greater_than_1(nnefop.args["stride"]) and utils.has_greater_than_1(nnefop.args["dilation"]):
        utils.print_error(
            "Custom stride AND dilation is not supported by tf.nn.depthwise_conv2d_native_backprop_input!")
    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop)

def convert_conv_grad_input(nnefop, converter):
    groups = nnefop.args["groups"]
    nnef_input_shape = nnefop.args["orig_input_shape"]
    d = len(nnef_input_shape) - 2

    if not nnefop.args["stride"]:
        nnefop.args["stride"] = [1 for _ in range(d)]
    if not nnefop.args["dilation"]:
        nnefop.args["dilation"] = [1 for _ in range(d)]

    dilated = utils.has_greater_than_1(nnefop.args["dilation"])

    if groups == 1:
        if d == 2:
            partial_convert_conv_grad_input_normal_2d(nnefop, converter)
        elif d == 3 and not dilated:
            partial_convert_conv_grad_input_normal_3d(nnefop, converter)
        else:
            utils.print_error("Unsupported conv_grad_input: {}d, {}".format(d, "dilated" if dilated else "not dilated"))
    else:
        if groups in [0, nnef_input_shape[-1]] and d == 2:
            partial_convert_conv_grad_input_planewise(nnefop, converter)
        else:
            # TODO consider implementing this
            utils.print_error(
                "Grouped deconvolutions are only supported if they can be converted to "
                + "tf.nn.depthwise_conv2d_native_backprop_input.")
            return

def partial_convert_conv_grad_filter_normal(nnefop, converter):
    d = len(dog.get_shape_safe(nnefop.args["orig_input"])) - 2

    tfop = dog.OperationNode("tf.nn.conv2d_backprop_filter" if d == 2 else "tf.nn.conv3d_backprop_filter_v2")
    tfop.add_arg("input", converter.get_tfdn_safe(nnefop.args["orig_input"]))
    tfop.add_arg("filter_sizes", nnefop.args["orig_filter_shape"])
    tfop.add_arg("out_backprop", converter.get_tfdn_safe(nnefop.args["output_grad"]))
    tfop.add_arg("strides", spatial_nhwc(nnefop.args["stride"],
                                         in_spatial=True, out_spatial=False, empty_value=[1] * (d + 2)))
    tfop.add_arg("padding", calculate_tfpadding_for_deconv(nnefop))
    if utils.has_greater_than_1(nnefop.args["dilation"]):
        tfop.add_arg("dilations", spatial_nhwc(nnefop.args["dilation"],
                                               in_spatial=True, out_spatial=False, empty_value=[1] * (d + 2)))
    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop)

def partial_convert_conv_grad_filter_planewise(nnefop, converter):
    tfop = dog.OperationNode("tf.nn.depthwise_conv2d_native_backprop_filter")
    tfop.add_arg("input", converter.get_tfdn_safe(nnefop.args["orig_input"]))
    tfop.add_arg("filter_sizes", nnefop.args["orig_filter_shape"])
    tfop.add_arg("out_backprop", converter.get_tfdn_safe(nnefop.args["output_grad"]))
    tfop.add_arg("strides", spatial_nhwc(nnefop.args["stride"],
                                         in_spatial=True, out_spatial=False, empty_value=[1] * 4))
    tfop.add_arg("padding", calculate_tfpadding_for_deconv(nnefop))
    if utils.has_greater_than_1(nnefop.args["dilation"]):
        tfop.add_arg("rate",
                     spatial_nhwc(nnefop.args["dilation"], in_spatial=True, out_spatial=True, empty_value=None))
    if utils.has_greater_than_1(nnefop.args["stride"]) and utils.has_greater_than_1(nnefop.args["dilation"]):
        utils.print_error(
            "Custom stride AND dilation is not supported by tf.nn.depthwise_conv2d_native_backprop_filter!")
    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop)

def convert_conv_grad_filter(nnefop, converter):
    groups = nnefop.args["groups"]
    nnef_input_shape = dog.get_shape_safe(nnefop.args["orig_input"])
    d = len(nnef_input_shape) - 2

    if not nnefop.args["stride"]:
        nnefop.args["stride"] = [1 for _ in range(d)]
    if not nnefop.args["dilation"]:
        nnefop.args["dilation"] = [1 for _ in range(d)]

    if groups == 1:
        if d in [2, 3]:
            partial_convert_conv_grad_filter_normal(nnefop, converter)
        else:
            utils.print_error("Unsupported conv_grad_filter: {}d".format(d))
    else:
        if groups in [0, nnef_input_shape[-1]] and d == 2:
            partial_convert_conv_grad_filter_planewise(nnefop, converter)
        else:
            # TODO consider implementing this
            utils.print_error(
                "Grouped convolutions are only supported if they can be converted to "
                + "tf.nn.depthwise_conv2d_native_backprop_filter.")
            return

def partial_convert_deconv_to_conv2d_transpose_or_conv3d_transpose(nnefop, converter):
    d = len(dog.get_shape_safe(nnefop.args["input"])) - 2

    tfop = dog.OperationNode("tf.nn.conv2d_transpose" if d == 2 else "tf.nn.conv3d_transpose")
    tfop.add_arg("value", converter.get_tfdn_safe(nnefop.args["input"]))
    tfop.add_arg("filter", converter.get_tfdn_safe(nnefop.args["filter"]))
    if nnefop.args["output_shape"]:
        tfop.add_arg("output_shape", nnefop.args["output_shape"])
    tfop.add_arg("strides", spatial_nhwc(nnefop.args["stride"],
                                         in_spatial=True, out_spatial=False, empty_value=[1] * (d + 2)))
    tfop.add_arg("padding", calculate_tfpadding_for_deconv(nnefop))
    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop)

def partial_convert_deconv_to_atrous_conv2d_transpose(nnefop, converter):
    tfdilation = spatial_nhwc(nnefop.args["dilation"], in_spatial=True, out_spatial=True, empty_value=[1] * 2)

    if len(np.unique(tfdilation)) > 1:
        utils.print_error("Cannot specify different x and y dilation in tf.nn.atrous_conv2d_transpose.")

    tfop = dog.OperationNode("tf.nn.atrous_conv2d_transpose")
    tfop.add_arg("value", converter.get_tfdn_safe(nnefop.args["input"]))
    tfop.add_arg("filters", converter.get_tfdn_safe(nnefop.args["filter"]))
    tfop.add_arg("output_shape", nnefop.args["output_shape"])
    tfop.add_arg("rate", tfdilation[0])
    tfop.add_arg("padding", calculate_tfpadding_for_deconv(nnefop))

    if utils.has_greater_than_1(nnefop.args["stride"]):
        utils.print_error("Cannot use stride>1 in tf.nn.atrous_conv2d_transpose.")

    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop)

def partial_convert_deconv_to_depthwise_conv2d_native_backprop_input(nnefop, converter):
    tfop = dog.OperationNode("tf.nn.depthwise_conv2d_native_backprop_input")
    tfop.add_arg("input_sizes", nnefop.args["output_shape"])
    tfop.add_arg("filter", converter.get_tfdn_safe(nnefop.args["filter"]))
    tfop.add_arg("out_backprop", converter.get_tfdn_safe(nnefop.args["input"]))
    tfop.add_arg("strides", spatial_nhwc(nnefop.args["stride"],
                                         in_spatial=True, out_spatial=False, empty_value=[1] * 4))
    tfop.add_arg("padding", calculate_tfpadding_for_deconv(nnefop))
    if utils.has_greater_than_1(nnefop.args["dilation"]):
        tfop.add_arg("rate",
                     spatial_nhwc(nnefop.args["dilation"],
                                  in_spatial=True, out_spatial=True, empty_value=None))
    if utils.has_greater_than_1(nnefop.args["stride"]) and utils.has_greater_than_1(nnefop.args["dilation"]):
        utils.print_error(
            "Custom stride AND dilation is not supported by tf.nn.depthwise_conv2d_native_backprop_input!")
    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop)

def convert_deconv(nnefop, converter):
    groups = nnefop.args["groups"]
    nnef_input_shape = dog.get_shape_safe(nnefop.args["input"])
    d = len(nnef_input_shape) - 2

    if not nnefop.args["stride"]:
        nnefop.args["stride"] = [1 for _ in range(d)]
    if not nnefop.args["dilation"]:
        nnefop.args["dilation"] = [1 for _ in range(d)]

    is_dilated = utils.has_greater_than_1(nnefop.args["dilation"])

    if groups == 1:
        if not is_dilated and d in [2, 3]:
            partial_convert_deconv_to_conv2d_transpose_or_conv3d_transpose(nnefop, converter)
        elif is_dilated and d == 2:
            partial_convert_deconv_to_atrous_conv2d_transpose(nnefop, converter)
        else:
            utils.print_error(
                "{} dimensional{} deconv is not supported by TensorFlow.".format(d, " dilated" if is_dilated else ""))
    else:
        if groups in [0, nnef_input_shape[-1]] and d == 2:
            partial_convert_deconv_to_depthwise_conv2d_native_backprop_input(nnefop, converter)
        else:
            # TODO consider implementing this
            utils.print_error(
                "Grouped deconvolutions are only supported if they can be converted to "
                + "tf.nn.depthwise_conv2d_native_backprop_input.")
            return

def generic_convert_pooling(nnefop, converter):
    rank = len(dog.get_shape_safe(nnefop.args["input"]))

    if not nnefop.args["stride"]:
        nnefop.args["stride"] = [1 for _ in range(rank)]
    # if not nnefop.args["dilation"]:
    #     nnefop.args["dilation"] = [1 for _ in range(rank)]
    if nnefop.args["dilation"]:
        utils.print_error("Dilated pool is not supported for now")

    tfargname_input = "input" if nnefop.name in ["max_pool_with_index", "argmax_pool"] else "value"

    tfop = dog.OperationNode(nnefopname_to_tf(nnefop.name))
    tfop.add_arg(tfargname_input, converter.get_tfdn_safe(nnefop.args["input"]))
    tfop.add_arg("ksize", nnefop.args["size"])
    tfop.add_arg("strides", nnefop.args["stride"])
    tfop.add_arg("padding", nnefop.args["_tf_padding"])  # TODO padding?

    if nnefop.name == "argmax_pool":
        tfop.add_result("output", converter.make_tfdn(nnefop.result, discriminator="output"))
        tfop.add_result("argmax", converter.make_tfdn(nnefop.result))
    elif nnefop.name == "max_pool_with_index":
        tfop.add_result("output", converter.make_tfdn(nnefop.results["output"]))
        tfop.add_result("argmax", converter.make_tfdn(nnefop.results["index"]))
    else:
        tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop)

def generic_convert_pool_grad(nnefop, converter):
    if nnefop.name == "avg_pool_grad":
        rank = len(nnefop.args["orig_input_shape"])
    else:
        rank = len(dog.get_shape_safe(nnefop.args["orig_input"]))

    if not nnefop.args["stride"]:
        nnefop.args["stride"] = [1 for _ in range(rank)]
    if not nnefop.args["dilation"]:
        nnefop.args["dilation"] = [1 for _ in range(rank)]

    tfop = dog.OperationNode(nnefopname_to_tf(nnefop.name))
    if nnefop.name == "max_pool_grad_with_index":
        tfop.add_arg("input", converter.get_tfdn_safe(nnefop.args["orig_input"]))
        tfop.add_arg("grad", converter.get_tfdn_safe(nnefop.args["output_grad"]))
        tfop.add_arg("argmax", converter.get_tfdn_safe(nnefop.args["orig_index"]))
    elif nnefop.name == "max_pool_grad":
        tfop.add_arg("orig_input", converter.get_tfdn_safe(nnefop.args["orig_input"]))
        tfop.add_arg("orig_output", converter.get_tfdn_safe(nnefop.args["orig_output"]))
        tfop.add_arg("grad", converter.get_tfdn_safe(nnefop.args["output_grad"]))
    elif nnefop.name == "avg_pool_grad":
        tfop.add_arg("orig_input_shape", nnefop.args["orig_input_shape"])
        tfop.add_arg("grad", converter.get_tfdn_safe(nnefop.args["output_grad"]))
    else:
        assert False

    tfop.add_arg("ksize", nnefop.args["size"])
    tfop.add_arg("strides", spatial_nhwc(nnefop.args["stride"],
                                         in_spatial=False, out_spatial=False, empty_value=[1] * rank))
    tfop.add_arg("padding", nnefop.args["_tf_padding"])
    tfop.add_result("result", converter.make_tfdn(nnefop.results["input_grad"]))
    converter.add_tfop(tfop)

def generic_convert_upsample_downsample(nnefop, converter):
    nnef_input_shape = dog.get_shape_safe(nnefop.args["input"])
    nnef_factors = nnefop.args["factor"]
    is_downsample = "down" in nnefop.name
    output_size = [(int(i / f) if is_downsample else int(i * f)) for i, f in zip(nnef_input_shape[1:-1], nnef_factors)]

    tfop = dog.OperationNode(nnefopname_to_tf(nnefop.name))
    tfop.add_arg("images", converter.get_tfdn_safe(nnefop.args["input"]))
    tfop.add_arg("size", output_size)
    if nnefop.name == "multilinear_upsample":
        # TODO implement full border and method support
        if nnefop.args["border"].lower() != "replicate":
            utils.print_error("border={} is unsupported in multilinear_upsample.".format(nnefop.args["border"]))
        if nnefop.args["method"].lower() not in ["aligned", "asymmetric"]:
            utils.print_error("method={} is unsupported in multilinear_upsample.".format(nnefop.args["method"]))
        tfop.add_arg("align_corners", nnefop.args["method"].lower() == 'aligned')
    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop, nnefop)

def generic_convert_reduce(nnefop, converter):
    tfopname = nnefopname_to_tf(nnefop.name)
    if nnefop.name == "sum_reduce" and nnefop.args["normalize"]:
        tfopname = "tf.mean_reduce"

    tfinput = converter.get_tfdn_safe(nnefop.args["input"])
    if isinstance(tfinput, float):  # TF can't reduce a single value
        tfinput = [tfinput]

    tfop = dog.OperationNode(tfopname)
    tfop.add_arg("input_tensor", tfinput)
    tfop.add_arg("axis", nnefop.args["axes"])
    tfop.add_arg(tf_compat.reduce_arg_keepdims, True)
    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop, nnefop)

def generic_convert_argminmax_reduce(nnefop, converter):
    if len(nnefop.args["axes"]) != 1:
        utils.print_error("argmin/argmax is only supported for exactly one axis")
    axis = nnefop.args["axes"][0]

    tfinput = converter.get_tfdn_safe(nnefop.args["input"])
    if isinstance(tfinput, float):  # TF can't reduce a single value
        tfinput = [tfinput]

    tfop_argmax = dog.OperationNode(nnefopname_to_tf(nnefop.name))
    tfop_argmax.add_arg("input", tfinput)
    tfop_argmax.add_arg("axis", axis)
    tfop_argmax.add_result("result", converter.make_tfdn(nnefop.result, discriminator="squeezed"))
    converter.add_tfop(tfop_argmax, nnefop)

    tfop_expand_dims = dog.OperationNode("tf.expand_dims")
    tfop_expand_dims.add_arg("input", tfop_argmax.result)
    tfop_expand_dims.add_arg("axis", axis)
    tfop_expand_dims.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop_expand_dims, nnefop)

def convert_moments(nnefop, converter):
    tfop = dog.OperationNode("tf.nn.moments")
    tfop.add_arg("x", converter.get_tfdn_safe(nnefop.args["input"]))
    tfop.add_arg("axes", nnefop.args["axes"])
    tfop.add_arg("keep_dims", True)
    tfop.add_result("mean", converter.make_tfdn(nnefop.results["mean"]))
    tfop.add_result("variance", converter.make_tfdn(nnefop.results["variance"]))
    converter.add_tfop(tfop, nnefop)

def generic_convert_activation(nnefop, converter):
    tfop = dog.OperationNode(nnefopname_to_tf(nnefop.name))
    tfop.add_arg("features", converter.get_tfdn_safe(nnefop.args["x"]))
    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop, nnefop)

def convert_softmax(nnefop, converter):
    tfop = dog.OperationNode("tf.nn.softmax")
    tfop.add_arg("logits", converter.get_tfdn_safe(nnefop.args["x"]))
    rank = len(dog.get_shape_safe(nnefop.args["x"]))
    nnefaxes = nnefop.args["axes"]
    if len(nnefaxes) != 1:  # btw using rank-1 causes error on tf 1.0.0, while -1 is ok
        utils.print_error("In softmax 'axes' parameter is only supported if it has exactly one value in it.")
    if nnefaxes[0] != rank - 1:
        tfop.add_arg(tf_compat.softmax_arg_axis, nnefaxes[0])
    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop, nnefop)

def convert_leaky_relu(nnefop, converter):
    tfop = dog.OperationNode("tf.nn.leaky_relu")
    tfop.add_arg("features", converter.get_tfdn_safe(nnefop.args["x"]))
    tfop.add_arg("alpha", converter.get_tfdn_safe(nnefop.args["alpha"]))
    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop, nnefop)

def convert_prelu(nnefop, converter):
    convert_leaky_relu(nnefop, converter)

def convert_local_response_normalization(nnefop, converter):
    nnefsize = list(nnefop.args["size"])
    if len(nnefsize) < 2:
        utils.print_error("arg 'size' of local_response_normalization must have at least 2 elements")
        nnefsize = [1, 11]
    depth_size = nnefsize[-1]
    nnefsize[-1] = 1
    if utils.has_greater_than_1(nnefsize):
        utils.print_error("local_response_normalization only supported when only the last element of 'size' is > 1")

    depth_radius = (depth_size - 1) / 2
    alpha = nnefop.args["alpha"] / depth_size

    tfop = dog.OperationNode("tf.nn.local_response_normalization")
    tfop.add_arg("input", converter.get_tfdn_safe(nnefop.args["input"]))
    tfop.add_arg("depth_radius", depth_radius)
    tfop.add_arg("bias", nnefop.args["bias"])
    tfop.add_arg("alpha", alpha)
    tfop.add_arg("beta", nnefop.args["beta"])
    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop, nnefop)

def convert_l2_normalization(nnefop, converter):
    # TODO epsilon and bias are not really the same, is this a problem?
    tfop = dog.OperationNode("tf.nn.l2_normalize")
    tfop.add_arg("x", converter.get_tfdn_safe(nnefop.args["input"]))
    tfop.add_arg(tf_compat.l2_normalize_arg_axis, nnefop.args["axes"])
    tfop.add_arg("epsilon", nnefop.args["bias"])
    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop, nnefop)

def convert_batch_normalization(nnefop, converter):
    tfop = dog.OperationNode("tf.nn.batch_normalization ")
    tfop.add_arg("x", converter.get_tfdn_safe(nnefop.args["input"]))
    tfop.add_arg("mean", converter.get_tfdn_safe(nnefop.args["mean"]))
    tfop.add_arg("variance", converter.get_tfdn_safe(nnefop.args["variance"]))
    tfop.add_arg("offset", nnefop.args["offset"])
    tfop.add_arg("scale", nnefop.args["scale"])
    tfop.add_arg("variance_epsilon", nnefop.args["epsilon"])
    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop, nnefop)

def convert_copy_n(nnefop, converter):
    if len(nnefop.result) != nnefop.args["times"]:
        utils.print_error("copy_n must have 'times' outputs")

    for nnefdn_result in nnefop.result:
        tfop = dog.OperationNode("tf.identity")
        tfop.add_arg("input", converter.get_tfdn_safe(nnefop.args["x"]))
        tfop.add_result("result", converter.make_tfdn(nnefdn_result))
        converter.add_tfop(tfop, nnefop)

def convert_add_n(nnefop, converter):
    tfop = dog.OperationNode("tf.add_n")
    tfop.add_arg("inputs", [converter.get_tfdn_safe(nnefdn) for nnefdn in nnefop.args["x"]])
    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop, nnefop)

def convert__nnef_pad(nnefop, converter):
    tfop = dog.OperationNode("tf.pad")
    tfop.add_arg("tensor", converter.get_tfdn_safe(nnefop.args["input"]))
    tfop.add_arg("paddings", nnefop.args["padding"])
    tf_pad_mode = nnefborder_to_tf(nnefop.args["border"])
    if tf_pad_mode != "CONSTANT":
        tfop.add_arg("mode", tf_pad_mode)
    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop, nnefop.extra.get("original_op", nnefop))

def convert__nnef_bias_add(nnefop, converter):
    tfop = dog.OperationNode("tf.nn.bias_add")
    tfop.add_arg("value", converter.get_tfdn_safe(nnefop.args["x"]))
    tfop.add_arg("bias", converter.get_tfdn_safe(nnefop.args["y"]))
    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop, nnefop.extra.get("original_op", nnefop))

def convert_zeros(nnefop, converter):
    tfshape = nnefop.args["shape"]
    tfop = dog.OperationNode("tf.zeros")
    tfop.add_arg("shape", tfshape)
    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop, nnefop)

def convert_ones(nnefop, converter):
    tfshape = nnefop.args["shape"]
    tfop = dog.OperationNode("tf.ones")
    tfop.add_arg("shape", tfshape)
    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop, nnefop)

def convert_zeros_like(nnefop, converter):
    tfop = dog.OperationNode("tf.zeros_like")
    tfop.add_arg("tensor", converter.get_tfdn_safe(nnefop.args["reference"]))
    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop, nnefop)

def convert_ones_like(nnefop, converter):
    tfop = dog.OperationNode("tf.ones_like")
    tfop.add_arg("tensor", converter.get_tfdn_safe(nnefop.args["reference"]))
    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop, nnefop)

def convert_slice(nnefop, converter):
    rank = len(dog.get_shape_safe(nnefop.args["input"]))
    axes = nnefop.args["axes"]
    begin = nnefop.args["begin"]
    end = nnefop.args["end"]

    size = [
        -1 if e == -1 else e - b
        for b, e in zip(begin, end)
    ]

    tfbegin, tfsize = utils.zip_inverse(2, [
        (begin[axes.index(i)], size[axes.index(i)]) if i in axes else (0, -1)
        for i in range(rank)
    ])

    tfop = dog.OperationNode("tf.slice")
    tfop.add_arg("input_", converter.get_tfdn_safe(nnefop.args["input"]))
    tfop.add_arg("begin", tfbegin)
    tfop.add_arg("size", tfsize)
    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop, nnefop)

def convert_stack(nnefop, converter):
    nnefdns = nnefop.args["values"]
    tfop = dog.OperationNode("tf.stack")
    tfop.add_arg("values", [converter.get_tfdn_safe(nnefdn) for nnefdn in nnefdns])
    tfaxis = nnefop.args['axis']
    tfop.add_arg("axis", tfaxis)
    tfop.add_result("result", converter.make_tfdn(nnefop.result))
    converter.add_tfop(tfop, nnefop)

def convert_unstack(nnefop, converter):
    nnefaxis = nnefop.args["axis"]
    nnefdn_value = nnefop.args["value"]
    nnefshape = dog.get_shape_safe(nnefdn_value)

    tfop = dog.OperationNode("tf.unstack")
    tfop.add_arg("value", converter.get_tfdn_safe(nnefdn_value))
    tfop.add_arg("num", len(nnefop.result))
    tfop.add_arg("axis", nnefaxis)

    for i in range(len(nnefop.result)):
        tfop.add_result("result{}".format(i), converter.make_tfdn(nnefop.result[i]))

    converter.add_tfop(tfop, nnefop)

def convert_box(nnefop, converter):
    # only padding
    if (not utils.has_not_equal_1(nnefop.args["size"])
            and nnefop.args["border"] in ["constant", "reflect", "reflect-even"]
            and not utils.has_not_equal_1(nnefop.args["stride"])
            and not utils.has_not_equal_1(nnefop.args["dilation"])):
        tfop = dog.OperationNode("tf.pad")
        tfop.add_arg("tensor", converter.get_tfdn_safe(nnefop.args["input"]))
        tfop.add_arg("paddings", nnefop.args["padding"])
        tfmode = nnefborder_to_tf(nnefop.args["border"])
        if tfmode != 'CONSTANT':
            tfop.add_arg("mode", tfmode)
        tfop.add_result("result", converter.make_tfdn(nnefop.result))
        converter.add_tfop(tfop, nnefop)
    else:
        utils.print_error("Box is not yet fully supported, your use is invalid.")

def generic_convert_unsupported(nnefop, converter):
    utils.print_error("{} is not supported".format(nnefop.name))

generic_converters = {
    "add": (generic_convert_binary, "tf.add"),
    "sub": (generic_convert_binary, "tf.subtract"),
    "mul": (generic_convert_binary, "tf.multiply"),
    "div": (generic_convert_binary, "tf.divide"),
    "pow": (generic_convert_binary, "tf.pow"),
    "lt": (generic_convert_binary, "tf.less"),
    "gt": (generic_convert_binary, "tf.greater"),
    "le": (generic_convert_binary, "tf.less_equal"),
    "ge": (generic_convert_binary, "tf.greater_equal"),
    "eq": (generic_convert_binary, "tf.equal"),
    "ne": (generic_convert_binary, "tf.not_equal"),
    "and": (generic_convert_binary, "tf.logical_and"),
    "or": (generic_convert_binary, "tf.logical_or"),
    "min": (generic_convert_binary, "tf.minimum"),
    "max": (generic_convert_binary, "tf.maximum"),
    "exp": (generic_convert_unary, "tf.exp"),
    "log": (generic_convert_unary, "tf.log"),
    "abs": (generic_convert_unary, "tf.abs"),
    "sign": (generic_convert_unary, "tf.sign"),
    "rcp": (generic_convert_unary, "tf.reciprocal"),
    "neg": (generic_convert_unary, "tf.negative"),
    "floor": (generic_convert_unary, "tf.floor"),
    "ceil": (generic_convert_unary, "tf.ceil"),
    "round": (generic_convert_unary, "tf.round"),
    "sqr": (generic_convert_unary, "tf.square"),
    "sqrt": (generic_convert_unary, "tf.sqrt"),
    "rsqrt": (generic_convert_unary, "tf.rsqrt"),
    "not": (generic_convert_unary, "tf.logical_not"),
    "sigmoid": (generic_convert_unary, "tf.sigmoid"),
    "sinh": (generic_convert_unary, "tf.sinh"),
    "cosh": (generic_convert_unary, "tf.cosh"),
    "tanh": (generic_convert_unary, "tf.tanh"),
    "nearest_downsample": (generic_convert_upsample_downsample, "tf.image.resize_nearest_neighbor"),
    "area_downsample": (generic_convert_upsample_downsample, "tf.image.resize_area"),
    "nearest_upsample": (generic_convert_upsample_downsample, "tf.image.resize_nearest_neighbor"),
    "multilinear_upsample": (generic_convert_upsample_downsample, "tf.image.resize_bilinear"),
    "sum_reduce": (generic_convert_reduce, "tf.reduce_sum"),
    "min_reduce": (generic_convert_reduce, "tf.reduce_min"),
    "max_reduce": (generic_convert_reduce, "tf.reduce_max"),
    "mean_reduce": (generic_convert_reduce, "tf.reduce_mean"),
    "argmax_reduce": (generic_convert_argminmax_reduce, "tf.argmax"),
    "argmin_reduce": (generic_convert_argminmax_reduce, "tf.argmin"),
    "elu": (generic_convert_activation, "tf.nn.elu"),
    "relu": (generic_convert_activation, "tf.nn.relu"),
    "softplus": (generic_convert_activation, "tf.nn.softplus"),
    "argmax_pool": (generic_convert_pooling, "tf.nn.max_pool_with_argmax"),
    "max_pool_with_index": (generic_convert_pooling, "tf.nn.max_pool_with_argmax"),
    "max_pool": (generic_convert_pooling, "tf.nn.max_pool"),
    "avg_pool": (generic_convert_pooling, "tf.nn.avg_pool"),
    "max_pool_grad": (generic_convert_pool_grad, tf_compat.gen_nn_ops_max_pool_grad_name),
    "max_pool_grad_with_index": (generic_convert_pool_grad, tf_compat.gen_nn_ops_max_pool_grad_with_argmax_name),
    "avg_pool_grad": (generic_convert_pool_grad, tf_compat.gen_nn_ops_avg_pool_grad_name),
    "debox": (generic_convert_unsupported, ""),
    "sample": (generic_convert_unsupported, ""),
    "desample": (generic_convert_unsupported, ""),
    "local_mean_normalization": (generic_convert_unsupported, ""),
    "local_variance_normalization": (generic_convert_unsupported, ""),
    "local_contrast_normalization": (generic_convert_unsupported, ""),
    "linear_quantize": (generic_convert_unsupported, ""),
    "logarithmic_quantize": (generic_convert_unsupported, ""),
    "binary_quantize": (generic_convert_unsupported, ""),
    "ternary_quantize": (generic_convert_unsupported, "")
}

def nnefopname_to_tf(nnefopname):
    return generic_converters[nnefopname][1]

DefaultConverters = {fun.__name__[len("convert_"):]: fun for fun in utils.get_functions(prefix="convert_")}
for k, v in generic_converters.items():
    DefaultConverters[k] = v[0]


import numpy

from .spatialfilters import Xdawn
from .utils.covariance import covariances, covariances_EP, cospectrum
from sklearn.base import BaseEstimator, TransformerMixin

###############################################################################

class Covariances(BaseEstimator, TransformerMixin):

    """
    compute the covariances matrices

    """

    def __init__(self, estimator='scm'):
        self.estimator = estimator

    def fit(self, X, y=None):
        return self

    def transform(self, X):

        covmats = covariances(X, estimator=self.estimator)
        return covmats

    def fit_transform(self, X, y=None):
        return self.transform(X)

###############################################################################

class ERPCovariances(BaseEstimator, TransformerMixin):

    """
    Compute special form ERP cov mat

    """

    def __init__(self, classes=None, estimator='scm', svd=None):
        self.classes = classes
        self.estimator = estimator
        self.svd = svd

        if svd is not None:
            if not isinstance(svd,int):
                raise TypeError('svd must be None or int')

    def fit(self, X, y):

        if self.classes is not None:
            classes = self.classes
        else:
            classes = numpy.unique(y)

        self.P = []
        for c in classes:
            # Prototyped responce for each class
            P = numpy.mean(X[y == c, :, :], axis=0)

            # Apply svd if requested
            if self.svd is not None:
                U, s, V = numpy.linalg.svd(P)
                P = numpy.dot(U[:, 0:self.svd].T,P)

            self.P.append(P)

        self.P = numpy.concatenate(self.P, axis=0)
        return self

    def transform(self, X):

        covmats = covariances_EP(X, self.P, estimator=self.estimator)
        return covmats

    def fit_transform(self, X, y):
        self.fit(X, y)
        return self.transform(X)

###############################################################################

class XdawnCovariances(BaseEstimator, TransformerMixin):

    """
    Compute xdawn, project the signal and compute the covariances

    """

    def __init__(self, nfilter=4, applyfilters=True, classes=None,
                 estimator='scm'):
        self.Xd = Xdawn(nfilter=nfilter, classes=classes)
        self.applyfilters = applyfilters
        self.estimator = estimator

    def fit(self, X, y):
        self.Xd.fit(X, y)
        return self

    def transform(self, X):
        if self.applyfilters:
            X = self.Xd.transform(X)

        covmats = covariances_EP(X, self.Xd.evokeds_, estimator=self.estimator)
        return covmats

    def fit_transform(self, X, y):
        self.fit(X, y)
        return self.transform(X)

###############################################################################

class CospCovariances(BaseEstimator, TransformerMixin):

    """
    compute the cospectral covariance matrices

    """

    def __init__(
            self,
            window=128,
            overlap=0.75,
            fmin=None,
            fmax=None,
            fs=None,
            phase_correction=False):
        self._window = _nextpow2(window)
        self._overlap = overlap
        self._fmin = fmin
        self._fmax = fmax
        self._fs = fs

        self._phase_corr = phase_correction

    def fit(self, X, y=None):
        return self

    def transform(self, X):

        Nt, Ne, _ = X.shape
        out = []

        for i in range(Nt):
            S = cospectrum(X[i], window=self._window, overlap=self._overlap,
                           fmin=self._fmin, fmax=self._fmax, fs=self._fs)
            out.append(S.real)

        return numpy.array(out)

    def fit_transform(self, X, y=None):
        return self.transform(X)

##########################################################################
def _nextpow2(i):
    n = 1
    while n < i:
        n *= 2
    return n


import os
import datetime
import IPython
import IPython.display
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
import tensorflow as tf
from WindowGenerator import *

mpl.rcParams['figure.figsize'] = (8, 6)
mpl.rcParams['axes.grid'] = False

df = pd.read_csv("orders_autumn_2020.txt")

# Combine all events into days
# and datetime conversion from string datatype to useful signal data
date_time = pd.to_datetime(df['TIMESTAMP'], format='%Y-%m-%d %H:%M:%S')
df['TIMESTAMP'] = date_time
groupby = df.groupby(pd.Grouper(key = 'TIMESTAMP', freq = 'D'))
df = groupby.mean()
df['COUNTS'] = groupby.size()

timestamp = df.index.map(datetime.datetime.timestamp)
year = (365.2425)*24*60*60
df['Year sin'] = np.sin(timestamp * (2 * np.pi / year))
df['Year cos'] = np.cos(timestamp * (2 * np.pi / year))

# This model will not incorporate location data so we can remove these features
df = df.drop(columns = ['USER_LONG','USER_LAT','VENUE_LONG','VENUE_LAT'])

#Plot some interesting features
#Plot the data
plt.figure(figsize = (12, 8))
plt.suptitle("FEATURE VISUALIZATION")
plot_cols = ['COUNTS', 'PRECIPITATION','ITEM_COUNT']
plot_features = df[plot_cols]
_ = plot_features.plot(subplots=True)

# Plotting feature correlation to identify any strong correlations in data.
# No badly correlating features found.
plt.figure(figsize = (12, 8))
plt.suptitle("FEATURE CORRELATION")
cor = df.corr()
sns.heatmap(cor,annot=True,cmap=plt.cm.Reds)
plt.show()

# Splitting data into training, validation, and testing sets.
# NOTE: This testing set is not directly used in this model and it is extracted later.
# The training and validation data have 20% overlap, because the daily training data is limited.
column_dict = {name : i for i, name in enumerate(df.columns)}
n = len(df)
train_df  = df[0:int(n*0.7)]
val_df = df[int(n*0.5):]
test_df = df[int(n*0.5):]

num_features = df.shape[1]
df.fillna(df.mean(),inplace = True)

#Normalize the training data
training_mean = train_df.mean()
training_std = train_df.std()

train_df = (train_df - training_mean) / training_std
val_df = (val_df - training_mean) / training_std
test_df = (test_df - training_mean) / training_std

#Visualize the distribution of features
df_std = (df - training_mean) / training_std
df_std = df_std.select_dtypes([np.number])
df_std = df_std.melt(var_name = "Features", value_name = "Normalized")
plt.figure(figsize=(12, 6))
ax = sns.violinplot(x = "Features", y = "Normalized", data = df_std)
_ = ax.set_xticklabels(df.keys(), rotation=90)

# Initialize window for training using WindowGenerator class.
IN_STEPS = 23
OUT_STEPS = 7
WINDOW = IN_STEPS + OUT_STEPS
multi_window = WindowGenerator(input_width = IN_STEPS,
                               label_width = OUT_STEPS,
                               shift = OUT_STEPS,
                               train_df = train_df,
                               val_df = val_df,
                               test_df = test_df)

CONV_WIDTH = 3
multi_conv_model = tf.keras.Sequential([
    # Shape [batch, time, features] => [batch, CONV_WIDTH, features]
    tf.keras.layers.Lambda(lambda x: x[:, -CONV_WIDTH:, :]),
    # Shape => [batch, 1, conv_units]
    tf.keras.layers.Conv1D(256, activation='relu', kernel_size=(CONV_WIDTH)),
    # Shape => [batch, 1,  out_steps*features]
    tf.keras.layers.Dense(OUT_STEPS*num_features,
                          kernel_initializer=tf.initializers.zeros),
    # Shape => [batch, out_steps, features]
    tf.keras.layers.Reshape([OUT_STEPS, num_features])
])

# Train the data and plot the validation results using WindowGenerator.
history = compile_and_fit(multi_conv_model, multi_window)
IPython.display.clear_output()
multi_window.plot(multi_conv_model, plot_col = 'COUNTS')

# Make the forecast by sampling the latest test data.
data = np.array(test_df[-IN_STEPS:],dtype=np.float32)
ds = tf.keras.preprocessing.timeseries_dataset_from_array(
        data = data,
        targets = None,
        sequence_length = IN_STEPS,
        sequence_stride = 1,
        shuffle = True,
        batch_size = 32)

early_stopping = tf.keras.callbacks.EarlyStopping(monitor='val_loss', patience=2, mode='min')
predictions = multi_conv_model.predict(ds, verbose = 0, callbacks=[early_stopping])
j = df.columns.get_loc('COUNTS')
count_predictions = tf.reshape(predictions[0,:,j],[1,OUT_STEPS]).numpy()*training_std[j]+training_mean[j]

#Plot the predicted values
plt.figure(figsize = (12, 8))
plt.suptitle('PREDICTION FOR NEXT 7 DAYS')
plt.plot((test_df['COUNTS'][-IN_STEPS:]).apply(lambda x: x*training_std[j] + training_mean[j]), label='Inputs', marker='.', zorder= -10)
prediction_indices = pd.date_range(start= test_df.index[-1] + pd.DateOffset(1), periods = OUT_STEPS, freq = 'D')
plt.scatter(prediction_indices, count_predictions,
                  marker='X', edgecolors='k', label='Predictions',
                  c='#ff7f0e', s=64)
plt.xticks(rotation=90)
plt.legend()
plt.xlabel('Date')

#Plot the training and validation loss for each epoch.
plt.figure(figsize = (12, 8))
plt.plot(history.history['loss'], label='MAE (training data)')
plt.plot(history.history['val_loss'], label='MAE (validation data)')
plt.ylabel('MAE value')
plt.xlabel('No. epoch')
plt.legend(loc="upper left")
plt.show()