from typing import List, Optional

# External Dependencies:
from aws_cdk import CfnOutput, CfnParameter, Duration, RemovalPolicy, Stack
from aws_cdk.aws_iam import (
    ManagedPolicy,
    PolicyDocument,
)
import aws_cdk.aws_s3 as s3
import aws_cdk.aws_ssm as ssm
from constructs import Construct

# Local Dependencies:
from annotation import AnnotationInfra
from pipeline import ProcessingPipeline
from pipeline.iam_utils import (
    S3Statement,
    SsmParameterReadStatement,
    StateMachineExecuteStatement,
)

class PipelineDemoStack(Stack):
    """Deployable CDK stack bundling the core OCR pipeline construct with supporting demo resources

    This stack bundles the core ProcessingPipeline construct with the additional resources required
    to deploy and use it for the demo: Such as the project ID SSM parameter, input data bucket,
    SageMaker permissions policy, and the infrastructure for custom annotation UIs in SageMaker
    Ground Truth.

    It also creates several CloudFormation stack outputs to help users find important created
    resources.
    """

    def __init__(
        self,
        scope: Construct,
        construct_id: str,
        default_project_id: str,
        use_thumbnails: bool,
        enable_sagemaker_autoscaling: bool = False,
        build_sagemaker_ocrs: List[str] = [],
        deploy_sagemaker_ocrs: List[str] = [],
        use_sagemaker_ocr: Optional[str] = None,
        **kwargs,
    ) -> None:
        """Create a PipelineDemoStack

        Parameters
        ----------
        scope :
            As per aws_cdk.Stack
        construct_id :
            As per aws_cdk.Stack
        default_project_id :
            The `ProjectId` is a CFn stack parameter that prefixes created SSM parameters and
            allows SageMaker notebooks to look up the parameters for the deployed stack. If you're
            deploying straight from `cdk deploy`, then the value you specify here will be used. If
            you're `cdk synth`ing a CloudFormation template, then this will be the default value
            for the ProjectId parameter.
        use_thumbnails :
            Set `True` to build the stack with support for visual (page thumbnail image) model
            input features, or `False` to omit the thumbnailing step. Pipelines deployed with
            `use_thumbnails=True` will fail if a thumbnailer endpoint is not set up (see SageMaker
            notebooks). Pipelines deployed with `use_thumbnails=False` cannot fully utilize model
            architectures that use page images for inference (such as LayoutLMv2+, etc).
        enable_sagemaker_autoscaling :
            Set True to enable auto-scale-to-zero on any SageMaker endpoints created by the stack.
            Turning this on should improve cost-efficiency for workloads which are often idle, but
            will introduce cold-start delays to affected stages of the pipeline so may not be ideal
            during development. This setting does not affect endpoints created *outside* the stack
            and later plumbed in to the pipeline (i.e. endpoints deployed from notebooks).
        build_sagemaker_ocrs :
            List of alternative (SageMaker-based) OCR engine names to build container images and
            SageMaker Models for in the deployed stack. By default ([]), none will be included. See
            `CUSTOM_OCR_ENGINES` in pipeline/ocr/sagemaker_ocr.py for supported engines.
        deploy_sagemaker_ocrs :
            List of alternative OCR engine names to deploy SageMaker endpoints for in the stack. Any
            names in here must also be included in `build_sagemaker_ocrs`. Default []: Support
            Amazon Textract OCR only.
        use_sagemaker_ocr :
            Optional alternative OCR engine name to use in the deployed document pipeline. If set
            and not empty, this must also be present in `build_sagemaker_ocrs` and
            `deploy_sagemaker_ocrs`. Default None: Use Amazon Textract for initial document OCR.
        **kwargs :
            As per aws_cdk.Stack
        """
        super().__init__(scope, construct_id, **kwargs)

        # Could consider just directly using the stack ID for this, but then if you were to vend
        # the stack through e.g. AWS Service Catalog you may not have control over setting a nice
        # readable stack ID:
        self.project_id_param = CfnParameter(
            self,
            "ProjectId",
            allowed_pattern=r"[a-zA-Z0-9]+(\-[a-zA-Z0-9]+)*",
            constraint_description="Alphanumeric with internal hyphens allowed",
            default=default_project_id,
            description=(
                "ID to look up this stack's resources from SageMaker notebooks, used in folder "
                "prefixes for SSM parameters."
            ),
            max_length=25,
            min_length=3,
        )

        self.annotation_infra = AnnotationInfra(self, "AnnotationInfra")

        self.input_bucket = s3.Bucket(
            self,
            "PipelineInputBucket",
            auto_delete_objects=True,
            block_public_access=s3.BlockPublicAccess(
                block_public_acls=True,
                block_public_policy=True,
                ignore_public_acls=True,
                restrict_public_buckets=True,
            ),
            encryption=s3.BucketEncryption.S3_MANAGED,
            enforce_ssl=True,
            lifecycle_rules=[
                s3.LifecycleRule(enabled=True, expiration=Duration.days(7)),
            ],
            removal_policy=RemovalPolicy.DESTROY,
            cors=[
                # CORS permissions are required for the A2I human review UI to retrieve objects:
                s3.CorsRule(
                    allowed_headers=["*"],
                    allowed_methods=[s3.HttpMethods.GET],
                    allowed_origins=[
                        "https://mturk-console-template-preview-hooks.s3.amazonaws.com",
                    ],
                ),
            ],
        )
        self.pipeline = ProcessingPipeline(
            self,
            "ProcessingPipeline",
            input_bucket=self.input_bucket,
            ssm_param_prefix=f"/{self.project_id_param.value_as_string}/config/",
            use_thumbnails=use_thumbnails,
            enable_sagemaker_autoscaling=enable_sagemaker_autoscaling,
            build_sagemaker_ocrs=build_sagemaker_ocrs,
            deploy_sagemaker_ocrs=deploy_sagemaker_ocrs,
            use_sagemaker_ocr=use_sagemaker_ocr,
        )
        self.data_science_policy = ManagedPolicy(
            self,
            "PipelineDataSciencePolicy",
            document=PolicyDocument(
                statements=[
                    S3Statement(
                        grant_write=True,
                        resources=[self.input_bucket],
                        sid="ReadWritePipelineInputBucket",
                    ),
                    StateMachineExecuteStatement(
                        resources=[self.pipeline.plain_textract_state_machine],
                        sid="RunPlainTextractStateMachine",
                    ),
                ]
                + self.annotation_infra.get_data_science_policy_statements()
                + self.pipeline.config_read_write_statements()
                # In the notebooks we'll use the same execution role for the trained model/endpoint
                # as the notebook itself runs with - so need to grant the role the required perms
                # for reading/writing relevant S3 buckets and publishing to SNS in the pipeline:
                + self.pipeline.sagemaker_model_statements(),
            ),
        )

        # We'd like the stack to push some useful outputs to CloudFormation, and defining them here
        # rather than in the lower-level constructs will keep the constructs flexible/reusable.
        #
        # We override the auto-generated logical IDs to make the names simple to find in console.
        self.data_science_policy_output = CfnOutput(
            self,
            "DataSciencePolicyName",
            description=(
                "Name of the IAM policy with permissions needed for the SageMaker notebooks to "
                "access this stack's resources. Add this policy to your SageMaker execution role."
            ),
            value=self.data_science_policy.managed_policy_name,
        )
        self.data_science_policy_output.override_logical_id("DataSciencePolicyName")
        self.input_bucket_name_output = CfnOutput(
            self,
            "InputBucketName",
            description="Name of the S3 bucket to which input documents should be uploaded",
            value=self.pipeline.input_bucket.bucket_name,
        )
        self.input_bucket_name_output.override_logical_id("InputBucketName")
        self.pipeline_statemachine_output = CfnOutput(
            self,
            "PipelineStateMachine",
            description="ARN of the State Machine for the end-to-end OCR pipeline",
            value=self.pipeline.state_machine.state_machine_arn,
        )
        self.pipeline_statemachine_output.override_logical_id("PipelineStateMachine")
        self.textract_statemachine_output = CfnOutput(
            self,
            "PlainTextractStateMachine",
            description="ARN of the State Machine for *only* running Textract (no enrichments)",
            value=self.pipeline.plain_textract_state_machine.state_machine_arn,
        )
        self.textract_statemachine_output.override_logical_id("PlainTextractStateMachine")
        self.model_param_output = CfnOutput(
            self,
            "SageMakerEndpointParamName",
            description="SSM parameter to configure the pipeline's SageMaker endpoint name",
            value=self.pipeline.sagemaker_endpoint_param.parameter_name,
        )
        self.model_param_output.override_logical_id("SageMakerEndpointParamName")
        self.thumbnail_param_output = CfnOutput(
            self,
            "ThumbnailEndpointParamName",
            description=(
                "SSM parameter to configure the pipeline's Thumbnail generation endpoint name"
            ),
            value="undefined"
            if self.pipeline.thumbnail_endpoint_param is None
            else self.pipeline.thumbnail_endpoint_param.parameter_name,
        )
        self.thumbnail_param_output.override_logical_id("ThumbnailEndpointParamName")
        self.entity_config_param_output = CfnOutput(
            self,
            "EntityConfigParamName",
            description=(
                "JSON configuration describing the field types to be extracted by the pipeline"
            ),
            value=self.pipeline.entity_config_param.parameter_name,
        )
        self.entity_config_param_output.override_logical_id("EntityConfigParamName")
        self.a2i_role_arn_output = CfnOutput(
            self,
            "A2IHumanReviewExecutionRoleArn",
            description="ARN of the execution Role to use for Amazon A2I human review workflows",
            value=self.pipeline.review_a2i_role.role_arn,
        )
        self.a2i_role_arn_output.override_logical_id("A2IHumanReviewExecutionRoleArn")
        self.workflow_param_output = CfnOutput(
            self,
            "A2IHumanReviewFlowParamName",
            description="SSM parameter to configure the pipeline's A2I review workflow ARN",
            value=self.pipeline.review_workflow_param.parameter_name,
        )
        self.workflow_param_output.override_logical_id("A2IHumanReviewFlowParamName")
        self.reviews_bucket_name_output = CfnOutput(
            self,
            "A2IHumanReviewBucketName",
            description="Name of the S3 bucket to which A2I reviews should be stored",
            value=self.pipeline.human_reviews_bucket.bucket_name,
        )
        self.reviews_bucket_name_output.override_logical_id("A2IHumanReviewBucketName")

        # While these CloudFormation outputs are nice for the CFn console, we'd also like to be
        # able to automatically look up project resources from SageMaker notebooks. To support
        # this, we'll create additional SSM params used just to *retrieve* static attributes of the
        # stack - rather than configuration points like the ProcessingPipeline construct's params.
        static_param_prefix = f"/{self.project_id_param.value_as_string}/static"
        self.sm_image_build_role_ssm_param = ssm.StringParameter(
            self,
            "SMImageBuildRoleSSMParam",
            string_value=self.annotation_infra.sm_image_build_role.role_name,
            description=(
                "Name of the CodeBuild execution role to use in SMStudio Image Build commands"
            ),
            parameter_name=f"{static_param_prefix}/SMDockerBuildRole",
            simple_name=False,
        )
        self.preproc_image_param = ssm.StringParameter(
            self,
            "PreprocImageSSMParam",
            description="URI of the thumbnail generator container image pre-created by the stack",
            parameter_name=f"{static_param_prefix}/PreprocImageURI",
            simple_name=False,
            string_value=self.pipeline.preproc_image.image_uri,
        )
        self.input_bucket_ssm_param = ssm.StringParameter(
            self,
            "InputBucketNameSSMParam",
            string_value=self.input_bucket.bucket_name,
            description="Name of the S3 bucket to which input documents should be uploaded",
            parameter_name=f"{static_param_prefix}/InputBucket",
            simple_name=False,
        )
        self.reviews_bucket_ssm_param = ssm.StringParameter(
            self,
            "ReviewsBucketNameSSMParam",
            string_value=self.pipeline.human_reviews_bucket.bucket_name,
            description="Name of the S3 bucket to which human reviews should be stored",
            parameter_name=f"{static_param_prefix}/ReviewsBucket",
            simple_name=False,
        )
        self.pipeline_statemachine_ssm_param = ssm.StringParameter(
            self,
            "PipelineStateMachineSSMParam",
            string_value=self.pipeline.state_machine.state_machine_arn,
            description="ARN of the State Machine for the end-to-end OCR pipeline",
            parameter_name=f"{static_param_prefix}/PipelineStateMachine",
            simple_name=False,
        )
        self.textract_statemachine_ssm_param = ssm.StringParameter(
            self,
            "PlainTextractStateMachineSSMParam",
            string_value=self.pipeline.plain_textract_state_machine.state_machine_arn,
            description="ARN of the State Machine for *only* running Textract (no enrichments)",
            parameter_name=(
                f"/{self.project_id_param.value_as_string}/static/PlainTextractStateMachine"
            ),
            simple_name=False,
        )
        self.enrichment_results_bucket_ssm_param = ssm.StringParameter(
            self,
            "EnrichmentModelResultsBucketSSMParam",
            description=(
                "Name of the S3 bucket to which SageMaker (async) model results should be stored"
            ),
            parameter_name=f"{static_param_prefix}/ModelResultsBucket",
            simple_name=False,
            string_value=self.pipeline.enriched_results_bucket.bucket_name,
        )
        self.thumbnails_callback_topic_ssm_param = ssm.StringParameter(
            self,
            "ThumbnailsCallbackTopicSSMParam",
            description="ARN of the SNS Topic to use for thumbnail images generation callback",
            parameter_name=f"{static_param_prefix}/ThumbnailsCallbackTopicArn",
            simple_name=False,
            string_value=(
                self.pipeline.thumbnail_sns_topic.topic_arn
                if self.pipeline.thumbnail_sns_topic
                else "undefined"
            ),
        )
        self.enrichment_callback_topic_ssm_param = ssm.StringParameter(
            self,
            "EnrichmentModelCallbackTopicSSMParam",
            description="ARN of the SNS Topic to use for callback in SageMaker Async Inference",
            parameter_name=f"{static_param_prefix}/ModelCallbackTopicArn",
            simple_name=False,
            string_value=(
                self.pipeline.sagemaker_sns_topic.topic_arn
                if self.pipeline.sagemaker_sns_topic
                else "undefined"
            ),
        )
        self.a2i_role_arn_param = ssm.StringParameter(
            self,
            "A2IExecutionRoleArnParam",
            string_value=self.pipeline.review_a2i_role.role_arn,
            description="ARN of the execution role which A2I human review workflows should use",
            parameter_name=f"{static_param_prefix}/A2IExecutionRoleArn",
            simple_name=False,
        )

        self.data_science_policy.add_statements(
            SsmParameterReadStatement(
                resources=[
                    self.sm_image_build_role_ssm_param,
                    self.input_bucket_ssm_param,
                    self.reviews_bucket_ssm_param,
                    self.pipeline_statemachine_ssm_param,
                    self.preproc_image_param,
                    self.textract_statemachine_ssm_param,
                    self.thumbnails_callback_topic_ssm_param,
                    self.enrichment_callback_topic_ssm_param,
                    self.enrichment_results_bucket_ssm_param,
                    self.a2i_role_arn_param,
                ],
                sid="ReadStaticPipelineParams",
            ),
        )


import json
import time
import random
import numpy
from typing import Union
from cv2 import absdiff
from neopixel import NeoPixels, Color
# from communication.esp_serial import esp
from webcam import Webcam, read_im, store_im
from config import storage, NUM_PIXELS, NUM_SNAP_FRAMES

np = NeoPixels(NUM_PIXELS)

class Snapper:
    """
    Snapper is used to create all images needed to measure 1 side of the tree.
    Snapper.run() is used to create a series of images.
    First each pixel in the tree is assigned a unique number. This number is translated to a binary number.
    Next for each bit the pixel is turned on in case the bit is high, turned off if the bit is low.
    When making images, first a base image is created with all leds off. The sequence of images representing the bit
    is shot after, and from each image the base image is subtracted to make maximum contrast betwwen led on and led off.
    """
    def __init__(self, series_name: str, snap_color: Union[Color, str] = 'white', num_pixels=NUM_PIXELS):
        self.series_name = series_name
        self.snap_color = snap_color
        self.num_pixels = num_pixels

        (storage / series_name).mkdir(exist_ok=True, parents=True)

    def snap_monochrome(self, name: str, color: Union[Color, str]):
        """
        Create an image in which all the pixels have the same color.
        :param name: Name of the image
        :param color: 1 Color to use for all the pixels in the tree.
        :return:
        """
        for pixel in NeoPixels.pixels:
            pixel.color = color

        # esp.write()
        NeoPixels.send_by_wifi()
        # Hardcoded sleep for now. This sleep could be replaced by some code calculating
        # and waiting the exact time it takes to render the frame.
        time.sleep(0.2)
        store_im(f'{self.series_name}/{name}', Webcam.snap())

    def snap_ground(self):
        """Snap a pixture with all leds off."""
        self.snap_monochrome('ground', 'black')

    def snap_full_on(self):
        """Snap a pixture with all leds on"""
        self.snap_monochrome('all', Color(100, 100, 100))

    def snap_leds(self):
        """Snap N pixtures. In each pixture all leds turn either on or off, depending on their "foto number".
        The foto number is a "random" number each led is assigned.
        With the N pixures, each led will flash a binary pattern.
        This pattern is the binary representation of the foto number of the led.
        :return:
        """
        for n in range(NUM_SNAP_FRAMES):
            for pixel in NeoPixels.pixels:
                if pixel.bit_n_set(n):
                    pixel.color = self.snap_color
                else:
                    pixel.color = 'black'
            # esp.write()
            NeoPixels.send_by_wifi()
            NeoPixels.send_by_wifi()
            time.sleep(0.2)
            store_im(f'{self.series_name}/{n}', Webcam.snap())

    def make_example(self):
        """
        Create an example image showing all snapped foto's in 1 images
        :return:
        """
        ground = read_im(f'{self.series_name}/ground')
        images = [absdiff(read_im(f'{self.series_name}/{n}'), ground) for n in range(8)]

        store_im(f'{self.series_name}/example', numpy.concatenate(
            (numpy.concatenate((images[0], images[1], images[2], images[3]), axis=1),
             numpy.concatenate((images[4], images[5], images[6], images[7]), axis=1)), axis=0))

    def assign_pixel_id(self):
        """
        Assign each pixel a unique random number
        """
        pixel = 0
        id_dict = {}
        for num in random.sample(range(16, 2**NUM_SNAP_FRAMES), self.num_pixels):
            NeoPixels.pixels[pixel].foto_number = num
            id_dict[num] = pixel
            pixel += 1

        with (storage / self.series_name / 'numbers.txt').open('w') as number_file:
            number_file.write(json.dumps(id_dict))

    def run(self):
        """
        Run the sequence needed to measure 1 side of the tree
        """
        self.assign_pixel_id()
        self.snap_ground()
        self.snap_full_on()
        self.snap_leds()
        self.make_example()

if __name__ == "__main__":
    Snapper('slaapkamer5').run()


import sys

# check for compatible Python version
if sys.version_info.major < 3:
	print('\nERROR: this script requires Python 3, using {}'.format(sys.version.split()[0]))
	sys.exit(1)

import errno, os, subprocess
from subprocess import CalledProcessError, DEVNULL
from PIL import Image

# FIXME: use "PIL" or "wand" module (ImageMagick) for converting images

WIN32 = sys.platform == 'win32'
exe = os.path.basename(sys.argv[0])
if WIN32:
	dir_temp = os.getenv('TEMP')
else:
	dir_temp = '/tmp'

def showUsage():
	usage = '\nUsage:\n\t{} <source> <target>'.format(exe)
	usage += '\n\nArguments:\n\n\tsource:  Input image.\n\ttarget:  Output image'
	print(usage)

def getCommand(cmd_name):
	cmd = None

	try:
		if WIN32:
			cmd = subprocess.check_output(('where', cmd_name,), stderr=DEVNULL).decode('utf-8').strip(' \t\r\n')
			# user first executable found
			cmd = cmd.split('\r\n')[0]
		else:
			cmd = subprocess.check_output(('which', cmd_name,), stderr=DEVNULL).decode('utf-8').strip(' \t\r\n')
	except CalledProcessError:
		pass

	return cmd

args = tuple(sys.argv[1:])

if len(args) != 2:
	print('\nERROR: requires exactly two arguments')
	showUsage()
	sys.exit(1)

source = args[0]
target = args[1]

if os.path.isdir(source):
	print('\nERROR: source is a directory: {}'.format(source))
	sys.exit(errno.EISDIR)

if os.path.isdir(target):
	print('\nERROR: target is a directory: {}'.format(target))
	sys.exit(errno.EISDIR)

if not os.path.isfile(source):
	print('\nERROR: could not find source file: {}'.format(source))
	sys.exit(errno.ENOENT)

sys.stdout.write('\nChecking for "convert" executable ...')

# FIXME: Windows has a "convert" executable that is not related to ImageMagick
cmd_convert = getCommand('convert')
if not cmd_convert:
	print('\nERROR: could not find "convert" command\n       Please install ImageMagick: https://imagemagick.org/')
	sys.exit(errno.ENOENT)

print(' {}'.format(cmd_convert))

sys.stdout.write('\nChecking image dimensions ...')

# PIL
img = Image.open(source)

print(' {}'.format(img.size))

# cropping width & height
crop_wh = '{}x{}'.format(img.size[0], int(img.size[1] / 4))

sys.stdout.write('\nCreating temporary pieces ...')

f_basename = '.'.join(os.path.basename(source).split('.')[:-1])
f_suffix = source.split('.')[-1]

proc = subprocess.Popen((cmd_convert, source, '-define', 'png:format=png32', '-crop', crop_wh, os.path.join(dir_temp, '{}-TMP.{}'.format(f_basename, f_suffix)),))
out, err = proc.communicate()

temp_pieces = []
for I in (2, 3, 1, 0,):
	p = os.path.join(dir_temp, '{}-TMP-{}.{}'.format(f_basename, I, f_suffix))

	if not os.path.isfile(p):
		print('\n\nERROR: failed to created temporary images in "{}"'.format(dir_temp))
		sys.exit(errno.ENOENT)

	temp_pieces.append(p)

sys.stdout.write(' done!\n\nCreating new image from temporary pieces ...')

join_args = tuple([cmd_convert,] + temp_pieces + ['-define', 'png:format=png32', '-append', target,])

proc = subprocess.Popen(join_args)
out, err = proc.communicate()

print(' done!')


from django.shortcuts import render
from order.models import Order
from parse_data.models import Product
from user_admin.models import User
from django.core.paginator import Paginator, EmptyPage, PageNotAnInteger
import plotly.graph_objs as go
from plotly.offline import plot

def is_admin(request):
    user_id = request.COOKIES.get('userid', None)

    if user_id:
        try:
            user = User.objects.get(id=user_id)
            if user.is_admin:
                return True
        except User.DoesNotExist:
            pass

    return False

def list(request):
    if not is_admin(request):
        return render(request, 'not_admin.html')
    order_list = Order.objects.all().order_by('id')
    # Pagination
    paginator = Paginator(order_list, 10)  # Display 10 products per page
    page = request.GET.get('page') or 1  # If the requested page number is not set or is an empty string, set to 1

    try:
        orders = paginator.page(page)
    except PageNotAnInteger:
        # If the requested page number is not an integer, the first page is displayed
        orders = paginator.page(1)
    except EmptyPage:
        # If the requested page number is out of range, the last page will be displayed
        orders = paginator.page(paginator.num_pages)

    return render(request, 'admin_list.html',{'orders':orders})

def userList(request):
    if not is_admin(request):
        return render(request, 'not_admin.html')
    user_list = User.objects.all().order_by('id')
    # Pagination
    paginator = Paginator(user_list, 10)  # Display 10 products per page
    page = request.GET.get('page') or 1  # If the requested page number is not set or is an empty string, set to 1

    try:
        users = paginator.page(page)
    except PageNotAnInteger:
        # If the requested page number is not an integer, the first page is displayed
        users = paginator.page(1)
    except EmptyPage:
        # If the requested page number is out of range, the last page will be displayed
        users = paginator.page(paginator.num_pages)

    return render(request, 'user_list.html',{'users':users})

def dashboard(request):
    if not is_admin(request):
        return render(request, 'not_admin.html')
    orders = Order.objects.all().order_by('order_date')
    data = [
        go.Scatter(
            x=[order.order_date for order in orders],
            y=[order.product_quantity for order in orders],
            name='Quantity',
            mode='lines'
        ),
        go.Scatter(
            x=[order.order_date for order in orders],
            y=[order.product_price for order in orders],
            name='Price',
            mode='lines'
        )
    ]
    layout = go.Layout(
        title='Sales Over Time',
        xaxis=dict(title='Date'),
        yaxis=dict(title='Sales')
    )
    chart = plot(go.Figure(data=data, layout=layout), output_type='div')

    return render(request, 'dashboard.html', {'chart': chart})


from controller import Robot
import pika
import json
import string
import ast
import threading
import pathlib
import time
import os
import sys
from message_subscriber import MessageSubscriber

# This is a "workaround" to be able to import the class when the controllers are running in the
# Webots "sandbox". So we need to add the current working directory as well as "two parents steps up".
cbaa_path = os.path.join(os.getcwd(), os.pardir)
cbaa_path = os.path.join(cbaa_path, os.pardir)
sys.path.insert(0, cbaa_path)
from task_allocation.cbaa import CBAA

class ControllerSuperclass:
    def __init__(self, name, robot_type):
        '''
        Super class for the simpleactions generators, with the most common functionalities. Each robot type
        class extends this class
        '''
        # create the Robot instance.
        self.robot = Robot()
        self.robot_name = name
        self.robot_type = robot_type
        # get the time step of the current world.
        self.timestep = int(self.robot.getBasicTimeStep())

        self.available_simpleactions_functions = {}
        self.simpleactions = []

        # List of simpleaction from the configuration file
        self.config_simpleactions_names_cost = {}

        # Initiate the topics subscriber
        self.binding_key = f"{name}_queue"
        self.exchange = "routing_exchange"
        self.exchange_type = "direct"
        self.simpleactions_subscriber = MessageSubscriber(
            self.binding_key, self.exchange, self.exchange_type, self.execute_simpleactions_callback
        )

        # Create subscriber for listening to the CBAA task allocation updates
        # we can give it an empty binding key, because the 'fanout' exchange will ignore its value
        # either way
        cbaa_exchange_name = "cbaa_initiate_exchange"
        cbaa_exchange_type = "fanout"
        self.cbaa_subscriber = MessageSubscriber(
            "", cbaa_exchange_name, cbaa_exchange_type, self.initiate_cbaa_callback
        )

        # Initiate thread for listening to the CBAA task allocation initiation from the server
        cbaa_initiation_communication = threading.Thread(target=self.cbaa_subscriber.subscription)
        cbaa_initiation_communication.start()

        # Create a Subscription for receiving bids from the other robots
        self.cbaa_bids_exchange_name = "cbaa_bids_exchange"
        self.cbaa_bids_subscriber = MessageSubscriber(
            "", self.cbaa_bids_exchange_name, cbaa_exchange_type, self.receive_cbaa_bids_callback
        )
        # Initiate a thread for receiving the bids from the other robots
        cbaa_bids_communication = threading.Thread(target=self.cbaa_bids_subscriber.subscription)
        cbaa_bids_communication.start()

        # Read the config data to get the available simpleactions for this robot type
        self.read_config_data()

        # TODO this should somehow be dynamic
        # self.n_robots = 4
        self.consensus_based_auction_algorithm = CBAA(
            # self.robot_name, self.test_avail_simpleactions, self.n_robots
            "cbaa",self.robot_name, self.config_simpleactions_names_cost, self.n_robots
        )

        print(f"Super class initiated. {self.robot_name}")

    def read_config_data(self):
        """
        Reads the configuration file
        """
        with open(pathlib.Path.cwd().parent.parent / 'config.json') as json_config_file:
            data = json.load(json_config_file)
            self.n_robots = len(data["robots"])
            for x in data["robots"][self.robot_type]["simpleactions"]:
                self.config_simpleactions_names_cost[x["name"]] = float(x["cost"])

    def add_available_simpleaction(self, name, function_reference):
        self.available_simpleactions_functions[name] = function_reference

    def execute_simpleactions_callback(self, ch, method, properties, body):
        print(f"{self.robot_name} callback: %r" % body)
        # TODO as for now, the incoming messages are functions calls, separated by ","
        # simpleactions.extend(body.decode('utf-8').split(","))

        # Decode the JSON back to a list
        new_simpleactions = json.loads(body.decode('utf-8'))
        self.simpleactions.extend(new_simpleactions)
        print(f'{self.robot_name} Simpleactions = {self.simpleactions}, type={type(self.simpleactions)}')

        # Now execute the simpleactions
        # for i in range(len(simpleactions)):
        while self.simpleactions:
            sim_act = self.simpleactions.pop(0)
            print(f"sim_act = {sim_act}")
            print(f"Function name = {sim_act['function_name']}, Args = {sim_act['args']}")
            function_name = sim_act['function_name']
            args = sim_act['args']
            # Retrieve the args
            # function_name = sim_act.split('(')[0]
            # args = sim_act.split('(')[1].split(')')[0]
            #
            if not args:
                self.available_simpleactions_functions[function_name]()
                continue
            # Convert the arguments
            if all([x in string.digits for x in args]):
                print(f"{args} is a number")
                args = int(args)
            elif "[" in args or "]" in args:
                # ast.literal_eval will safely evaluate the string representation of a list (with the square brackets
                # as well) to a python list
                args = ast.literal_eval(args)

            self.available_simpleactions_functions[function_name](args)

        print(f"{self.robot_name} finished callback function")

    def receive_cbaa_bids_callback(self, ch, method, properties, body):
        all_bids = json.loads(body)
        print(f"{self.robot_name} received: {all_bids}")
        for other_robot_name, bids in all_bids.items():
            if self.robot_name == other_robot_name:
                # Skip the bid of this robot
                continue
            self.consensus_based_auction_algorithm.receive_other_winning_bids(other_robot_name, bids)

    def publish_bids(self, bids):
        connection = pika.BlockingConnection(
            pika.ConnectionParameters(host='localhost'))
        channel = connection.channel()

        channel.exchange_declare(exchange=self.cbaa_bids_exchange_name, exchange_type='fanout')

        # message = f"Bids by {self.robot_name}: {bids}"
        message = {self.robot_name : bids}
        channel.basic_publish(exchange=self.cbaa_bids_exchange_name, routing_key='', body=json.dumps(message))
        # print(" [x] Sent %r" % message)
        connection.close()

    def initiate_cbaa_callback(self, ch, method, properties, body):
        print(f'{self.robot_name} initiate cbaa callback function, received {body}')
        # time.sleep(1)
        # self.publish_bids()
        new_available_tasks = json.loads(body.decode('utf-8'))
        # print(f"new_available_tasks = {new_available_tasks}")
        # new_task_lists = []
        # for task in new_available_tasks:
        #     new_task_lists.append(task["simpleactions"])
        # print(f"new_task_list = {new_task_lists}")
        self.consensus_based_auction_algorithm.add_task_list(new_available_tasks)
        # self.consensus_based_auction_algorithm.add_task_list(new_task_lists)

        # Phase 1
        self.consensus_based_auction_algorithm.select_task()

        # Phase 2
        # self.consensus_based_auction_algorithm.update_task()
        bids = self.consensus_based_auction_algorithm.get_winning_bids()
        self.publish_bids(bids)

        self.consensus_based_auction_algorithm.confirm_all_bids(self.robot_name)

        self.consensus_based_auction_algorithm.update_task()

        # Finally, publish the results
        # winning_robots = self.consensus_based_auction_algorithm.winning_robots
        # self.send_cbaa_result_to_server(winning_robots)
        self.consensus_based_auction_algorithm.post_results()
        # After posting the consensus to the server, reset the values in the CBAA Class
        self.consensus_based_auction_algorithm.cleanup()

    # def send_cbaa_result_to_server(self, winning_robots):
    #     """
    #     Publish the winning robots to the server, using the 'routing_exchange' exchange which is commonly used
    #     to publish the simpleactions to the server
    #     """
    #     connection = pika.BlockingConnection(pika.ConnectionParameters(host='localhost'))
    #     channel = connection.channel()
    #
    #     exchange_name = 'routing_exchange'
    #     routing_key = 'server_queue'
    #     message = json.dumps(winning_robots)
    #
    #     channel.exchange_declare(exchange=exchange_name, exchange_type='direct')
    #     channel.basic_publish(exchange=exchange_name, routing_key=routing_key, body=message)
    #     print(f"{self.robot_name} published it winning bid list")
    #     connection.close()


from typing import Optional

import boto3

from app.config import AWS_REGION
from .data import AUTHORS
from moto import mock_dynamodb

@mock_dynamodb
def get_dynamodb_table(table_name: str):
    dynamodb = boto3.resource("dynamodb", region_name=AWS_REGION)
    return dynamodb.Table(table_name)

@mock_dynamodb
def setup_db() -> None:
    dynamodb = boto3.client("dynamodb", region_name=AWS_REGION)

    dynamodb.create_table(
        TableName='Authors',
        BillingMode='PAY_PER_REQUEST',
        KeySchema=[
            {
                'AttributeName': 'name',
                'KeyType': 'HASH'
            }
        ],
        AttributeDefinitions=[
            {
                'AttributeName': 'name',
                'AttributeType': 'S'
            },
            {
                'AttributeName': 'country',
                'AttributeType': 'S'
            }
        ],
        GlobalSecondaryIndexes=[
            {
                'IndexName': 'GSI1',
                'KeySchema': [
                    {
                        'AttributeName': 'name',
                        'KeyType': 'HASH'
                    },
                    {
                        'AttributeName': 'country',
                        'KeyType': 'RANGE'
                    }
                ],
                'Projection': {
                    'ProjectionType': 'ALL'
                }
            }
        ],
    )

@mock_dynamodb
def add_item(table_name: str, item) -> None:
    table = get_dynamodb_table(table_name)
    table.put_item(Item=item)

@mock_dynamodb
def get_item(table_name: str, keys) -> Optional[dict]:
    table = get_dynamodb_table(table_name)
    response = table.get_item(Key=keys)

    if "Item" not in response:
        return None

    return response["Item"]

@mock_dynamodb
def populate_authors():

    for author in AUTHORS:
        add_item("Authors", author)


from ModelHelper.Common.CommonUtils import get, get_valid
from ModelHelper.Common.CommonModels.ModelFactory import BackboneFactory
from torch.autograd import Variable
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

class AbstractClassifyModel(nn.Module):
    def __init__(self, **kwargs):
        super(AbstractClassifyModel, self).__init__()
        backbone = get_valid('backbone', kwargs)
        backbone_factory = BackboneFactory()
        kwargs['model_name'] = backbone
        self.backbone = backbone_factory.get_model(**kwargs)

    def forward(self, **kwargs):
        pass

# NTS Net
class ProposalNet(nn.Module):
    def __init__(self, **kwargs):
        super(ProposalNet, self).__init__()
        self.nts_fc_ratio = get_valid("nts_fc_ratio", kwargs)
        self.down1 = nn.Conv2d(512*self.nts_fc_ratio, 128, 3, 1, 1)
        self.down2 = nn.Conv2d(128, 128, 3, 2, 1)
        self.down3 = nn.Conv2d(128, 128, 3, 2, 1)
        self.ReLU = nn.ReLU()
        self.tidy1 = nn.Conv2d(128, 6, 1, 1, 0)
        self.tidy2 = nn.Conv2d(128, 6, 1, 1, 0)
        self.tidy3 = nn.Conv2d(128, 9, 1, 1, 0)

    def forward(self, x):
        batch_size = x.size(0)
        d1 = self.ReLU(self.down1(x))
        d2 = self.ReLU(self.down2(d1))
        d3 = self.ReLU(self.down3(d2))
        t1 = self.tidy1(d1).view(batch_size, -1)
        t2 = self.tidy2(d2).view(batch_size, -1)
        t3 = self.tidy3(d3).view(batch_size, -1)
        return torch.cat((t1, t2, t3), dim=1)

class NTSClassifyModel(AbstractClassifyModel):
    def __init__(self, **kwargs):
        kwargs['get_layer4_feature'] = True
        kwargs['get_fc_feature'] = True
        kwargs['get_fc_output'] = True
        self.nts_fc_ratio = get_valid("nts_fc_ratio", kwargs)
        self.use_gpu = get('use_gpu', kwargs, True)
        self.class_num = get_valid('class_num', kwargs)
        super(NTSClassifyModel, self).__init__(**kwargs)

        self.top_n = get('top_n', kwargs, 6)
        self.crop_size = get('crop_size', kwargs, (448, 448))
        self.cat_num = get('cat_num', kwargs, 4)
        self._default_anchors_setting = (
            dict(layer='p3', stride=32, size=48, scale=[2 ** (1. / 3.), 2 ** (2. / 3.)], aspect_ratio=[0.667, 1, 1.5]),
            dict(layer='p4', stride=64, size=96, scale=[2 ** (1. / 3.), 2 ** (2. / 3.)], aspect_ratio=[0.667, 1, 1.5]),
            dict(layer='p5', stride=128, size=192, scale=[1, 2 ** (1. / 3.), 2 ** (2. / 3.)],
                 aspect_ratio=[0.667, 1, 1.5]),
        )

        self.backbone.avgpool = nn.AdaptiveAvgPool2d(1)
        self.backbone.fc = nn.Linear(512*self.nts_fc_ratio, self.class_num)
        self.proposal_net = ProposalNet(**kwargs)
        self.concat_net = nn.Linear(512*self.nts_fc_ratio * (self.cat_num + 1), self.class_num)
        self.partcls_net = nn.Linear(512 * self.nts_fc_ratio, self.class_num)
        _, edge_anchors, _ = self.generate_default_anchor_maps()
        self.pad_side = 224
        self.edge_anchors = (edge_anchors + 224).astype(np.int)

    def generate_default_anchor_maps(self, anchors_setting=None):
        """
        generate default anchor

        :param anchors_setting: all informations of anchors
        :param input_shape: shape of input images, e.g. (h, w)
        :return: center_anchors: # anchors * 4 (oy, ox, h, w)
                 edge_anchors: # anchors * 4 (y0, x0, y1, x1)
                 anchor_area: # anchors * 1 (area)
        """
        input_shape = self.crop_size
        if anchors_setting is None:
            anchors_setting = self._default_anchors_setting

        center_anchors = np.zeros((0, 4), dtype=np.float32)
        edge_anchors = np.zeros((0, 4), dtype=np.float32)
        anchor_areas = np.zeros((0,), dtype=np.float32)
        input_shape = np.array(input_shape, dtype=int)

        for anchor_info in anchors_setting:

            stride = anchor_info['stride']
            size = anchor_info['size']
            scales = anchor_info['scale']
            aspect_ratios = anchor_info['aspect_ratio']

            output_map_shape = np.ceil(input_shape.astype(np.float32) / stride)
            output_map_shape = output_map_shape.astype(np.int)
            output_shape = tuple(output_map_shape) + (4,)
            ostart = stride / 2.
            oy = np.arange(ostart, ostart + stride * output_shape[0], stride)
            oy = oy.reshape(output_shape[0], 1)
            ox = np.arange(ostart, ostart + stride * output_shape[1], stride)
            ox = ox.reshape(1, output_shape[1])
            center_anchor_map_template = np.zeros(output_shape, dtype=np.float32)
            center_anchor_map_template[:, :, 0] = oy
            center_anchor_map_template[:, :, 1] = ox
            for scale in scales:
                for aspect_ratio in aspect_ratios:
                    center_anchor_map = center_anchor_map_template.copy()
                    center_anchor_map[:, :, 2] = size * scale / float(aspect_ratio) ** 0.5
                    center_anchor_map[:, :, 3] = size * scale * float(aspect_ratio) ** 0.5

                    edge_anchor_map = np.concatenate((center_anchor_map[..., :2] - center_anchor_map[..., 2:4] / 2.,
                                                      center_anchor_map[..., :2] + center_anchor_map[..., 2:4] / 2.),
                                                     axis=-1)
                    anchor_area_map = center_anchor_map[..., 2] * center_anchor_map[..., 3]
                    center_anchors = np.concatenate((center_anchors, center_anchor_map.reshape(-1, 4)))
                    edge_anchors = np.concatenate((edge_anchors, edge_anchor_map.reshape(-1, 4)))
                    anchor_areas = np.concatenate((anchor_areas, anchor_area_map.reshape(-1)))

        return center_anchors, edge_anchors, anchor_areas

    @staticmethod
    def hard_nms(cdds, topn=10, iou_thresh=0.25):
        if not (type(cdds).__module__ == 'numpy' and len(cdds.shape) == 2 and cdds.shape[1] >= 5):
            raise TypeError('edge_box_map should be N * 5+ ndarray')

        cdds = cdds.copy()
        indices = np.argsort(cdds[:, 0])
        cdds = cdds[indices]
        cdd_results = []

        res = cdds

        while res.any():
            cdd = res[-1]
            cdd_results.append(cdd)
            if len(cdd_results) == topn:
                return np.array(cdd_results)
            res = res[:-1]

            start_max = np.maximum(res[:, 1:3], cdd[1:3])
            end_min = np.minimum(res[:, 3:5], cdd[3:5])
            lengths = end_min - start_max
            intersec_map = lengths[:, 0] * lengths[:, 1]
            intersec_map[np.logical_or(lengths[:, 0] < 0, lengths[:, 1] < 0)] = 0
            iou_map_cur = intersec_map / ((res[:, 3] - res[:, 1]) * (res[:, 4] - res[:, 2]) + (cdd[3] - cdd[1]) * (
                    cdd[4] - cdd[2]) - intersec_map)
            res = res[iou_map_cur < iou_thresh]

        return np.array(cdd_results)

    def forward(self, **kwargs):
        super(NTSClassifyModel, self).forward(**kwargs)
        image = get_valid('image', kwargs)

        feature_list = self.backbone(image)
        resnet_out = feature_list[2]
        rpn_feature = feature_list[0]
        feature = feature_list[1]

        x_pad = F.pad(image, (self.pad_side, self.pad_side, self.pad_side, self.pad_side), mode='constant', value=0)
        batch = image.size(0)
        # we will reshape rpn to shape: batch * nb_anchor
        rpn_score = self.proposal_net(rpn_feature.detach())
        all_cdds = [
            np.concatenate((x.reshape(-1, 1), self.edge_anchors.copy(), np.arange(0, len(x)).reshape(-1, 1)), axis=1)
            for x in rpn_score.data.cpu().numpy()]
        top_n_cdds = [self.hard_nms(x, topn=self.top_n, iou_thresh=0.25) for x in all_cdds]
        top_n_cdds = np.array(top_n_cdds)
        top_n_index = top_n_cdds[:, :, -1].astype(np.int)
        if self.use_gpu:
            top_n_index = torch.from_numpy(top_n_index).cuda()
        else:
            top_n_index = torch.from_numpy(top_n_index)
        top_n_prob = torch.gather(rpn_score, dim=1, index=top_n_index.long())
        if self.use_gpu:
            part_imgs = torch.zeros([batch, self.top_n, 3, 224, 224]).cuda()
        else:
            part_imgs = torch.zeros([batch, self.top_n, 3, 224, 224])
        for i in range(batch):
            for j in range(self.top_n):
                [y0, x0, y1, x1] = top_n_cdds[i][j, 1:5].astype(np.int)
                part_imgs[i:i + 1, j] = F.interpolate(x_pad[i:i + 1, :, y0:y1, x0:x1], size=(224, 224), mode='bilinear',
                                                      align_corners=True)
        part_imgs = part_imgs.view(batch * self.top_n, 3, 224, 224)
        feature_list2 = self.backbone(part_imgs.detach())
        part_features = feature_list2[1]
        part_feature = part_features.view(batch, self.top_n, -1)
        part_feature = part_feature[:, :self.cat_num, ...].contiguous()
        part_feature = part_feature.view(batch, -1)
        concat_out = torch.cat([part_feature, feature], dim=1)
        concat_logits = self.concat_net(concat_out)
        raw_logits = resnet_out
        part_logits = self.partcls_net(part_features).view(batch, self.top_n, -1)
        return [raw_logits, concat_logits, part_logits, top_n_index, top_n_prob]

    @staticmethod
    def list_loss(logits, targets):
        temp = F.log_softmax(logits, -1)
        loss = [-temp[i][targets[i].item()] for i in range(logits.size(0))]
        return torch.stack(loss)

    @staticmethod
    def ranking_loss(score, targets, proposal_num, **kwargs):
        use_gpu = get('use_gpu', kwargs, True)
        if use_gpu:
            loss = Variable(torch.zeros(1).cuda())
        else:
            loss = Variable(torch.zeros(1))
        batch_size = score.size(0)
        for i in range(proposal_num):
            if use_gpu:
                targets_p = (targets > targets[:, i].unsqueeze(1)).type(torch.cuda.FloatTensor)
            else:
                targets_p = (targets > targets[:, i].unsqueeze(1)).type(torch.FloatTensor)
            pivot = score[:, i].unsqueeze(1)
            loss_p = (1 - pivot + score) * targets_p
            loss_p = torch.sum(F.relu(loss_p))
            loss += loss_p
        return loss / batch_size

class ResnetClassifyModel(AbstractClassifyModel):
    def __init__(self, **kwargs):
        kwargs['get_fc_feature'] = True
        super(ResnetClassifyModel, self).__init__(**kwargs)
        self.class_num = get_valid('class_num', kwargs)
        fc_input_num = get_valid('fc_input_num', kwargs)
        self.fc = nn.Linear(fc_input_num, self.class_num)

    def forward(self, **kwargs):
        super(ResnetClassifyModel, self).forward(**kwargs)
        image = get_valid('image', kwargs)
        feature_list = self.backbone(image)
        output = self.fc(feature_list[0])
        return output


import pytest

from application.crud import room as crud_room, booking as crud_booking
from application.database.models import Booking
from application.dto import BookingIn, RoomIn

def test_booking_create_success(client, test_db, room_input, booking_input):
    room = crud_room.save_room(test_db, RoomIn(**room_input))

    input_data = {**booking_input, 'room_id': room.id}
    response = client.post('/booking/', json=input_data)

    assert response.status_code == 201

    data = response.json()
    expected_data = BookingIn(**input_data)
    assert 'id' in data

    booking = test_db.query(Booking).get(data['id'])
    assert booking.begin_date == expected_data.begin_date
    assert booking.end_date == expected_data.end_date
    assert booking.room_id == expected_data.room_id

@pytest.mark.parametrize('field', ['begin_date', 'end_date', 'room_id'])
def test_booking_create_without_required_field(client, booking_input, field):
    input_data = {**booking_input}
    del input_data[field]
    response = client.post('/booking/', json=input_data)

    assert response.status_code == 422

@pytest.mark.parametrize('begin_date', [None, '', '2021-01-0', '2021-02-30'])
def test_booking_create_invalid_begin_date(client, booking_input, begin_date):
    response = client.post('/booking/', json={**booking_input, 'begin_date': begin_date})

    assert response.status_code == 422

@pytest.mark.parametrize('end_date', [None, '', '2021-01-0', '2021-02-30'])
def test_booking_create_invalid_end_date(client, booking_input, end_date):
    response = client.post('/booking/', json={**booking_input, 'end_date': end_date})

    assert response.status_code == 422

@pytest.mark.parametrize('room_id', [None, '', 'N'])
def test_booking_create_invalid_room_id(client, booking_input, room_id):
    response = client.post('/booking/', json={**booking_input, 'room_id': room_id})

    assert response.status_code == 422

def test_booking_create_end_date_preceds_begin_date(client):
    response = client.post('/booking/', json={
        'begin_date': '2021-01-07',
        'end_date': '2021-01-01',
        'room_id': 1
    })

    assert response.status_code == 422

@pytest.mark.parametrize('dates', [
    # inside of the date range
    {'begin_date': '2021-01-07', 'end_date': '2021-01-13'},
    {'begin_date': '2021-01-06', 'end_date': '2021-01-14'},
    {'begin_date': '2021-01-07', 'end_date': '2021-01-14'},
    # outside of the date range
    {'begin_date': '2021-01-06', 'end_date': '2021-01-15'},
    {'begin_date': '2021-01-06', 'end_date': '2021-01-14'},
    {'begin_date': '2021-01-07', 'end_date': '2021-01-15'},
    # crosses the left border
    {'begin_date': '2021-01-06', 'end_date': '2021-01-13'},
    {'begin_date': '2021-01-06', 'end_date': '2021-01-07'},
    # crosses the right border
    {'begin_date': '2021-01-13', 'end_date': '2021-01-15'},
    {'begin_date': '2021-01-14', 'end_date': '2021-01-15'},
])
def test_booking_create_overlapping_dates(client, test_db, room_input, dates):
    room = crud_room.save_room(test_db, RoomIn(**room_input))

    booking_1 = {
        'begin_date': '2021-01-07',
        'end_date': '2021-01-14',
        'room_id': room.id
    }
    crud_booking.save_booking(test_db, BookingIn(**booking_1))

    booking_2 = {
        **dates,
        'room_id': room.id
    }
    response = client.post('/booking/', json=booking_2)

    assert response.status_code == 422


from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

from database_setup import Genre, Base, Movie, User

engine = create_engine('sqlite:///Movie_Genre.db')

Base.metadata.bind = engine

DBSession = sessionmaker(bind=engine)
session = DBSession()

# Create dummy user
User1 = User(name="John Doe", email="johndoe@gmail.com",
             picture='https://nicolealvares.files.'
                     'wordpress.com/2013/10/making-of-the'
                     '-worlds-first-camera-thumb.jpg')

session.add(User1)
session.commit()

# Movie comedy
movies = Genre(user_id=1, name="Comedy")

session.add(movies)
session.commit()

Item2 = Movie(user_id=1, name="Smallfoot",
              description="Mythical Creatures, Unlikely Friendships",
              price="$9.99", ratings="PG", genre=movies)

session.add(Item2)
session.commit()

Item1 = Movie(user_id=1, name="Teen Titans Go! To the Movies",
              description="Big Break, Heroic Mission, Unlikely Heroes",
              price="$9.99", ratings="PG", genre=movies)

session.add(Item1)
session.commit()

Item2 = Movie(user_id=1, name="The House with a Clock in Its Walls",
              description="Curses and Spells, Wizards and Magicians",
              price="$5.50", ratings="PG", genre=movies)

session.add(Item2)
session.commit()

Item3 = Movie(user_id=1, name="Ferdinand",
              description="Coming Home, Non-Traditional Families",
              price="$3.99", ratings="PG", genre=movies)

session.add(Item3)
session.commit()

Item4 = Movie(user_id=1, name="Spaceballs",
              description="Space Wars, Robots and Androids, Unlikely Heroes",
              price="$3.99", ratings="PG", genre=movies)

session.add(Item4)
session.commit()

Item5 = Movie(user_id=1, name="Space Jam",
              description="Basketball Players, Evil Aliens, Heroic Mission",
              price="$12.99", ratings="PG", genre=movies)

session.add(Item5)
session.commit()

Item6 = Movie(user_id=1, name="Shrek the Third",
              description="Fantasy Lands, Talking Animals, Crowned Heads, "
                          "Curses and Spells, Fish Out of Water, "
                          "Mythical Creatures",
              price="$.99", ratings="NR", genre=movies)

session.add(Item6)
session.commit()

Item7 = Movie(user_id=1, name="Dr. Seuss' The Grinch",
              description="Conspiracies, Crime Sprees, Metamorphosis, "
                          "Unlikely Friendships, Unlikely Heroes",
              price="$13.49", ratings="PG", genre=movies)

session.add(Item7)
session.commit()

# Action movies
movies2 = Genre(user_id=1, name="Action")

session.add(movies2)
session.commit()

Item1 = Movie(user_id=1, name="Incredibles 2",
              description="Heroic Mission, Unlikely Heroes",
              price="$19.99", ratings="PG", genre=movies2)

session.add(Item1)
session.commit()

Item2 = Movie(user_id=1, name="Christopher Robin",
              description="Starting Over, Toys Come to Life, "
                          "Unlikely Friendships",
              price="$25", ratings="PG", genre=movies2)

session.add(Item2)
session.commit()

Item3 = Movie(user_id=1, name="Moana",
              description="Chosen One, Curses and Spells, "
                          "Heroic Mission, Lost Worlds, "
                          "Priceless Artifacts and Prized Objects",
              price="$15", ratings="PG", genre=movies2)

session.add(Item3)
session.commit()

Item4 = Movie(user_id=1, name="Sherlock Gnomes",
              description="Amateur Sleuths, Kidnapping, Star Detectives, "
                          "Toys Come to Life",
              price="$12", ratings="PG", genre=movies2)

session.add(Item4)
session.commit()

Item5 = Movie(user_id=1, name="Johnny English Strikes Again",
              description="Betrayal, Computer Paranoia, "
                          "Hijackings, Virtual Reality.",
              price="$14", ratings="PG", genre=movies2)

session.add(Item5)
session.commit()

Item6 = Movie(user_id=1, name="The Jungle Book ",
              description="	Monkeys, Survival in the Wilderness, "
                          "Talking Animals.",
              price="$12", ratings="PG", genre=movies2)

session.add(Item6)
session.commit()

# Thriller movies
movies = Genre(user_id=1, name="Thriller")

session.add(movies)
session.commit()

Item1 = Movie(user_id=1, name="Ocean's 8", description="Jewel Theft.",
                     price="$8.99", ratings="PG13", genre=movies)

session.add(Item1)
session.commit()

Item2 = Movie(user_id=1, name="The Darkest Minds",
              description="Escape From Prison, Future Dystopias, "
                          "On the Run, Redemption.",
              price="$6.99", ratings="PG13", genre=movies)

session.add(Item2)
session.commit()

Item3 = Movie(user_id=1, name="The Post",
              description="	Conspiracies, Fighting the System, "
                          "Office Politics, Scandals and Cover-Up",
              price="$9.95", ratings="PG13", genre=movies)

session.add(Item3)
session.commit()

Item4 = Movie(user_id=1, name="Non-Stop",
              description="	Air Disasters, Hijackings, "
                          "Race Against Time, Trapped or Confined.",
              price="$6.99", ratings="PG13", genre=movies)

session.add(Item4)
session.commit()

# Horror Movies
movies = Genre(user_id=1, name="Horror")

session.add(movies)
session.commit()

Item1 = Movie(user_id=1, name="Halloween",
              description="Crime Sprees, Escape From Prison, "
                          "Haunted By the Pas.",
              price="$2.99", ratings="R", genre=movies)

session.add(Item1)
session.commit()

Item2 = Movie(user_id=1, name="The Nun",
              description="Demonic Possession, "
                          "Members of the Clergy, Suicide",
              price="$5.99", ratings="R", genre=movies)

session.add(Item2)
session.commit()

Item3 = Movie(user_id=1, name="Hotel Artemis",
              description="Future Dystopias, Secret Organizations",
              price="$4.50", ratings="R", genre=movies)

session.add(Item3)
session.commit()

Item4 = Movie(user_id=1, name="Bad Times at the El Royale",
              description="Crime Sprees, Keeping a Secret",
              price="$6.95", ratings="R", genre=movies)

session.add(Item4)
session.commit()

Item5 = Movie(user_id=1, name="Red Sparrow",
              description="Femmes Fatales, Switching Sides, "
                          "Traitorous Spies",
              price="$7.95", ratings="R", genre=movies)

session.add(Item5)
session.commit()

# War movie
movies = Genre(user_id=1, name="War")

session.add(movies)
session.commit()

Item1 = Movie(user_id=1, name="Fury",
              description="	Heroic Mission, Military Life",
              price="$13.95", ratings="R", genre=movies)

session.add(Item1)
session.commit()

Item2 = Movie(user_id=1, name="	Zero Dark Thirty",
              description="Assassination Plots, Heroic Mission, Terrorism",
              price="$4.95", ratings="R", genre=movies)

session.add(Item2)
session.commit()

Item3 = Movie(user_id=1, name="Hacksaw Ridge",
              description="Great Battles, Message From God, Military Life",
              price="$6.95", ratings="R", genre=movies)

session.add(Item3)
session.commit()

Item4 = Movie(user_id=1, name="	Dunkirk",
              description="Great Battles, Heroic Mission, "
                          "War At Sea, War in the Sky",
              price="$13.95", ratings="PG13", genre=movies)

session.add(Item4)
session.commit()

Item5 = Movie(user_id=1, name="Behind Enemy Lines",
              description="Behind Enemy Lines, Daring Rescues, "
                          "Heroic Mission",
              price="$7.95", ratings="PG13", genre=movies)

session.add(Item5)
session.commit()

# SciFi movie
movies = Genre(user_id=1, name="SciFi")

session.add(movies)
session.commit()

Item1 = Movie(user_id=1, name="The Predator",
              description="Evil Aliens",
              price="$9.95", ratings="R", genre=movies)

session.add(Item1)
session.commit()

Item2 = Movie(user_id=1, name="Ready Player One",
              description="Chicken cooked in Marsala "
                          "wine sauce with mushrooms",
              price="$7.95", ratings="PG", genre=movies)

session.add(Item2)
session.commit()

Item3 = Movie(user_id=1, name="Potstickers",
              description="Bounty Hunters, Computer Paranoia, "
                          "Future Dystopias, Priceless.",
              price="$6.50", ratings="PG13", genre=movies)

session.add(Item3)
session.commit()

Item4 = Movie(user_id=1, name="Sorry to Bother You",
              description="Assumed Identities, Big Break, "
                          "Double Life, Schemes and Ruses",
              price="$16.75", ratings="PG13", genre=movies)

session.add(Item4)
session.commit()

# Documentary
movies = Genre(user_id=1, name="Documentary")

session.add(movies)
session.commit()

Item9 = Movie(user_id=1, name="WHAT'S THE MATTER WITH KANSAS",
              description="Down on Their Luck, Religious "
                          "Zealotry, Underdogs",
              price="$8.99", ratings="NR", genre=movies)

session.add(Item9)
session.commit()

Item1 = Movie(user_id=1, name="Woodstock",
              description="Generation Gap, Bohemian Life",
              price="$2.99", ratings="R", genre=movies)

session.add(Item1)
session.commit()

Item2 = Movie(user_id=1, name="Undefeated",
              description="Documentary",
              price="$10.95", ratings="PG13", genre=movies)

session.add(Item2)
session.commit()

Item3 = Movie(user_id=1, name="Get Thrashed: The Story of Thrash Metal",
              description="Explore the blistering rise, brutal fall, and "
                          "lasting impact of thrash metal",
              price="$7.50", ratings="PG13", genre=movies)

session.add(Item3)
session.commit()

Item4 = Movie(user_id=1, name="Where Was God?",
              description="This inspirational documentary follows the "
                          "aftermath of the May 20, 2013 tornado "
                          "that devastated Moore, Oklahoma",
              price="$8.95", ratings="PG", genre=movies)

session.add(Item4)
session.commit()

session.add(Item2)
session.commit()

Item10 = Movie(user_id=1, name="MANNY",
               description="In Training, Boxers",
               price="$1.99", ratings="PG13", genre=movies)

session.add(Item10)
session.commit()

# Anime Movie
movies = Genre(user_id=1, name="Anime")

session.add(movies)
session.commit()

Item1 = Movie(user_id=1, name="Spirited Away ",
              description="Journey of Self-Discovery, Mythical Creatures, "
                          "Fantasy Lands",
              price="$5.95", ratings="PG", genre=movies)

session.add(Item1)
session.commit()

Item2 = Movie(user_id=1, name="Princess Mononoke",
              description="Heroic Mission, Mythical "
                          "Creatures, Righting the Wronged.",
              price="$17.99", ratings="PG", genre=movies)

session.add(Item2)
session.commit()

movies = Genre(user_id=1, name="Fantasy")
session.add(movies)
session.commit()

Item1 = Movie(user_id=1, name="Dog Days",
              description="Expecting a Baby, Intersecting Lives, "
                          "Looking For Love, Man's Best Friend",
              price="$5.95", ratings="PG", genre=movies)

session.add(Item1)
session.commit

Item1 = Movie(user_id=1, name="The Shape of Water",
              description="Living With Disability, "
                          "Mutants, On the Run, Workplace Romance",
              price="$16.95", ratings="R", genre=movies)

session.add(Item1)
session.commit()

Item1 = Movie(user_id=1, name="	Isle of Dogs",
              description="Conspiracies, Daring Rescues, "
                          "Future Dystopias, Talking Animals",
              price="$14.25", ratings="R", genre=movies)

session.add(Item1)
session.commit()

print("added menu items!")


from __future__ import print_function, division

from math import exp, expm1
import argparse

from jsonctmctree.interface import process_json_in

# high population boundary behaviors
WRAP = 'wrap'
ABSORB = 'absorb'
BLOCK = 'block'

def gen_speciation_triples(lam, n, boundary):
    for i in range(1, n-1):
        yield [1, i], [1, i+1], i*lam
    if boundary == WRAP:
        yield [1, n-1], [0, 0], (n-1)*lam

def gen_extinction_triples(mu, n, boundary):
    yield [1, 1], [0, 0], mu
    for i in range(2, n-1):
        yield [1, i], [1, i-1], i*mu
    if boundary in {WRAP, BLOCK}:
        yield [1, n-1], [1, n-2], (n-1)*mu

def get_scene(edge_rates, mu, lam, n, boundary):
    leaves = [2, 3, 4]
    tree = dict(
            row_nodes = [0, 0, 1, 1],
            column_nodes = [1, 2, 3, 4],
            edge_rate_scaling_factors = edge_rates,
            edge_processes = [0, 0, 0, 0])
    speciation_triples = list(gen_speciation_triples(lam, n, boundary))
    extinction_triples = list(gen_extinction_triples(mu, n, boundary))
    triples = speciation_triples + extinction_triples
    row_states, column_states, transition_rates = zip(*triples)
    scene = dict(
            node_count = 5,
            process_count = 1,
            state_space_shape = [2, n],
            tree = tree,
            root_prior = dict(
                states = [[1, 1]],
                probabilities = [1.0],
                ),
            process_definitions = [dict(
                row_states = row_states,
                column_states = column_states,
                transition_rates = transition_rates,
                )],
            observed_data = dict(
                nodes = leaves,
                variables = [0, 0, 0],
                iid_observations = [[1, 1, 1]],
                )
            )
    return scene

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--n', type=int, default=1000, help='population size')
    parser.add_argument('--mu', type=float, default=0.3, help='extinction')
    parser.add_argument('--lam', type=float, default=0.5, help='speciation')
    args = parser.parse_args()

    n = args.n
    mu = args.mu
    lam = args.lam

    edge_rates = [5, 10, 5, 5]
    wrap_scene = get_scene(edge_rates, mu, lam, n, WRAP)
    absorb_scene = get_scene(edge_rates, mu, lam, n, ABSORB)

    # Log likelihood.
    logl_request = dict(property='snnlogl')

    # Unweighted sum over observations and over edges,
    # and weighted sum over transitions consisting of the unweighted sum
    # over transitions corresponding to extinction events.
    extinction_request = dict(
        property='ssntran',
        transition_reduction = dict(
            row_states = [[1, i] for i in range(2, n)],
            column_states = [[1, i-1] for i in range(2, n)],
            weights = [1]*(n-2),
        ))

    # Unweighted sum over observations, and weighted sum over states.
    extant_request = dict(
        property='snwnode',
        state_reduction = dict(
            states = [[1, i] for i in range(n)],
            weights = range(n),
            ))

    # Unweighted sum over observations, weighted sum over edges,
    # and weighted sum over states.
    dwell_request = dict(
        property='swwdwel',
        edge_reduction = dict(
            edges = [0, 1, 2, 3],
            weights = edge_rates),
        state_reduction = dict(
            states = [[1, i] for i in range(n)],
            weights = range(n)))

    # Compute only the likelihood for the absorbing high population boundary.
    j_out = process_json_in(dict(scene=absorb_scene, requests=[logl_request]))
    absorb_likelihood = exp(j_out['responses'][0])

    # Compute more stuff for the wrapping boundary.
    j_in = dict(
            scene=wrap_scene,
            requests=[
                logl_request,
                extinction_request,
                extant_request,
                dwell_request,
                ])
    j_out = process_json_in(j_in)

    logl, extinction, extant, dwell = j_out['responses']
    wrap_likelihood = exp(logl)
    print('gene population limit:', n)
    print('gene birth rate:', lam)
    print('gene death rate:', mu)
    print('likelihood:', wrap_likelihood)
    print('upper bound likelihood for unbounded population:', absorb_likelihood)
    print('unconditional probability of exceeding the population cap:',
            absorb_likelihood - wrap_likelihood)
    print('expected number of extinctions:', extinction)
    print('expected number of extant lineages at each node:')
    for i, x in enumerate(extant):
        print(i, ':', x)
    print('expected total size of the gene tree:', dwell)

main()


from bitstring import Bits, BitString, BitArray, ConstBitStream

import base64

class Addr:
    def __init__(self,bank,addr) -> None:
        self.bank = bank
        self.addr = addr

    def absolute_pos(self) -> int:
        return (((self.bank-1)*BANK_SIZE)+self.addr)

    @classmethod
    def convert_to_addr(cls, long_addr) -> None:
        bank = int(long_addr/BANK_SIZE)
        addr = (long_addr%BANK_SIZE)+(BANK_SIZE if bank > 0 else 0)
        return cls(bank,addr)

    def __str__(self) -> str:
        return f"{self.bank:#04X}:{self.addr:04X}"

    def __add__(self, other):
        if isinstance(other, int):
            diff = other
        elif isinstance(other, Addr):
            diff = abs(self.absolute_pos() - other.absolute_pos())
        return self.convert_to_addr(self.absolute_pos() + diff)

    def __sub__(self, other):
        if isinstance(other, int):
            diff = other
        elif isinstance(other, Addr):
            diff = abs(self.absolute_pos() - other.absolute_pos())
        return self.convert_to_addr(self.absolute_pos() - diff)

    def __eq__(self, other) -> bool:
        return self.absolute_pos() == other.absolute_pos()

    def __gt__(self, other) -> bool:
        return self.absolute_pos() > other.absolute_pos()

    def __lt__(self, other) -> bool:
        return self.absolute_pos() < other.absolute_pos()

    def __ge__(self, other) -> bool:
        return self.absolute_pos() >= other.absolute_pos()

    def __le__(self, other) -> bool:
        return self.absolute_pos() <= other.absolute_pos()

    def __ne__(self, other) -> bool:
        return self.absolute_pos() != other.absolute_pos()

class GBDataPacket:
    def __init__(self, addr, packet_size, data) -> None:
        self.addr = addr
        self.packet_size = packet_size
        self.data = data

    @classmethod
    def get_static_data(cls, addr, packet_size, length):
        ROM.bytepos = addr.absolute_pos()
        data = ROM.readlist([f'uint:{packet_size}']*length)
        return cls(addr,packet_size,data)

    @classmethod
    def get_var_data(cls, addr, packet_size, target, bytealigned=True):
        ROM.bytepos = addr.absolute_pos()
        data = ROM.readto(target,bytealigned)
        data_list = data.readlist([f'uint:{packet_size}']*int(data.len/packet_size))
        return cls(addr,packet_size,data_list)

    def collapse(self, rev=False) -> int:
        out = 0
        if rev:
            self.data.reverse()
        for val in self.data:
            out = out << self.packet_size
            out+=val
        if rev:
            self.data.reverse()
        return out

    def __str__(self) -> str:
        return f"{self.addr}  " + " ".join(map((lambda n: f"{n:02x}"), self.data))

    def raw_dump(self) -> str:
        out = ""
        out+=(f"Start:{self.addr} Finish:{self.addr+len(self.data)} Length:{(len(self.data))} 2BPP:{len(self.data)/16:0.0f} 1BPP:{len(self.data)/8:0.0f}\n")

        data_fmt = []
        for i in range(int(len(self.data)/16)):
            data_fmt.append(f"{(i*16):#07X} " + ' '.join(map(lambda n: f"{n:02X}", self.data[16*i:(16*i)+16])))

        out+=('\n'.join(data_fmt))
        if (len(self.data) % 16 != 0):
            out+=(f"\n{len(data_fmt)*16:#07X} " + ' '.join(map(lambda n: f"{n:02X}", self.data[len(data_fmt)*16:])))
        return out

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index) -> int:
        return self.data[index]

class Sprite:

    def __init__(self,addr,width,height,data) -> None:
        self.addr = addr
        self.width = width
        self.height = height
        self.data = data

    def __str__(self):
        return f"[Loc: {self.addr} => Width: {self.width}, Height: {self.height}]"

    def to_json(self) -> dict:
        return {'width': self.width, 'height': self.height, 'data': self.to_base64()}

    @classmethod
    def __expandRLEPacket(cls, bit_length, value) -> BitString:
        return BitString((bit_length+value+1)*2)

    @classmethod
    def __findRLEBoundry(cls, sprite_data) -> Bits:
        length_found = sprite_data.readto('0b0')
        return length_found

    @classmethod
    def __mode1(cls,bit_planes,width) -> list:
        bit_planes[1] = cls.__deltaDecode(bit_planes[1],width)
        bit_planes[0] = cls.__deltaDecode(bit_planes[0],width)
        return bit_planes

    @classmethod
    def __mode2(cls,bit_planes,width) -> list:
        bit_planes[1] = cls.__deltaDecode(bit_planes[1],width)
        bit_planes[0] = bit_planes[0] ^ bit_planes[1]
        return bit_planes

    @classmethod
    def __mode3(cls,bit_planes,width) -> list:
        bit_planes[1] = cls.__deltaDecode(bit_planes[1],width)
        bit_planes[0] = cls.__deltaDecode(bit_planes[0],width)
        bit_planes[0] = bit_planes[0] ^ bit_planes[1]
        return bit_planes

    @classmethod
    def ___translate(cls, arr,row_num,coloumn_num):
        matrix = [[0 for x in range(coloumn_num)] for y in range(row_num)]
        for row in range(row_num):
            for col in range(int(coloumn_num/8)):
                for i in range(8):
                    matrix[row][col+i]=arr[(row*col)+row+i]
        return matrix

    @classmethod
    def __fillMatrix(cls, arr,row_num, coloumn_num) -> BitArray:
        #Array math is hard touch numbers at own risk
        matrix = [[0 for x in range(coloumn_num*4)] for y in range(row_num*8)]
        for row in range(row_num*8):
            for col in range(coloumn_num*4):
                matrix[row][col]=(''.join(arr[((col*row_num*16)+(row*2)):((col*row_num*16)+(row*2))+2].bin))
            matrix[row] = ''.join(matrix[row])

        output = BitArray()
        for out_row in matrix:
            output.append('0b'+out_row)

        return output

    @classmethod
    def __bufferToList(cls, arr, row_num, coloumn_num) -> list:
        #1 byte per row per tile
        #1 byte per coloumn per tile
        bufList = [0] * row_num*BYTE
        column_bits = coloumn_num*BYTE
        for row in range(row_num*BYTE):
            bufList[row]=list(map(int,(','.join(arr[(row*column_bits):((row*column_bits)+column_bits)].bin).split(','))))
        return bufList

    @classmethod
    def __combineBuffers(cls,bit_planes,high_bit_plane) -> list:
        result = [[(bit_planes[high_bit_plane][i][j]<<1) + bit_planes[high_bit_plane^1][i][j]  for j in range(len(bit_planes[high_bit_plane][0]))] for i in range(len(bit_planes[1]))]
        return result

    @classmethod
    def __fillTileMatrix(cls, arr, sprite_height_tiles, sprite_width_tiles) -> list:
        tile_side_px = 8
        tile_size = tile_side_px*tile_side_px
        out = []
        for tile_row in range (sprite_height_tiles):
            for row in range(tile_side_px):
                temp = []
                for col in range (sprite_width_tiles):
                    temp.extend(arr[((tile_row*tile_size*sprite_width_tiles)+(col*tile_size)+(row*tile_side_px)):((tile_row*tile_size*sprite_width_tiles)+(col*tile_size)+(row*tile_side_px))+tile_side_px])
                out.append(temp)
        return out

    def print_pixels(self):
        for row in self.data:
            print(','.join(map(str,row)))

    def __to_bignum(self) -> int:
        output = 0
        for row in self.data:
            for col in row:
                output = output << 2
                output += col
        return output

    def to_base64(self) -> str:
        num = self.__to_bignum()
        num_bytes = num.to_bytes((int(self.height*self.width*TWO_BPP_TILE_SIZE)),'big')
        return base64.b64encode(num_bytes).decode()

    @classmethod
    def __deltaDecode(cls, arr, width) -> BitArray:
        output = BitArray()
        currentBit = 0
        for index, bit in enumerate(arr):
            if index % (width*8) == 0:
                currentBit = 0
            if bit:
                currentBit = (currentBit ^ 1)

            output.append('0b%s' % currentBit)
        return output

    @classmethod
    def __parseData(cls, packet_type, width, height, bit_plane):
        while bit_plane.len < (width*height*ONE_BPP_TILE_SIZE*BYTE):
            if packet_type == 0:
                length = cls.__findRLEBoundry(ROM)
                value = ROM.read((f"uint:{length.len}"))
                zero_bits = cls.__expandRLEPacket(length.uint,value)
                bit_plane.append(zero_bits)
                packet_type = 1
            else:
                data_packet = ROM.read('bin:2')
                if data_packet != '00':
                    bit_plane.append('0b'+data_packet)
                else:
                    packet_type = 0

    @classmethod
    def parse_pkmn_sprite(cls, addr) -> None:
        ROM.bytepos = addr.absolute_pos()
        width = ROM.read('uint:4')
        height = ROM.read('uint:4')
        high_bit_plane = ROM.read('uint:1')
        packet_type = ROM.read('uint:1')
        bit_planes = [BitArray(), BitArray()]
        cls.__parseData(packet_type,width,height,bit_planes[1])
        zip_mode = -1
        if ROM.peek('uint:1') == 0:
            zip_mode = ROM.read('uint:1')
        else:
            zip_mode = ROM.read('uint:2')
        packet_type = ROM.read('uint:1')

        cls.__parseData(packet_type,width,height,bit_planes[0])

        bit_planes[0] = cls.__fillMatrix(bit_planes[0],width,height)
        bit_planes[1] = cls.__fillMatrix(bit_planes[1],width,height)
        if zip_mode == 0:
            bit_planes = cls.__mode1(bit_planes,width)
        elif zip_mode == 2:
            bit_planes = cls.__mode2(bit_planes,width)
        else:
            bit_planes = cls.__mode3(bit_planes,width)

        bit_planes[0] = cls.__bufferToList(bit_planes[0],width,height)
        bit_planes[1] = cls.__bufferToList(bit_planes[1],width,height)

        sprite_data = cls.__combineBuffers(bit_planes,high_bit_plane)

        return cls(addr,width,height,sprite_data)

    @classmethod
    def decode1BPP(cls,addr,width,height):
        ROM.bytepos = addr.absolute_pos()
        bit_planes = [BitArray(), BitArray()]
        for i in range(width*height*BYTE):
            bit_planes[0].append(ROM.peek('bits:8'))
            bit_planes[1].append(ROM.read('bits:8'))

        for i in range(2):
            bit_planes[i] = cls.__fillTileMatrix(bit_planes[i],height,width)

        sprite_data = cls.__combineBuffers(bit_planes,1)

        return cls(addr,width,height,sprite_data)

    @classmethod
    def decode2BPP(cls,addr,width,height):
        ROM.bytepos = addr.absolute_pos()
        bit_planes = [BitArray(), BitArray()]
        for i in range(width*height*BYTE*2):
            bit_planes[0].append(ROM.read('bits:8'))
            bit_planes[1].append(ROM.read('bits:8'))

        for i in range(2):
            bit_planes[i] = cls.__fillTileMatrix(bit_planes[i],height,width)

        sprite_data = cls.__combineBuffers(bit_planes,1)

        return cls(addr,width,height,sprite_data)

    @classmethod
    def decode_base64_sprite(cls, base64_sprite,width,height):
        decoded_sprite_bytes = base64.b64decode(base64_sprite)

        print(base64_sprite)

        sprite_array = []

        for data in decoded_sprite_bytes:
            for i in range(3,-1,-1):
                sprite_array.append((data >> (i*2)) & 0b11)

        sprite = []
        for i in range(0,int(len(sprite_array)),width*8):
            sprite.append(sprite_array[i:i+(width*8)])

        return cls(Addr(0,0),width,height,sprite)

class GBText:
    STRING_END = 0x50
    ALPHABET = {
        0x00: "",           #charmap "<NULL>"
        0x49: "^",       #charmap "<PAGE>"
        #charmap "<PKMN>",    #  "<PK><MN>"
        #charmap "<_CONT>",   #  implements "<CONT>"
        #charmap "<SCROLL>",  $4c
        0x4E: "<",     #Next
        0x4F: " ",
        0x57: "#",
        0x50: "@",   #charmap "@" string terminator
        0x51: "*",
        0x52: "A1",
        0x53: "A2",
        0x54: "POKé", #This is fine to leave multichar as it was only short hand for all four characters anyway
        0x55: "+",
        0x58: "$",
        0x5F: "}",   #charmap "<DEXEND>"
        0x75: "…",
        0x7F: " ",
        0x80: "A",
        0x81: "B",
        0x82: "C",
        0x83: "D",
        0x84: "E",
        0x85: "F",
        0x86: "G",
        0x87: "H",
        0x88: "I",
        0x89: "J",
        0x8A: "K",
        0x8B: "L",
        0x8C: "M",
        0x8D: "N",
        0x8E: "O",
        0x8F: "P",
        0x90: "Q",
        0x91: "R",
        0x92: "S",
        0x93: "T",
        0x94: "U",
        0x95: "V",
        0x96: "W",
        0x97: "X",
        0x98: "Y",
        0x99: "Z",
        0x9A: "(",
        0x9B: ")",
        0x9C: ":",
        0x9D: ";",
        0x9E: "[",
        0x9F: "]",
        0xA0: "a",
        0xA1: "b",
        0xA2: "c",
        0xA3: "d",
        0xA4: "e",
        0xA5: "f",
        0xA6: "g",
        0xA7: "h",
        0xA8: "i",
        0xA9: "j",
        0xAA: "k",
        0xAB: "l",
        0xAC: "m",
        0xAD: "n",
        0xAE: "o",
        0xAF: "p",
        0xB0: "q",
        0xB1: "r",
        0xB2: "s",
        0xB3: "t",
        0xB4: "u",
        0xB5: "v",
        0xB6: "w",
        0xB7: "x",
        0xB8: "y",
        0xB9: "z",
        0xBA: "é",
        0xBB: u"\u1E0B", #ḋ to represent 'd as one letter
        0xBC: u"\u013A", #ĺ to represent 'l as one letter
        0xBD: u"\u1E61", #ṡ to represent 's as one letter
        0xBE: u"\u1E6B", #ṫ to represent 't as one letter
        0xBF: u"\u1E7F", #ṿ to represent 'v as one letter
        0xE0: "'",
        0xE1: u"\u1D18", #ᴘ to represent PK as one letter
        0xE2: u"\u1D0D", #ᴍ to represent MN as one letter
        0xE3: "-",
        0xE4: u"\u1E59", #ṙ to represent 'r as one letter
        0xE5: u"\u1E41", #ṁ to represent 'm as one letter
        0xE6: "?",
        0xE7: "!",
        0xE8: ".",
        0xEC: "=",
        0xED: ">",
        0xEE: "_",
        0xEF: "♂",

        0x60: "<BOLD_A>",  #  unused
        0x61: "<BOLD_B>",  #  unused
        0x62: "<BOLD_C>",  #  unused
        0x63: "<BOLD_D>",  #  unused
        0x64: "<BOLD_E>",  #  unused
        0x65: "<BOLD_F>",  #  unused
        0x66: "<BOLD_G>",  #  unused
        0x67: "<BOLD_H>",  #  unused
        0x68: "<BOLD_I>",  #  unused
        0x69: "<BOLD_V>",
        0x6A: "<BOLD_S>",
        0x6B: "<BOLD_L>",  #  unused
        0x6C: "<BOLD_M>",  #  unused
        0x6D: "<COLON>",   #  colon with tinier dots than ":"
        0x6E: "ぃ",         #  hiragana small i, unused
        0x6F: "ぅ",         #  hiragana small u, unused
        0x70: "‘",         #  opening single quote
        0x71: "’",         #  closing single quote
        0x72: "“",         #  opening quote
        0x73: "”",         #  closing quote
        0x74: "·",         #  middle dot, unused
        0x75: "…",         #  ellipsis
        0x76: "ぁ",         #  hiragana small a, unused
        0x77: "ぇ",         #  hiragana small e, unused
        0x78: "ぉ",         #  hiragana small o, unused

        0x79: "┌",
        0x7A: "─",
        0x7B: "┐",
        0x7C: "│",
        0x7D: "└",
        0x7E: "┘",
        0x7F: " ",

        0xF0: "¥",
        0xF1: "×",
        0xF2: "<DOT>",
        0xF3: "/",
        0xF4: ",",
        0xF5: "♀",
        0xF6: "0",
        0xF7: "1",
        0xF8: "2",
        0xF9: "3",
        0xFA: "4",
        0xFB: "5",
        0xFC: "6",
        0xFD: "7",
        0xFE: "8",
        0xFF: "9"
    }

    def decodeText(self) -> str:
        return list(map(self.ALPHABET.get, self.packet.data))

    def __init__(self,packet) -> None:
        self.packet = packet
        self.text =  self.decodeText()

    def __str__(self):
        return "".join(self.text).strip('@')

    def __len__(self):
        return len(self.packet)

#Constants that have hard pointers in Red/Blue
ROM = ConstBitStream(filename='pokered.gbc')
BANK_SIZE = 0x4000
TWO_BPP_TILE_SIZE = 16
ONE_BPP_TILE_SIZE = 8
BYTE = 8
BIT = 1
NYBBLE = 4
TWO_BPP = 2
ONE_BPP = 1

POKEMON_NAME_LENGTH = 10

END_FILE = Addr.convert_to_addr(ROM.len/8)

POKEDEX_ORDER_POINTER = Addr(0x10,0x5024)
POKEDEX_ENTRY_POINTER = Addr(0x10,0x447e)
POKEMON_DATA_POINTER  = Addr(0X0E,0x43DE)
POKEMON_NAME_POINTER  = Addr(0x07,0x421e)
MOVE_NAME_POINTER     = Addr(0x2C,0x4000)
MOVES_DATA_POINTER    = Addr(0x0E,0x4000)
TM_HM_LIST_POINTER    = Addr(0x04,0x7773)
FONT_START_POINTER    = Addr(0x04,0x5a80)
EVO_TABLE_POINTER     = Addr(0x0E,0x705C)

datamap = {'Index to Pokedex':  [],
           'Pokedex Entry Loc': [],
           'EVO Table':         []
}

for i in range(0,380,2):
    datamap["Pokedex Entry Loc"].append(GBDataPacket.get_static_data(POKEDEX_ENTRY_POINTER+i,BYTE,2).collapse(rev=True))
    datamap["Index to Pokedex"].append(GBDataPacket.get_static_data(POKEDEX_ORDER_POINTER+int(i/2),BYTE,1).collapse())
    datamap['EVO Table'].append(GBDataPacket.get_static_data(EVO_TABLE_POINTER+i,BYTE,2).collapse(rev=True))


from __future__ import print_function

import os
from glob import glob
from tqdm import trange
from itertools import chain

import torch
from torch import nn
import torch.nn.parallel
import torchvision.utils as vutils
from torch.autograd import Variable

from models import *
from data_loader import get_loader

def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        m.weight.data.normal_(0.0, 0.02)
    elif classname.find('BatchNorm') != -1:
        m.weight.data.normal_(1.0, 0.02)
        m.bias.data.fill_(0)

class Trainer(object):
    def __init__(self, config, a_data_loader, b_data_loader):
        self.config = config

        self.a_data_loader = a_data_loader
        self.b_data_loader = b_data_loader

        self.num_gpu = config.num_gpu
        self.dataset = config.dataset

        self.loss = config.loss
        self.lr = config.lr
        self.beta1 = config.beta1
        self.beta2 = config.beta2
        self.optimizer = config.optimizer
        self.batch_size = config.batch_size
        self.weight_decay = config.weight_decay
        self.cnn_type = config.cnn_type

        self.model_dir = config.model_dir
        self.load_path = config.load_path

        self.start_step = 0
        self.log_step = config.log_step
        self.max_step = config.max_step
        self.save_step = config.save_step

        self.build_model()

        if self.num_gpu == 1:
            self.G_AB.cuda()
            self.G_BA.cuda()
            self.D_A.cuda()
            self.D_B.cuda()

        elif self.num_gpu > 1:
            self.G_AB = nn.DataParallel(self.G_AB.cuda(),device_ids=range(self.num_gpu))
            self.G_BA = nn.DataParallel(self.G_BA.cuda(),device_ids=range(self.num_gpu))
            self.D_A = nn.DataParallel(self.D_A.cuda(),device_ids=range(self.num_gpu))
            self.D_B = nn.DataParallel(self.D_B.cuda(),device_ids=range(self.num_gpu))

        if self.load_path:
            self.load_model()

    def build_model(self):
        if self.dataset == 'toy':
            self.G_AB = GeneratorFC(2, 2, [config.fc_hidden_dim] * config.g_num_layer)
            self.G_BA = GeneratorFC(2, 2, [config.fc_hidden_dim] * config.g_num_layer)

            self.D_A = DiscriminatorFC(2, 1, [config.fc_hidden_dim] * config.d_num_layer)
            self.D_B = DiscriminatorFC(2, 1, [config.fc_hidden_dim] * config.d_num_layer)
        else:
            a_height, a_width, a_channel = self.a_data_loader.shape
            b_height, b_width, b_channel = self.b_data_loader.shape

            if self.cnn_type == 0:
                #conv_dims, deconv_dims = [64, 128, 256, 512], [512, 256, 128, 64]
                conv_dims, deconv_dims = [64, 128, 256, 512], [256, 128, 64]
            elif self.cnn_type == 1:
                #conv_dims, deconv_dims = [32, 64, 128, 256], [256, 128, 64, 32]
                conv_dims, deconv_dims = [32, 64, 128, 256], [128, 64, 32]
            else:
                raise Exception("[!] cnn_type {} is not defined".format(self.cnn_type))

            self.G_AB = GeneratorCNN(
                    a_channel, b_channel, conv_dims, deconv_dims, self.num_gpu)
            self.G_BA = GeneratorCNN(
                    b_channel, a_channel, conv_dims, deconv_dims, self.num_gpu)

            self.D_A = DiscriminatorCNN(
                    a_channel, 1, conv_dims, self.num_gpu)
            self.D_B = DiscriminatorCNN(
                    b_channel, 1, conv_dims, self.num_gpu)

            self.G_AB.apply(weights_init)
            self.G_BA.apply(weights_init)

            self.D_A.apply(weights_init)
            self.D_B.apply(weights_init)

    def load_model(self):
        print("[*] Load models from {}...".format(self.load_path))

        paths = glob(os.path.join(self.load_path, 'G_AB_*.pth'))
        paths.sort()

        if len(paths) == 0:
            print("[!] No checkpoint found in {}...".format(self.load_path))
            return

        idxes = [int(os.path.basename(path.split('.')[0].split('_')[-1])) for path in paths]
        self.start_step = max(idxes)

        if self.num_gpu == 0:
            map_location = lambda storage, loc: storage
        else:
            map_location = None

        G_AB_filename = '{}/G_AB_{}.pth'.format(self.load_path, self.start_step)
        self.G_AB.load_state_dict(torch.load(G_AB_filename, map_location=map_location))
        self.G_BA.load_state_dict(
            torch.load('{}/G_BA_{}.pth'.format(self.load_path, self.start_step), map_location=map_location))

        self.D_A.load_state_dict(
            torch.load('{}/D_A_{}.pth'.format(self.load_path, self.start_step), map_location=map_location))
        self.D_B.load_state_dict(
            torch.load('{}/D_B_{}.pth'.format(self.load_path, self.start_step), map_location=map_location))

        print("[*] Model loaded: {}".format(G_AB_filename))

    def train(self):
        d = nn.MSELoss()
        bce = nn.BCELoss()

        real_label = 1
        fake_label = 0

        real_tensor = Variable(torch.FloatTensor(self.batch_size))
        _ = real_tensor.data.fill_(real_label)

        fake_tensor = Variable(torch.FloatTensor(self.batch_size))
        _ = fake_tensor.data.fill_(fake_label)

        if self.num_gpu > 0:
            d.cuda()
            bce.cuda()

            real_tensor = real_tensor.cuda()
            fake_tensor = fake_tensor.cuda()

        if self.optimizer == 'adam':
            optimizer = torch.optim.Adam
        else:
            raise Exception("[!] Caution! Paper didn't use {} opimizer other than Adam".format(config.optimizer))

        optimizer_d = optimizer(
            chain(self.D_A.parameters(), self.D_B.parameters()),
            lr=self.lr, betas=(self.beta1, self.beta2), weight_decay=self.weight_decay)
        optimizer_g = optimizer(
            chain(self.G_AB.parameters(), self.G_BA.parameters()),
            lr=self.lr, betas=(self.beta1, self.beta2))

        A_loader, B_loader = iter(self.a_data_loader), iter(self.b_data_loader)
        valid_x_A, valid_x_B = self._get_variable(A_loader.next()), self._get_variable(B_loader.next())

        vutils.save_image(valid_x_A.data, '{}/valid_x_A.png'.format(self.model_dir))
        vutils.save_image(valid_x_B.data, '{}/valid_x_B.png'.format(self.model_dir))

        for step in trange(self.start_step, self.max_step):
            try:
                x_A, x_B = A_loader.next(), B_loader.next()
            except StopIteration:
                A_loader, B_loader = iter(self.a_data_loader), iter(self.b_data_loader)
                x_A, x_B = A_loader.next(), B_loader.next()
            if x_A.size(0) != x_B.size(0):
                print("[!] Sampled dataset from A and B have different # of data. Try resampling...")
                continue

            x_A, x_B = self._get_variable(x_A), self._get_variable(x_B)

            batch_size = x_A.size(0)
            real_tensor.data.resize_(batch_size).fill_(real_label)
            fake_tensor.data.resize_(batch_size).fill_(fake_label)

            # update D network
            self.D_A.zero_grad()
            self.D_B.zero_grad()

            x_AB = self.G_AB(x_A).detach()
            x_BA = self.G_BA(x_B).detach()

            x_ABA = self.G_BA(x_AB).detach()
            x_BAB = self.G_AB(x_BA).detach()

            if self.loss == "log_prob":
                l_d_A_real, l_d_A_fake = bce(self.D_A(x_A), real_tensor), bce(self.D_A(x_BA), fake_tensor)
                l_d_B_real, l_d_B_fake = bce(self.D_B(x_B), real_tensor), bce(self.D_B(x_AB), fake_tensor)
            elif self.loss == "least_square":
                l_d_A_real, l_d_A_fake = \
                    0.5 * torch.mean((self.D_A(x_A) - 1)**2), 0.5 * torch.mean((self.D_A(x_BA))**2)
                l_d_B_real, l_d_B_fake = \
                    0.5 * torch.mean((self.D_B(x_B) - 1)**2), 0.5 * torch.mean((self.D_B(x_AB))**2)
            else:
                raise Exception("[!] Unkown loss type: {}".format(self.loss))

            l_d_A = l_d_A_real + l_d_A_fake
            l_d_B = l_d_B_real + l_d_B_fake

            l_d = l_d_A + l_d_B

            l_d.backward()
            optimizer_d.step()

            # update G network
            self.G_AB.zero_grad()
            self.G_BA.zero_grad()

            x_AB = self.G_AB(x_A)
            x_BA = self.G_BA(x_B)

            x_ABA = self.G_BA(x_AB)
            x_BAB = self.G_AB(x_BA)

            l_const_A = d(x_ABA, x_A)
            l_const_B = d(x_BAB, x_B)

            if self.loss == "log_prob":
                l_gan_A = bce(self.D_A(x_BA), real_tensor)
                l_gan_B = bce(self.D_B(x_AB), real_tensor)
            elif self.loss == "least_square":
                l_gan_A = 0.5 * torch.mean((self.D_A(x_BA) - 1)**2)
                l_gan_B = 0.5 * torch.mean((self.D_B(x_AB) - 1)**2)
            else:
                raise Exception("[!] Unkown loss type: {}".format(self.loss))

            l_g = l_gan_A + l_gan_B + l_const_A + l_const_B

            l_g.backward()
            optimizer_g.step()

            if step % self.log_step == 0:
                print("[{}/{}] Loss_D: {:.4f} Loss_G: {:.4f}". \
                      format(step, self.max_step, l_d.data[0], l_g.data[0]))

                print("[{}/{}] l_d_A_real: {:.4f} l_d_A_fake: {:.4f}, l_d_B_real: {:.4f}, l_d_B_fake: {:.4f}". \
                      format(step, self.max_step, l_d_A_real.data[0], l_d_A_fake.data[0],
                             l_d_B_real.data[0], l_d_B_fake.data[0]))

                print("[{}/{}] l_const_A: {:.4f} l_const_B: {:.4f}, l_gan_A: {:.4f}, l_gan_B: {:.4f}". \
                      format(step, self.max_step, l_const_A.data[0], l_const_B.data[0],
                             l_gan_A.data[0], l_gan_B.data[0]))

                self.generate_with_A(valid_x_A, self.model_dir, idx=step)
                self.generate_with_B(valid_x_B, self.model_dir, idx=step)

            if step % self.save_step == self.save_step - 1:
                print("[*] Save models to {}...".format(self.model_dir))

                torch.save(self.G_AB.state_dict(), '{}/G_AB_{}.pth'.format(self.model_dir, step))
                torch.save(self.G_BA.state_dict(), '{}/G_BA_{}.pth'.format(self.model_dir, step))

                torch.save(self.D_A.state_dict(), '{}/D_A_{}.pth'.format(self.model_dir, step))
                torch.save(self.D_B.state_dict(), '{}/D_B_{}.pth'.format(self.model_dir, step))

    def generate_with_A(self, inputs, path, idx=None):
        x_AB = self.G_AB(inputs)
        x_ABA = self.G_BA(x_AB)

        x_AB_path = '{}/{}_x_AB.png'.format(path, idx)
        x_ABA_path = '{}/{}_x_ABA.png'.format(path, idx)

        vutils.save_image(x_AB.data, x_AB_path)
        print("[*] Samples saved: {}".format(x_AB_path))

        vutils.save_image(x_ABA.data, x_ABA_path)
        print("[*] Samples saved: {}".format(x_ABA_path))

    def generate_with_B(self, inputs, path, idx=None):
        x_BA = self.G_BA(inputs)
        x_BAB = self.G_AB(x_BA)

        x_BA_path = '{}/{}_x_BA.png'.format(path, idx)
        x_BAB_path = '{}/{}_x_BAB.png'.format(path, idx)

        vutils.save_image(x_BA.data, x_BA_path)
        print("[*] Samples saved: {}".format(x_BA_path))

        vutils.save_image(x_BAB.data, x_BAB_path)
        print("[*] Samples saved: {}".format(x_BAB_path))

    def generate_infinitely(self, inputs, path, input_type, count=10, nrow=2, idx=None):
        if input_type.lower() == "a":
            iterator = [self.G_AB, self.G_BA] * count
        elif input_type.lower() == "b":
            iterator = [self.G_BA, self.G_AB] * count

        out = inputs
        for step, model in enumerate(iterator):
            out = model(out)

            out_path = '{}/{}_x_{}_#{}.png'.format(path, idx, input_type, step)
            vutils.save_image(out.data, out_path, nrow=nrow)
            print("[*] Samples saved: {}".format(out_path))

    def test(self):
        batch_size = self.config.sample_per_image
        A_loader, B_loader = iter(self.a_data_loader), iter(self.b_data_loader)

        test_dir = os.path.join(self.model_dir, 'test')
        if not os.path.exists(test_dir):
            os.makedirs(test_dir)

        step = 0
        while True:
            try:
                x_A, x_B = self._get_variable(A_loader.next()), self._get_variable(B_loader.next())
            except StopIteration:
                print("[!] Test sample generation finished. Samples are in {}".format(test_dir))
                break

            vutils.save_image(x_A.data, '{}/{}_x_A.png'.format(test_dir, step))
            vutils.save_image(x_B.data, '{}/{}_x_B.png'.format(test_dir, step))

            self.generate_with_A(x_A, test_dir, idx=step)
            self.generate_with_B(x_B, test_dir, idx=step)

            self.generate_infinitely(x_A, test_dir, input_type="A", count=10, nrow=4, idx=step)
            self.generate_infinitely(x_B, test_dir, input_type="B", count=10, nrow=4, idx=step)

            step += 1

    def _get_variable(self, inputs):
        if self.num_gpu > 0:
            out = Variable(inputs.cuda())
        else:
            out = Variable(inputs)
        return out