#!/usr/bin/env python3

import os
import sys
import time

# import PyQt5 # hint for pyinstaller
from PyQt5.QtWidgets import QApplication
from PyQt5.QtMultimedia import QCameraInfo, QCamera

"""
Copyright (C) 2019 Denis Polygalov,
Laboratory for Circuit and Behavioral Physiology,
RIKEN Center for Brain Science, Saitama, Japan.

This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation; either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program; if not, a copy is available at
http://www.fsf.org/
"""

# TODO change output format to json instead of print()-ing each line

def captureMode2str(mode):
    if mode == QCamera.CaptureStillImage:
        return "CaptureStillImage"
    elif mode == QCamera.CaptureViewfinder:
        return "CaptureViewfinder"
    elif mode == QCamera.CaptureVideo:
        return "CaptureVideo"
    else:
        return "UNKNOWN CAPTURE MODE"
#

def lockStatus2str(status):
    if status == QCamera.Unlocked:
        return "Unlocked"
    elif status == QCamera.Searching:
        return "Searching"
    elif status == QCamera.Locked:
        return "Locked"
    else:
        return "UNKNOWN LOCK STATUS"
#

def state2str(state):
    if state == QCamera.UnloadedState:
        return "UnloadedState"
    elif state == QCamera.LoadedState:
        return "LoadedState"
    elif state == QCamera.ActiveState:
        return "ActiveState"
    else:
        return "UNKNOWN STATE"
#

def status2str(status):
    if status == QCamera.ActiveStatus:
        return "ActiveStatus"
    elif status == QCamera.StartingStatus:
        return "StartingStatus"
    elif status == QCamera.StandbyStatus:
        return "StandbyStatus"
    elif status == QCamera.LoadedStatus:
        return "LoadedStatus"
    elif status == QCamera.LoadingStatus:
        return "LoadingStatus"
    elif status == QCamera.UnloadingStatus:
        return "UnloadingStatus"
    elif status == QCamera.StoppingStatus:
        return "StoppingStatus"
    elif status == QCamera.UnloadedStatus:
        return "UnloadedStatus"
    elif status == QCamera.UnavailableStatus:
        return "UnavailableStatus"
    else:
        return "UNKNOWN STATUS"
#

def lockType2str(status):
    # TODO not sure this way is the correct one
    l_locks = []
    if (status == QCamera.NoLock):
        l_locks.append("NoLock")
    if (status == QCamera.LockExposure):
        l_locks.append("LockExposure")
    if (status == QCamera.LockWhiteBalance):
        l_locks.append("LockWhiteBalance")
    if (status == QCamera.LockFocus):
        l_locks.append("LockFocus")
    return ",".join(l_locks)
#

def frameRateRange2str(l_frate_ranges):
    l_out = []
    for oc_range in l_frate_ranges:
        l_out.append("%.1f - %.1f" % (oc_range.minimumFrameRate, oc_range.maximumFrameRate))
    return ", ".join(l_out)
#

def frameResolution2str(l_resolutions):
    l_out = []
    for oc_res in l_resolutions:
        l_out.append("(%i, %i)" % (oc_res.width(), oc_res.height()))
    return "\n\t\t".join(l_out)
#

def camera_sync_load_and_start(oc_qcamera):
    cam_status = oc_qcamera.status()
    if cam_status != QCamera.UnloadedStatus:
        print("ERROR: unexpected camera status: %s" % status2str(cam_status))
        sys.exit(-1)
    # print("INFO: attempting to LOAD camera from %s..." % status2str(cam_status))
    i_sec_cnt = 0
    oc_qcamera.load()
    while True:
        cam_status = oc_qcamera.status()
        # print("INFO: current status: %s" % status2str(cam_status))
        if cam_status == QCamera.LoadedStatus: break
        else:
            time.sleep(1)
            i_sec_cnt += 1
            if i_sec_cnt >= 10:
                print("ERROR: unable to load camera")
                sys.exit(-1)
    # print("INFO: attempting to START camera from %s..." % status2str(cam_status))
    i_sec_cnt = 0
    oc_qcamera.start()
    while True:
        cam_status = oc_qcamera.status()
        # print("INFO: current status: %s" % status2str(cam_status))
        if cam_status == QCamera.ActiveStatus: break
        else:
            time.sleep(1)
            i_sec_cnt += 1
            if i_sec_cnt >= 10:
                print("ERROR: unable to start camera")
                sys.exit(-1)
#

def camera_sync_stop_and_unload(oc_qcamera):
    cam_status = oc_qcamera.status()
    if cam_status != QCamera.ActiveStatus:
        print("ERROR: unexpected camera status: %s" % status2str(cam_status))
        sys.exit(-1)
    # print("INFO: attempting to STOP camera from %s..." % status2str(cam_status))
    i_sec_cnt = 0
    oc_qcamera.stop()
    while True:
        cam_status = oc_qcamera.status()
        # print("INFO: current status: %s" % status2str(cam_status))
        if cam_status == QCamera.LoadedStatus: break
        else:
            time.sleep(1)
            i_sec_cnt += 1
            if i_sec_cnt >= 10:
                print("ERROR: unable to stop camera")
                sys.exit(-1)
    # print("INFO: attempting to (UN)load camera from %s..." % status2str(cam_status))
    i_sec_cnt = 0
    oc_qcamera.unload()
    while True:
        cam_status = oc_qcamera.status()
        # print("INFO: current status: %s" % status2str(cam_status))
        if cam_status == QCamera.UnloadedStatus: break
        else:
            time.sleep(1)
            i_sec_cnt += 1
            if i_sec_cnt >= 10:
                print("ERROR: unable to unload camera")
                sys.exit(-1)
#

def main():
    app = QApplication(sys.argv) # must be called in order to be able to use all the stuff below(!)
    print("The application is executed on: %s" % app.platformName())
    l_cameras = QCameraInfo.availableCameras()
    if len(l_cameras) == 0:
        print("ERROR: no cameras found")
        sys.exit(-1)
    else:
        print("INFO: found %i camera(s)" % len(l_cameras))
        print()

    for i_idx, oc_caminfo in enumerate(l_cameras):
        # QCameraInfo - derived properties
        print("Camera #%i:" % i_idx)
        print("\t deviceName: %s" % oc_caminfo.deviceName())
        print("\t description: %s" % oc_caminfo.description())
        print("\t isNull: %s" % oc_caminfo.isNull())
        print("\t orientation: %s" % oc_caminfo.orientation())
        print("\t position: %s" % oc_caminfo.position())

        # QCamera - derived properties
        oc_camera = QCamera(i_idx)
        camera_sync_load_and_start(oc_camera)

        oc_camera_exposure = oc_camera.exposure()
        print("\t captureMode: %s" % captureMode2str(oc_camera.captureMode()))
        print("\t exposure:")
        print("\t\t isAvailabe: %s" % oc_camera_exposure.isAvailable())
        if oc_camera_exposure.isAvailable():
            print("\t\t aperture: %s" % oc_camera_exposure.aperture())
            # TODO test this, add more entries

        oc_camera_focus = oc_camera.focus()
        print("\t focus:")
        print("\t\t isAvailabe: %s" % oc_camera_focus.isAvailable())
        if oc_camera_focus.isAvailable():
            print("\t\t aperture: %f" % oc_camera_focus.digitalZoom())
            # TODO test this, add more entries

        oc_cam_img_proc = oc_camera.imageProcessing()
        print("\t imageProcessing:")
        print("\t\t isAvailabe: %s" % oc_cam_img_proc.isAvailable())
        if oc_cam_img_proc.isAvailable():
            print("\t\t brightness: %f" % oc_cam_img_proc.brightness())
            print("\t\t contrast: %f" % oc_cam_img_proc.contrast())
            print("\t\t denoisingLevel: %f" % oc_cam_img_proc.denoisingLevel())
            print("\t\t manualWhiteBalance: %f" % oc_cam_img_proc.manualWhiteBalance())
            print("\t\t saturation: %f" % oc_cam_img_proc.saturation())
            print("\t\t sharpeningLevel: %f" % oc_cam_img_proc.sharpeningLevel())

        print("\t isVideoCaptureSupported: %s" % oc_camera.isCaptureModeSupported(QCamera.CaptureVideo))
        print("\t lockStatus: %s" % lockStatus2str(oc_camera.lockStatus()))
        print("\t requestedLocks: %s" % lockType2str(oc_camera.requestedLocks()))
        print("\t state: %s" % state2str(oc_camera.state()))
        print("\t status: %s" % status2str(oc_camera.status()))
        print("\t supportedLocks: %s" % lockType2str(oc_camera.supportedLocks()))
        print("\t supportedViewfinderFrameRateRanges: %s" % frameRateRange2str(oc_camera.supportedViewfinderFrameRateRanges()))
        print("\t supportedViewfinderPixelFormats: %s" % repr(oc_camera.supportedViewfinderPixelFormats()))
        print("\t supportedViewfinderResolutions: \n\t\t%s" % frameResolution2str(oc_camera.supportedViewfinderResolutions()))
        print("\t len(supportedViewfinderSettings): %s" % len(oc_camera.supportedViewfinderSettings()))

        # QCameraViewfinderSettings - derived properties
        oc_vf_settings = oc_camera.viewfinderSettings()
        if oc_vf_settings.isNull():
            print("\t viewfinderSettings: not supported")
            camera_sync_stop_and_unload(oc_camera)
            print()
            continue
        print("\t maximumFrameRate: %f" % oc_vf_settings.maximumFrameRate())
        print("\t minimumFrameRate: %f" % oc_vf_settings.minimumFrameRate())
        print("\t resolution: %s" % frameResolution2str([oc_vf_settings.resolution()]))
        # TODO the rest of methods...
        camera_sync_stop_and_unload(oc_camera)
        print()
    #
#

if __name__ == '__main__':
    s_qt_plugin_path = os.path.join(os.getcwd(), 'PyQt5', 'Qt', 'plugins')
    if os.path.isdir(s_qt_plugin_path):
        os.environ['QT_PLUGIN_PATH'] = s_qt_plugin_path
        print("INFO: QT_PLUGIN_PATH is set to ", os.environ['QT_PLUGIN_PATH'])
    else:
        print("INFO: QT_PLUGIN_PATH was not altered")
    main()
    print("Press ENTER to continue...")
    input()
#


# -*- coding: utf-8 -*-
from __future__ import unicode_literals

from django.db import models, migrations
from django.conf import settings

class Migration(migrations.Migration):

    dependencies = [
        migrations.swappable_dependency(settings.AUTH_USER_MODEL),
        ('research', '0001_initial'),
    ]

    operations = [
        migrations.AddField(
            model_name='studyparticipation',
            name='created',
            field=models.DateTimeField(null=True, blank=True),
        ),
        migrations.AddField(
            model_name='studyparticipation',
            name='created_by',
            field=models.ForeignKey(related_name='created_research_studyparticipation_subrecords', blank=True, to=settings.AUTH_USER_MODEL, null=True),
        ),
        migrations.AddField(
            model_name='studyparticipation',
            name='updated',
            field=models.DateTimeField(null=True, blank=True),
        ),
        migrations.AddField(
            model_name='studyparticipation',
            name='updated_by',
            field=models.ForeignKey(related_name='updated_research_studyparticipation_subrecords', blank=True, to=settings.AUTH_USER_MODEL, null=True),
        ),
    ]


from random import randint
from BoardClasses import Move
from BoardClasses import Board
from math import sqrt, log
import sys
#The following part should be completed by students.
#Students can modify anything except the class name and exisiting functions and varibles.

class Node():
    #class for nodes in MCTS
    def __init__(self, color, move = None, parent = None):
        self.move = move #the move parent made to get to this node. If None, must be root node
        self.color = color #color of whose move it is. Always inverse to the parents
        self.num_simulations: int = 0 #si in slides
        self.num_parent_wins: int = 0 #wi in slides
        self.children = [] #list of child nodes
        self.parent = parent #if root, should be null
        self.current_child = 0 #index of what child we are currently expanding
        self.uct = 0

    def get_uct(self) -> float:
        return self.uct

    def get_current_child(self):
        'returns the current child to be expanded'
        self.current_child += 1
        return self.children[self.current_child - 1]

    def add_children(self, moves: [[Move]], opp_color):
        'extends the list of children when given a list of moves'
        self.children.extend((Node(opp_color, move, self)  for piece in moves for move in piece))

    def back_propogate(self, winning_color): # returns the number of undos that need to be performed
        'adds the result to the stats and propogates it up to the parent'
        self.num_simulations += 1
        if winning_color != self.color: # this is a win for the move made associated with this node
            self.num_parent_wins += 1 #if the node is black, and the winning color is white, it should count as a win since the move associated with this node was made by a white piece

        if self.parent:
            self.uct = (1.4 * sqrt(log(self.parent.num_simulations+1) / self.num_simulations)) + (self.num_parent_wins / self.num_simulations)
            return self.parent.back_propogate(winning_color) + 1 # it will alternate back and forth if it was a win, since half the nodes will be opponent moves that will count this as a loss
        return 0

    def pick_best_move(self) -> Move: #this function should only be called by the root node
        'returns the best move'
        return max(self.children, key = lambda x: x.num_parent_wins / x.num_simulations).move

    def is_leaf_node(self) -> bool:
        'returns whether the current node is a leaf node'
        return len(self.children) == 0

    def is_fully_expanded(self) -> bool:
        'returns whether all the children of this node have been expanded into'
        return len(self.children) == self.current_child

class StudentAI():

    def __init__(self,col,row,p):
        self.col = col
        self.row = row
        self.p = p
        self.board = Board(col,row,p)
        self.board.initialize_game()
        self.color = ''
        self.opponent = {1:2,2:1}
        self.color = 2
    def get_move(self,move):
        print('best move ')
        if len(move) != 0:
            self.board.make_move(move,self.opponent[self.color])
        else:
            self.color = 1

        #print(f'color = {self.color}')
        moves = self.board.get_all_possible_moves(self.color)
        #best_move, heuristic = self.recur_best_move(moves, 3, self.color)
        best_move = self.monte_carlo()
        self.board.make_move(best_move, self.color)
        return best_move

    def move_kings(self, move, color) -> bool: # returns a boolean for if the move kings or not
        if color == 2:
            if move[-1][0] == 0:
                return True
            return False
        if color == 1:
            if move[-1][0] == self.board.row - 1:
                return True
            return False

    def get_heuristic(self, color): # returns the piece differential for the selected color
        heuristic = 0
        for row in range(self.board.row):
                for col in range(self.board.col):
                    if self.board.board[row][col].is_king and self.board.board[row][col].color == color:
                        heuristic += 3

        if color == 2: #we are counting the white differential
            return heuristic + self.board.white_count - self.board.black_count

        return heuristic + self.board.black_count - self.board.white_count

    def get_opponent_best_move(self, move): # returns the best move the opponent can make
        opp_color = 1
        if self.color == 1:
            opp_color = 2 #this ensures we always know the opponents color for testing
        self.board.make_move(move, self.color)

        moves = self.board.get_all_possible_moves(opp_color)

        heuristic = None
        best_move = None
        for row in moves:
            for move in row:
                self.board.make_move(move, opp_color)

                temp = self.get_heuristic(opp_color)

                if heuristic == None or temp > heuristic:
                    heuristic = temp
                    best_move = move
                self.board.undo()

        self.board.undo()
        return heuristic

    def find_best_move(self, moves):
        heuristic = None
        best_move = None
        for piece in moves:
            for move in piece:
                if self.move_kings(move, self.color):
                    row, col = move[0]
                    if not self.board.board[row][col].is_king:
                        best_move = move
                        self.board.make_move(best_move, self.color)
                        return best_move

                opp_hueristic = self.get_opponent_best_move(move)
                if heuristic == None or opp_hueristic < heuristic:
                    heuristic = opp_hueristic
                    best_move = move

        self.board.make_move(best_move, self.color)
        return best_move

    def recur_best_move(self, moves, depth, color):
        '''
        print(f'depth = {depth} color = {color}')
        print(moves)
        print()
        '''
        if depth == 0:
            return (None, self.get_heuristic(self.color))
        elif len(moves) == 0:
            return (None, float('-inf'))
        heuristic = None
        best_move = None
        for piece in moves:
            for move in piece:
                #print(f'heuristic = {heuristic} depth = {depth}')
                self.board.make_move(move, color)
                #print(move)
                curr_moves = self.board.get_all_possible_moves(self.opponent[color])

                move1, temp = self.recur_best_move(curr_moves, depth-1, self.opponent[color])
                if color == self.color:
                    #Maximizing
                    if heuristic is None or temp > heuristic:
                        heuristic = temp
                        best_move = move
                else:
                    #Minimizing
                    if heuristic is None or temp < heuristic:
                        heuristic = temp
                        best_move = move

                self.board.undo()
                #print(f'ending heuristic = {heuristic} depth = {depth}')

        #print(f'returning {heuristic} depth = {depth}')
        return (best_move, heuristic)

    def find_next_node(self, root):  # this is selection AND expansion
        'returns the node for MCTS to simulate off of'
        curr_node = root

        while not curr_node.is_leaf_node() and curr_node.is_fully_expanded(): #selection; keeps going till a node is terminal or has unexplored children
            curr_node = self.best_uct(curr_node) # best child UCT
            #print(f"Making Move: {curr_node.move}")
            self.board.make_move(curr_node.move, self.opponent[curr_node.color]) #keeps the board up to date for rollout purposes

        if curr_node.is_leaf_node(): #we need to add the children of the current node into the tree if it has no children
            curr_node.add_children(self.board.get_all_possible_moves(curr_node.color), self.opponent[curr_node.color]) # if terminal node, adds the children

        result =  curr_node.get_current_child() # returns new unexplored node
        #print(f"Making Move: {result.move}")
        self.board.make_move(result.move, self.opponent[result.color])
        return result

    def pick_random_move(self, moves):
        index = randint(0,len(moves)-1)
        inner_index =  randint(0,len(moves[index])-1)
        return moves[index][inner_index]

    def simulation(self, node): #rollout for MCTS, random moves
        current_color = self.opponent[node.color]
        num_moves = 0
        while not self.board.is_win(current_color):
            current_color = self.opponent[current_color] # flip color
            moves = self.board.get_all_possible_moves(current_color)
            self.board.make_move(self.pick_random_move(moves), current_color) #make random move
            num_moves+=1

        for x in range(num_moves):
            self.board.undo()
        return current_color

    def best_uct(self, node):
        return max(node.children, key = lambda x: x.get_uct())

    def monte_carlo(self, iterations = 500):

        root = Node(self.color)
        root.add_children(self.board.get_all_possible_moves(self.color), self.opponent[self.color])
        for _x in range(iterations):
            #print(f'mcts loop {_x+1}')
            node = self.find_next_node(root)
            winning_color = self.simulation(node) #do simulation here
            undos = node.back_propogate(winning_color)

            for _x in range(undos):
                self.board.undo()

        return root.pick_best_move()


#!/usr/bin/env python

"""The setup script."""

from setuptools import setup, find_packages

file_requirements = open('requirements.txt', 'r')
requirements = file_requirements.readlines()
requirements = [line.replace("\n", "") for line in requirements]

setup(
    author="Camilo Caceres F",
    python_requires='>=3.8',
    classifiers=[
        'Development Status :: 2 - Pre-Alpha',
        'Intended Audience :: Developers',
        'License :: OSI Approved :: GNU General Public License v3 (GPLv3)',
        'Natural Language :: English',
        'Programming Language :: Python :: 3',
        'Programming Language :: Python :: 3.8',
    ],
    description="Image API converter from blackboard to B/W.",
    install_requires=requirements,
    license="GNU General Public License v3",
    include_package_data=True,
    keywords='blackboard_photo2cleanimage',
    name='blackboard_photo2cleanimage',
    packages=find_packages(include=['blackboard_photo2cleanimage', 'blackboard_photo2cleanimage.*']),
    test_suite='tests',
    url='https://github.com/camilo-cf/blackboard_photo2cleanimage',
    version='0.1.0',
    zip_safe=False,
)


r"""Helper library for handling packets from the microstrain device.

Helper dataclasses are introduced:
  MipsField: are strung together as a list and packed into the payload portion
    of the packet.
  MipsPacket: the top level container that describes the payload and holds all
    the data.
  ReplyFormats: provide a mapping between packets and the data to make it easier
    to decode.
"""

import collections
from dataclasses import dataclass
import enum
import itertools
import logging
import struct
from typing import Any, Optional

# The sync headers are always the same, and indicate the start of a packet.
SYNC1 = 0x75
SYNC2 = 0x65
SYNC_MSG = bytearray([SYNC1, SYNC2])

class DataMessages(enum.Enum):
    """Small class to bundle interactions with the different message formats.

    Attributes:
      msg_ind: the packet descriptor for a message of that type.
      fmt_cmd: the field descriptor portion of the format command.
      poll_cmd: the field desecriptor portion of the poll command.
    """
    IMU = (0x80, 0x08, 0x01, 0x06)
    GPS = (0x81, 0x09, 0x02, 0x07)
    EKF = (0x82, 0x0a, 0x03, 0x0b)

    def __init__(self, msg_ind, fmt_cmd, poll_cmd, rate_cmd):
        self.msg_ind = msg_ind
        self.fmt_cmd = fmt_cmd
        self.poll_cmd = poll_cmd
        self.rate_cmd = rate_cmd

@dataclass
class MipsField:
    """Container that represents the payload data in the MipsPacket.

    For convenience, the data held in the MipsField is represented in both byte
    and as a list of integers. This gives some flexibility in how it is used in
    the code, and makes it easier to reference the documentation.

    Attributes:
      field_len: the length of the field in bytes.
      field_desc: the field descriptor explaining the data contents.
      data: the data of the field as a list of integers.
      data_bytes: the data of the field in byte format.
    """
    field_len: int
    field_desc: int
    data: Optional[list[int]] = None
    data_bytes: Optional[bytes] = None

    def __init__(self, field_len: int, field_desc: int,
                 data: Optional[list[int]] = None,
                 data_bytes: Optional[bytes] = None):
        self.field_len = field_len
        self.field_desc = field_desc
        # Ensure we load the data from one source, and make sure the alternate
        # form is consistent.
        if data is not None:
            self.data = list(data)
            self.data_bytes = bytearray(data)
        elif data_bytes is not None:
            self.data_bytes = data_bytes
            self.data = [i for i in data_bytes]
        else:
            self.data = []
            self.data_bytes = b''

    def __len__(self):
        """Returns length of field data."""
        data_len = len(self.data) if self.data else 0
        # Additional 2 for field length and descriptor byte.
        return 2 + data_len

    def __iter__(self):
        """Returns the contents of class (used for unpacking)."""
        content = [self.field_len, self.field_desc]
        if self.data:
            content += self.data
        return iter(content)

    @property
    def as_bytes(self) -> bytes:
        field_bytes = bytearray([self.field_len, self.field_desc])
        field_bytes += self.data_bytes
        return field_bytes

@dataclass
class MipsPacket:
    """Container for communication with the device.

    This class is responsible for properly composing the messages to send to the
    device. It stores the essential information only, skipping always needed
    items like the SYNC and checksum components.

    Attributes:
      desc_set: the descriptor set byte, explaining which message catalog is
        being sent to the device.
      payload: a list of MipsField objects, which contain the data to send or
        what was received from the device.
      recv_time: a field that can be used to store time a streamed packet
        arrived, used for decoding and writing out messages.
    """
    desc_set: int
    payload: list[MipsField]
    recv_time: Optional[float] = None

    def __init__(self, desc_set: int, payload: list[MipsField],
                 recv_time: Optional[float] = None):
        self.desc_set = desc_set
        self.payload = payload
        self.payload_len = sum([len(f) for f in self.payload])
        self.recv_time = recv_time

    @property
    def as_bytes(self) -> bytes:
        """Returns the full packet info as bytes (used to send to device)."""
        field_bytes = []
        for entry in self.payload:
            field_bytes.extend(entry)
        packet = bytearray([SYNC1, SYNC2, self.desc_set, self.payload_len]
                           + field_bytes)
        packet += fletcher_checksum(packet)
        return packet

    def convert_payload(self) -> list[collections.namedtuple]:
        """Converts payload fields to their respective objects."""
        entries = []
        for f in self.payload:
            try:
              value = ReplyFormats.process_field(self.desc_set, f)
              entries.append(value)
            except KeyError:
              logging.debug('Failed to process field %s', f.data)
        return entries

    def flatten_payload(self) -> list[Any]:
        """Flattens values in payload fields into a list."""
        processed_fields = self.convert_payload()
        return list(itertools.chain(*processed_fields))

    def as_log_str(self) -> str:
        t = struct.pack('>d', self.recv_time)
        return t + self.as_bytes

def split_payload_to_fields(payload: bytes) -> list[MipsField]:
    """Takes payload bytes and divides them into fields."""
    index = 0
    last_byte = len(payload)
    field_data = []
    while index < last_byte:
        field_len = payload[index]
        field_desc = payload[index + 1]
        data = payload[index + 2:index + field_len]
        field_data.append(MipsField(field_len, field_desc, data_bytes=data))
        index += field_len
    return field_data

class ReplyFormats:
    """Container of messages and mapping to formats for unpacking.

    This class contains a number of helper classes introduced to make it more
    obvious what's in the data received from the device. It maps the format
    to the struct unpacking instructions, as well as additional helpful info
    like the units or info of the terms.

    This class is intended to operate on individual fields, and contains
    class methods to:
        - Unpack a field's byte data into an object
        - Retrieve units for the field
    """
    # The ACK/NACK message consists of a command echo and error code.
    # This comes back when the field description is 0xf1.
    Ack = collections.namedtuple('Ack', ['cmd_echo', 'error_code'])

    # The device information, note there are lots of spaces in these
    # fields.
    DevInfo = collections.namedtuple('DevInfo',[
        'fw_ver', 'model_name', 'model_number', 'serial_num', 'lot_num',
        'dev_opts'])

    DataFmt = collections.namedtuple('DataFmt',
                                     ['descriptor', 'rate_decimation'])

    # Each set of information stores the format, class, and units of the data.
    Vector = collections.namedtuple('Vector', ['x', 'y', 'z'])
    base_set = {
        0x81: ('>H16s16s16s16s16s', DevInfo),  # Device info.
        0x83: ('>I', int),  # Built-in self test.
        0x9b: ('>fff', Vector),  # Gyro bias.
        0xf1: ('>BB', Ack),  # (N)Ack message.
    }

    Quat = collections.namedtuple('Quat', ['q0', 'q1', 'q2', 'q3'])
    Euler = collections.namedtuple('Euler', ['roll', 'pitch', 'yaw'])
    Matrix = collections.namedtuple('Matrix', ['M11', 'M12', 'M13',
                                               'M21', 'M22', 'M23',
                                               'M31', 'M32', 'M33'])
    Timestamp = collections.namedtuple('GpsTimestamp',
                                       ['time_of_week', 'week_num', 'flags'])
    AmbPressure = collections.namedtuple('AmbPressure', ['pressure'])

    imu_set = {
        0x04: ('>fff', Vector),  # Accelerometer vector.
        0x05: ('>fff', Vector),  # Gyro vector.
        0x06: ('>fff', Vector),  # Magnetometer vector.
        0x07: ('>fff', Vector),  # Delta theta vector.
        0x08: ('>fff', Vector),  # Delta velocity vector.
        0x09: ('>fffffffff', Matrix),  # CF orientation matrix.
        0x0a: ('>ffff', Quat),  # Complimentary filter (CF) quaternion.
        0x0c: ('>fff', Euler),  # CF Euler Angles.
        0x10: ('>fff', Vector),  # CF stabilized mag vector.
        0x11: ('>fff', Vector),  # CF stabilized accel vector.
        0x12: ('>dHH', Timestamp),  # GPS correlation timestamp.
        0x17: ('>f', AmbPressure),  # Scaled ambient pressure.
    }

    imu_units = {
        0x04: 3*['Accel g'],
        0x05: 3*['Gyro rad/s'],
        0x06: 3*['Mag gauss'],
        0x07: 3*['Delta Theta rad'],
        0x08: 3*['Delta Vel g*s'],
        0x09: 9*['n/a'],
        0x0a: ['q0', 'q1', 'q2', 'q3'],
        0x0c: 3*['rad'],
        0x10: 3*['CF Mag gauss'],
        0x11: 3*['CF Accel g'],
        0x12: ['sec', 'week', 'n/a'],
        0x17: ['milliBar'],
    }

    GpsLLH = collections.namedtuple(
        'GpsLLH',['lat', 'lon', 'ellipsoid_ht', 'msl_ht',
                  'horizontal_acc', 'vertical_acc', 'flags'
    ])
    PosECEF = collections.namedtuple(
        'PosECEF', ['x', 'y', 'z', 'position_acc', 'flags'])
    VelNED = collections.namedtuple('VelNED', [
        'north', 'east', 'down', 'speed', 'ground_speed', 'heading',
        'speed_acc', 'heading_acc', 'flags'
    ])
    VelECEF = collections.namedtuple('VelECEF',[
        'x', 'y', 'z', 'velocity_acc', 'flags'
    ])
    Dilution = collections.namedtuple('Dilution', [
        'geometric_dop', 'position_dop', 'horizontal_dop', 'vertical_dop',
        'time_dop', 'northing_dop', 'easting_dop', 'flags'
    ])
    UTC = collections.namedtuple('UTC', [
        'year', 'month', 'day', 'hour', 'minute', 'second', 'msec', 'flags'
    ])
    GPSTime = collections.namedtuple('GPSTime', [
        'time_of_week', 'week_num', 'flags'
    ])
    ClockInfo = collections.namedtuple('ClockInfo', [
        'clock_bias', 'clock_drift', 'acc_est', 'flags'
    ])
    GPSFix = collections.namedtuple('GPSFix', [
        'fix_type', 'num_sol_SVs', 'fix_flags', 'valid_flags'
    ])
    SvInfo = collections.namedtuple('SvInfo', [
        'channel', 'sv_id', 'carrier_noise_ratio', 'az', 'el', 'sv_flags',
        'valid_flags'
    ])
    GpsHwStatus = collections.namedtuple('GpsHwStatus', [
        'sensor_state', 'antenna_state', 'antenna_power', 'flags'
    ])

    gps_set = {
        0x03: ('>ddddffH', GpsLLH),
        0x04: ('>dddfH', PosECEF),
        0x05: ('>ffffffffH', VelNED),
        0x06: ('>ffffH', VelECEF),
        0x07: ('>fffffffH', Dilution),
        0x08: ('>HBBBBBIH', UTC),
        0x09: ('>dHH', GPSTime),
        0x0a: ('>dddH', ClockInfo),
        0x0b: ('>BBHH', GPSFix),
        0x0c: ('>BBH2s2sHH', SvInfo),
        0x0d: ('>BBBH', GpsHwStatus)
    }

    gps_units = {
        0x03: ['deg', 'deg', 'm', 'm', 'm', 'm', 'n/a'],
        0x04: 4*['m'] + ['n/a'],
        0x05: 5*['m/s'] + ['deg', 'm/s', 'deg', 'n/a'],
        0x06: 4*['m/s'] + ['n/a'],
        0x07: 8*['n/a'],
        0x08: ['years', 'months', 'days', 'hrs', 'min', 'sec', 'ms', 'n/a'],
        0x09: ['sec', 'week', 'n/a'],
        0x0a: ['sec', 'sec/sec', 'sec', 'n/a'],
        0x0b: ['n/a', 'count'] + 2*['n/a'],
        0x0c: ['chan_num', 'sv_id_num', 'dBHz', 'deg', 'deg'] + 2*['n/a'],
        0x0d: 4*['n/a']
    }

    FilterState = collections.namedtuple(
        'FilterState', ['state', 'dynamics_mode', 'flags'])
    EkfLLH = collections.namedtuple(
        'EkfLLH',['lat', 'lon', 'ellipsoid_ht', 'flags'])
    EkfVelNED = collections.namedtuple(
        'EkfVelNED', ['north', 'east', 'down', 'flags'])
    EkfQuat = collections.namedtuple('EkfQuat', Quat._fields + ('flags',))
    EkfMatrix = collections.namedtuple('EkfMatrix', Matrix._fields + ('flags',))
    EkfEuler = collections.namedtuple('EkfEuler', Euler._fields + ('flags',))
    EkfVector = collections.namedtuple('EkfVector', Vector._fields + ('flags',))
    UcLLH = collections.namedtuple('UcLLH',
        [f'uc_{f}_1sig' for f in ['north', 'east', 'down']] + ['flags'])
    UcNED = collections.namedtuple('UcNED', UcLLH._fields)
    UcEuler = collections.namedtuple('UcLLH',
        [f'uc_{f}_1sig' for f in ['roll', 'pitch', 'yaw']] + ['flags'])
    UcVec = collections.namedtuple('UcVec',
        ['uc_x_1sig', 'uc_y_1sig', 'uc_z_1sig', 'flags'])
    Wgs84 = collections.namedtuple('Wgs84', ['grav_mag', 'flags'])
    UcQuat = collections.namedtuple('UcQuat',
        [f'uc_{f}_1sig' for f in ['q0', 'q1', 'q2', 'q3']] + ['flags'])
    EkfHeading = collections.namedtuple('EkfHeading',
        ['true_heading', 'heading_uc_1sig', 'source', 'flags'])
    AtmModel = collections.namedtuple('AtmModel',
        ['geometric_alt', 'geopotential_alt', 'temp', 'pressure', 'density',
         'flags'])
    MagModel = collections.namedtuple('MagModel',
        ['intensity_north', 'intensity_east', 'intensity_down', 'inclination',
        'declination', 'flags'])
    PressureAlt = collections.namedtuple('PressureAlt', ['pressure', 'flags'])

    ekf_set = {
        0x01: ('>dddH', EkfLLH),  # LLH position.
        0x02: ('>fffH', EkfVelNED),  # NED velocity.
        0x03: ('>ffffH', EkfQuat),  # Orientation quaternion.
        0x04: ('>fffffffffH', EkfMatrix),  # Orientation matrix.
        0x05: ('>fffH', EkfEuler),  # Orientation euler angles.
        0x06: ('>fffH', EkfVector),  # Gyro bias.
        0x07: ('>fffH', EkfVector),  # Accel bias.
        0x08: ('>fffH', UcLLH),  # LLH position uncertainty.
        0x09: ('>fffH', UcNED),  # NED velocity uncertainty.
        0x0a: ('>fffH', UcEuler),  # Attitude uncertainty euler angles.
        0x0b: ('>fffH', UcVec),  # Gyro bias uncertainty.
        0x0c: ('>fffH', UcVec),  # Accel bias uncertainty.
        0x0d: ('>fffH', EkfVector),  # Linear acceleration.
        0x0e: ('>fffH', EkfVector),  # Compensated angular rate.
        0x0f: ('>fH', Wgs84),  # WGS84 local gravity magnitude.
        0x10: ('>HHH', FilterState),  # Filter status.
        0x11: ('>dHH', GPSTime),  # GPS timestamp.
        0x12: ('>ffffH', UcQuat),  # Attitude uncertainty, quat elements.
        0x13: ('>fffH', EkfVector),  # Gravity vector.
        0x14: ('>ffHH', EkfHeading),  # Heading update source state.
        0x15: ('>fffffH', MagModel),  # Magnetic model solution.
        0x16: ('>fffH', EkfVector),  # Gyro scale factor.
        0x17: ('>fffH', EkfVector),  # Accel scale factor.
        0x18: ('>fffH', UcVec),  # Gyro scale factor uncertainty.
        0x19: ('>fffH', UcVec),  # Accel scale factor uncertainty.
        0x1c: ('>fffH', EkfVector),  # Compensated linear acceleration.
        0x20: ('>fffffH', AtmModel),  # Standard atmosphere model.
        0x21: ('>fH', PressureAlt),  # Pressure altitude.
        0x30: ('>fffH', EkfVector),  # GPS Antenna offset correction.
        0x31: ('>fffH', UcVec),  # GPS Antenna offset correction uncertainty.
    }

    ekf_units = {
        0x01: 2*['deg'] + ['m', 'n/a'],
        0x02: 3*['m/s'] + ['n/a'],
        0x03: 5*['n/a'],
        0x04: 10*['n/a'],
        0x05: 3*['rad'] + ['n/a'],
        0x06: 3*['rad/s'] + ['n/a'],
        0x07: 3*['m/s2'] + ['n/a'],
        0x08: 3*['m'] + ['n/a'],
        0x09: 3*['m/s'] + ['n/a'],
        0x0a: 3*['rad'] + ['n/a'],
        0x0b: 3*['rad/s'] + ['n/a'],
        0x0c: 3*['m/s2'] + ['n/a'],
        0x0d: 3*['m/s2'] + ['n/a'],
        0x0e: 3*['rad/s'] + ['n/a'],
        0x0f: ['m/s2', 'n/a'],
        0x10: 3*['n/a'],
        0x11: ['sec', 'week', 'n/a'],
        0x12: 5*['n/a'],
        0x13: 3*['m/s2'] + ['n/a'],
        0x14: 2*['rad'] + 2*['n/a'],
        0x15: 3*['gauss'] + 2*['rad'] + ['n/a'],
        0x16: 3*['Gyro SF %/100'] + ['n/a'],
        0x17: 3*['Acc SF %/100'] + ['n/a'],
        0x18: 3*['Gyro SF %/100'] + ['n/a'],
        0x19: 3*['Acc SF %/100'] + ['n/a'],
        0x1c: 3*['m/s2'] + ['n/a'],
        0x20: ['m', 'm', 'decC', 'mBar', 'kg/m3', 'n/a'],
        0x21: ['m', 'n/a'],
        0x30: 3*['m'] + ['n/a'],
        0x31: 3*['m'] + ['n/a'],
    }

    format_map = {
        0x01: base_set,  # Back commands for all devices.
        0x0c: base_set,  # To pick up ack/nack for 3DM commands.
        0x0d: base_set,  # To pick up ack/nack for EKF commands.
        DataMessages.IMU.msg_ind: imu_set,
        DataMessages.GPS.msg_ind: gps_set,
        DataMessages.EKF.msg_ind: ekf_set,
    }

    @classmethod
    def process_field(self, desc_set: int,
                      mips_field: MipsField) -> tuple[any,...]:
        """Returns processed bytes according to field description."""
        message_set = self.format_map[desc_set]
        fmt, obj = message_set[mips_field.field_desc]
        return obj(*struct.unpack(fmt, mips_field.data_bytes))

    @classmethod
    def get_field_units(
        self, desc_set: int, mips_field: MipsField) -> list[str]:
        """Returns units for MipsField provided."""
        if desc_set == DataMessages.IMU.msg_ind:
            units = self.imu_units
        elif desc_set == DataMessages.GPS.msg_ind:
            units = self.gps_units
        elif desc_set == DataMessages.EKF.msg_ind:
            units = self.ekf_units
        return units[mips_field.field_desc]

    @classmethod
    def decode(self, data: collections.namedtuple) -> list[str]:
        """Converts data message into human readable form."""
        decoder_funcs = {
            self.FilterState.__name__: self._decode_ekf_status,
            self.GPSFix.__name__: self._decode_gps_fix_information,
            self.GpsHwStatus.__name__: self._decode_gps_hw_status,
        }
        func = decoder_funcs[type(data).__name__]
        return func(data)

    @classmethod
    def _apply_decoder(self, decoder: dict[int, str], val: int) -> list[str]:
        """Returns human readable description of integer value."""
        messages = []
        for key, msg in decoder.items():
            if key & val:
                messages.append(msg)
        return messages

    @classmethod
    def _decode_ekf_status(self, status: FilterState) -> list[str]:
        """Translates EKF messages into human readable form."""
        run_decoder = {
            0x0001: 'IMU Unavailable',
            0x0002: 'GPS Unavailable',
            0x0008: 'Matrix Singularity in calculation',
            0x0010: 'Position Covariance High Warning',
            0x0020: 'Velocity Covariance High Warning',
            0x0040: 'Attitude Covariance High Warning',
            0x0080: 'NAN in Solution',
            0x0100: 'Gyro bias estimate high warning',
            0x0200: 'Accel bias estimate high warning',
            0x0400: 'Gyro scale factor estimate high warning',
            0x0800: 'Accel scale factor estimate high warning',
            0x2000: 'GPS Antenna Offset Correction estimate high warning'
            }
        init_decoder = {
            0x1000: 'Attitude not initialized',
            0x2000: 'Position & Velocity not initialized'
            }
        mode_decoder = {
            0x00: 'startup',
            0x01: 'init',
            0x02: 'run',
            0x03: 'error'
        }
        state_str = self._apply_decoder(mode_decoder, status.state)[0]
        logging.debug(f'EKF in {state_str} phase.')

        decoder = {0x01: init_decoder,
                   0x02: run_decoder,
                   0x03: run_decoder}[status.state]
        return self._apply_decoder(decoder, status.flags)

    @classmethod
    def _decode_gps_fix_information(self, status: GPSFix) -> list[str]:
        """Parses fields of GPS fix information."""
        valid_decoder = {
            0x0001: "Fix type valid",
            0x0002: "Number of SVs valid",
            0x0004: "Fix flags valid",
        }
        type_decoder = {
            0x00: '3D Fix',
            0x01: '2D Fix',
            0x02: 'Time only',
            0x03: 'None',
            0x04: 'Invalid',
        }
        flag_decoder = {
            0x0001: 'SBAS corrections used',
            0x0002: 'Differential (DGPS) corrections used'
        }
        messages = self._apply_decoder(type_decoder, status.fix_type)
        messages.extend(self._apply_decoder(flag_decoder, status.fix_flags))
        messages.extend(self._apply_decoder(valid_decoder, status.valid_flags))
        return messages

    @classmethod
    def _decode_gps_hw_status(self, status: GpsHwStatus) -> list[str]:
        """Parses fields of GPS HW Status information."""
        valid_decoder = {
            0x0001: 'Sensor state valid',
            0x0002: 'Antenna state valid',
            0x0004: 'Antenna power valid',
        }
        sensor_state = {
            0x00: 'Sensor off',
            0x01: 'Sensor on',
            0x02: 'Sensor state unknown',
        }
        antenna_state = {
            0x01: 'Antenna init',
            0x02: 'Antenna short',
            0x03: 'Antenna open',
            0x04: 'Antenna good',
            0x05: 'Antenna state unknown',
        }
        antenna_power = {
            0x00: 'Antenna off',
            0x01: 'Antenna on',
            0x02: 'Antenna power unknown',
        }
        messages = self._apply_decoder(sensor_state, status.sensor_state)
        messages.extend(self._apply_decoder(antenna_state, status.antenna_state))
        messages.extend(self._apply_decoder(antenna_power, status.antenna_power))
        messages.extend(self._apply_decoder(valid_decoder, status.flags))
        return messages

def fletcher_checksum(packet: bytes) -> bytes:
    """Calculates fletcher checksum for bytes in packet."""
    check_sum = [sum(packet), sum(itertools.accumulate(packet))]
    return bytearray([0xff & check_sum[0], 0xff & check_sum[1]])

def process_mips_packets(
        packets: list[MipsPacket]) -> list[collections.namedtuple]:
    """Converts fields of MipsPackets to objects."""
    return list(itertools.chain.from_iterable(
        [p.convert_payload() for p in packets]))

def divide_mips_packets(
        packets: list[MipsPacket]) -> tuple[MipsPacket, MipsPacket, MipsPacket]:
    """Separates MipsPackets into IMU, GPS, and EKF formats."""
    imu_data = []
    gps_data = []
    ekf_data = []
    for p in packets:
        if p.desc_set == DataMessages.IMU.msg_ind:
            imu_data.append(p)
        elif p.desc_set == DataMessages.GPS.msg_ind:
            gps_data.append(p)
        elif p.desc_set == DataMessages.EKF.msg_ind:
            ekf_data.append(p)
    return imu_data, gps_data, ekf_data


import io
import os
from typing import Any, Dict, List, Optional, Set

import numpy as np
from PIL import Image
from pythomata.impl.symbolic import PropositionalInterpretation

from multinav.envs import sapientino_defs
from multinav.envs.base import AbstractFluents
from multinav.envs.temporal_goals import SapientinoGoal, SapientinoOfficeGoal
from multinav.helpers.general import classproperty
from multinav.helpers.gym import (
    Action,
    MyDiscreteEnv,
    Probability,
    State,
    Transition,
    Transitions,
    from_discrete_env_to_graphviz,
    iter_space,
)
from multinav.wrappers.temprl import FlattenAutomataStates, MyTemporalGoalWrapper

class AbstractSapientino(MyDiscreteEnv):
    """Abstract Sapientino environment."""

    def __init__(self, nb_colors: int, failure_probability: float = 0.1):
        """
        Initialize the environment.

        :param nb_colors: the number of color to consider.
        """
        self._nb_colors = nb_colors
        self._failure_probability = failure_probability
        model = self._make_transitions()
        isd = np.zeros(self.nb_states)
        isd[self.initial_state] = 1.0
        super().__init__(self.nb_states, self.nb_actions, model, isd)

    @property
    def nb_colors(self) -> int:
        """Get the number of colors."""
        return self._nb_colors

    @property
    def fail_prob(self) -> Probability:
        """Get the failure probability."""
        return self._failure_probability

    @property
    def initial_state(self) -> int:
        """Get the initial state."""
        return 0

    @property
    def nb_states(self) -> int:
        """
        Get the number of states.

        That is:
        - one state for the corridor
        - one state for each color
        NOTE: Some parts of this module assume that the corridor, or "blank"
        state is number 0. This is also the same convention of gym-sapientino.

        :return: the number of states.
        """
        return self.nb_colors + 1

    @property
    def nb_actions(self) -> int:
        """
        Get the number of actions.

        It includes:
        - go to to each color (+nb_colors)
        - go back to the corridor (+1)
        - visit color (+1)

        :return: the number of actions.
        """
        return self.nb_colors + 1 + 1

    def _is_legal_color(self, color_id: int):
        """Check that it is a legal color."""
        assert 0 <= color_id < self.nb_colors, f"{color_id} is not a legal color."

    def state_from_color(self, color_id: int):
        """
        Get the state from the color.

        Sum +1 because 0 is the initial state.
        """
        self._is_legal_color(color_id)
        return color_id + 1

    def action_goto_color_from_color(self, color_id: int) -> Action:
        """Get the action "goto color" from the color id."""
        assert 0 <= color_id < self.nb_colors, f"{color_id} is not a legal color."
        return color_id + 2

    @classproperty
    def goto_corridor(cls) -> Action:  # pylint: disable=no-self-argument
        """Get the action "goto corridor"."""
        return 0

    @classproperty
    def visit_color(cls) -> Action:  # pylint: disable=no-self-argument
        """Get the action "visit_color"."""
        return 1

    def _make_transitions(self) -> Transitions:
        """
        Make the model.

        :return: the transitions.
        """
        model: Transitions = {}

        for color_id in range(self.nb_colors):
            color_state = self.state_from_color(color_id)

            # from the corridor, you can go to any color
            goto_color_action = self.action_goto_color_from_color(color_id)
            new_transition = (1.0 - self.fail_prob, color_state, 0.0, False)
            fail_transition = (self.fail_prob, self.initial_state, 0.0, False)
            model.setdefault(self.initial_state, {}).setdefault(
                goto_color_action, []
            ).extend([new_transition, fail_transition])

            # if you visit a color, you remain in the same state.
            # NOTE: these two transitions are equivalent to a single transition
            #   it may be useful later, though.
            new_transition = (1.0 - self.fail_prob, color_state, 0.0, False)
            fail_transition = (self.fail_prob, color_state, 0.0, False)
            model.setdefault(color_state, {}).setdefault(self.visit_color, []).extend(
                [new_transition, fail_transition]
            )

            # from any color, you can go back to the corridor.
            new_transition = (1.0 - self.fail_prob, self.initial_state, 0.0, False)
            fail_transition = (self.fail_prob, color_state, 0.0, False)
            model.setdefault(color_state, {}).setdefault(self.goto_corridor, []).extend(
                [new_transition, fail_transition]
            )

            # TODO: experiments in AbstractSapientinoOffice require the
            # possibility to travel between pairs of colors (out and in of a
            # room). Adapted graph for this purpose.
            related_color = color_id + 1 if color_id % 2 == 0 else color_id - 1
            related_state = self.state_from_color(related_color)
            new_transition = (1.0 - self.fail_prob, related_state, 0.0, False)
            fail_transition = (self.fail_prob, color_state, 0.0, False)
            model.setdefault(color_state, {}).setdefault(
                self.action_goto_color_from_color(related_color), []
            ).extend([new_transition, fail_transition])

        return model

    def _action_to_string(self, action: Action):
        """From action to string."""
        self._is_legal_action(action)
        if action == self.goto_corridor:
            return "goto_corridor"
        if action == self.visit_color:
            return "visit"
        return f"goto_{action}"

    def _state_to_string(self, state: State):
        """From state to string.

        All sapientino environments interpret IDs as the same color.
        """
        self._is_legal_state(state)
        if state == 0:
            return "corridor"
        else:
            return sapientino_defs.int2color[state + 1]

    def render(self, mode="human"):
        """Render the environment (only rgb_array mode)."""
        assert mode == "rgb_array"
        graph = from_discrete_env_to_graphviz(
            self, state2str=self._state_to_string, action2str=self._action_to_string
        )
        image_file = io.BytesIO(graph.pipe(format="png"))
        array = np.array(Image.open(image_file))
        return array

class AbstractSapientinoTemporalGoal(MyDiscreteEnv):
    """
    Abstract Sapientino with Temporal Goals.

    We need this instead of using the temporal goal wrapper because we
    need to build an explicit model.
    """

    def __init__(
        self,
        *,
        tg_reward: float = 1.0,
        save_to: Optional[str] = None,
        **sapientino_kwargs,
    ):
        """Initialize the environment.

        :param tg_reward: reward supplied when the temporal goal is reached.
        :param save_to: path where the automaton temporal goal should be saved.
        """
        # Make AbstractSapientino
        self.sapientino_env = AbstractSapientino(**sapientino_kwargs)

        # Make a temporal goal
        nb_colors = sapientino_kwargs["nb_colors"]
        first_color = 2
        color_sequence = [
            sapientino_defs.int2color[i] for i in range(first_color, nb_colors + first_color)]
        self.fluents = Fluents(nb_colors=nb_colors)
        self.temporal_goal = SapientinoGoal(
            colors=color_sequence,
            fluents=self.fluents,
            reward=tg_reward,
            save_to=save_to,
        )

        # Build env with temporal goal
        self.temporal_env = FlattenAutomataStates(
            MyTemporalGoalWrapper(self.sapientino_env, [self.temporal_goal])
        )

        # compute model
        model: Transitions = {}
        initial_state = self.temporal_env.reset()  # because deterministic
        self._generate_transitions(model, initial_state)

        # Set discrete env
        nb_states = self.temporal_env.observation_space.nvec.prod()
        nb_actions = self.sapientino_env.nb_actions
        isd = np.zeros(nb_states)
        isd[0] = 1.0
        MyDiscreteEnv.__init__(self, nS=nb_states, nA=nb_actions, P=model, isd=isd)

        # Update observation space
        self.observation_space = self.temporal_env.observation_space

    def _generate_transitions(
        self,
        model: Transitions,
        state: State,
    ):
        """Recusively visit states and generate transitions."""
        if state in model:
            return

        # Transition function from state
        automaton = self.temporal_goal.automaton
        sapientino_tf = self.sapientino_env.P[state[0]]

        # For all actions
        model[state] = {}
        for action in sapientino_tf:
            new_transitions: List[Transition] = []
            model[state][action] = new_transitions

            # For all nondeterministic transitions
            for transition in sapientino_tf[action]:
                p, sap_state, sap_reward, _sap_done = transition

                # Possible evaluations
                fluents_list = self.fluents.evaluations_prob(sap_state, action)
                for fluents, fluents_p in fluents_list:

                    automaton_state = automaton.get_successor(state[1], fluents)

                    # Compose state
                    goal_reached = automaton_state in automaton.accepting_states
                    new_reward = sap_reward + (
                        0.0 if goal_reached else -1.0
                    )
                    new_done = goal_reached
                    new_state = (sap_state, automaton_state)

                    new_transitions.append(
                        (p * fluents_p, new_state, new_reward, new_done))

                    # Recurse
                    self._generate_transitions(model, new_state)

    def render(self, mode="human"):
        """Render with temporal goal; mode is ignored."""
        return self.temporal_env.render(mode="rgb_array")

    def reset(self):
        """Reset the environment."""
        return self.temporal_env.reset()

    def step(self, action):
        """Do a step on the environment."""
        return self.temporal_env.step(action)

class AbstractSapientinoOffice(AbstractSapientinoTemporalGoal):
    """AbstractSapientino with the office scenario temporal goal."""

    def __init__(
        self,
        *,
        nb_rooms: int,
        tg_reward: float,
        saved_automaton: str,
        save_to: Optional[str] = None,
        seed: int,
        **sapientino_kwargs,
    ):
        """Initialize the environment.

        :param nb_rooms: number of rooms to visit and get inside.
        :param tg_reward: reward supplied when the temporal goal is reached.
        :param saved_automaton: path to a saved DFA that represents the
            temporal goal. Transitions must be interpretations of fluents.
        :param save_to: path where the automaton temporal goal should be saved.
        """
        # Define temporal goal
        self.fluents = OfficeFluents(n_rooms=nb_rooms, seed=seed)
        self.temporal_goal = SapientinoOfficeGoal(
            n_rooms=nb_rooms,
            fluents=self.fluents,
            saved_automaton=saved_automaton,
            reward=tg_reward,
            save_to=save_to,
        )

        # Make AbstractSapientino
        sapientino_kwargs["nb_colors"] = nb_rooms * 2
        self.sapientino_env = AbstractSapientino(**sapientino_kwargs)

        # Build env with temporal goal
        self.temporal_env = FlattenAutomataStates(
            MyTemporalGoalWrapper(self.sapientino_env, [self.temporal_goal])
        )

        # compute model
        model: Transitions = {}
        initial_state = self.temporal_env.reset()  # because deterministic
        self._generate_transitions(model, initial_state)

        # Set discrete env
        nb_states = len(model)
        nb_actions = self.sapientino_env.nb_actions
        isd = np.zeros(nb_states)
        isd[0] = 1.0
        MyDiscreteEnv.__init__(self, nS=nb_states, nA=nb_actions, P=model, isd=isd)

        # Update observation space
        self.observation_space = self.temporal_env.observation_space

class Fluents(AbstractFluents):
    """Define the propositions for `AbstractSapientino`."""

    def __init__(self, nb_colors: int):
        """Initialize.

        :param nb_colors: The number of colors/rooms in the environment.
        """
        base_id = 2  # the firsts are for corridor and wall
        self.fluents = {
            sapientino_defs.int2color[i] for i in range(base_id, base_id + nb_colors)
        }

    def evaluate(self, obs: int, action: int) -> PropositionalInterpretation:
        """Respects AbstractFluents.evaluate.

        :param obs: assuming that the observation comes from an
            `AbstractSapientino` environment.
        :param action: the last action.
        """
        if action == AbstractSapientino.visit_color:
            fluents = {sapientino_defs.int2color[obs]}
            if obs == 0:  # blank/corridor
                fluents = set()
        else:
            fluents = set()
        return {f: f in fluents for f in self.fluents}

class OfficeFluents(AbstractFluents):
    """Define propositions for AbstractSapientino with Office goal."""

    _colors_to_room = {
        "red": "at1",
        "green": "in1",
        "blue": "at2",
        "yellow": "in2",
        "pink": "at3",
        "brown": "in3",
        "gray": "at4",
        "orange": "in4",
    }

    def __init__(self, n_rooms: int, seed: int):
        """Initialize.

        :param nb_rooms: number of rooms to navigate.
        """
        assert n_rooms < 5, "Can't support more than four rooms"

        # Define the propositional symbols
        self.fluents = {"bip", "person", "closed"}
        for i in range(1, n_rooms + 1):
            at_room, in_room = f"at{i}", f"in{i}"
            self.fluents.add(at_room)
            self.fluents.add(in_room)

        self._rng = np.random.default_rng(seed)
        self._n_rooms = n_rooms

    def evaluate(self, obs: int, action: int) -> Set[str]:
        """Respects.AbstractFluents.evaluate.

        :param obs: assuming that the observation comes from an
            `AbstractSapientino` environment.
        :param action: the last action.
        """
        fluents = set()
        if obs != 0:
            color = sapientino_defs.int2color[obs + 1]
            fluents.add(self._colors_to_room[color])
        if action == AbstractSapientino.visit_color:
            fluents.add("bip")

        # Doors and people are not in the observation. Valuate at random
        samples = self._rng.integers(0, 2, size=2)
        if samples[0] == 1:
            fluents.add("closed")
        if samples[1] == 1:
            fluents.add("person")

        return {f: f in fluents for f in self.fluents}

    def evaluations_prob(self, obs, action):
        fluents = self.evaluate(obs, action)
        fluents["closed"] = False
        fluents["person"] = False

        p = 1.0 / 4
        values = [(dict(fluents), p) for i in range(4)]
        values[1][0].update({"closed": True}), p
        values[2][0].update({"person": True}), p
        values[3][0].update({"closed": True, "person": True}), p
        return values

def make(params: Dict[str, Any], log_dir: Optional[str] = None):
    """Make the sapientino abstract state environment (agent teleports).

    :param params: a dictionary of parameters; see in this function the
        only ones that are used.
    :param log_dir: directory where logs can be saved.
    :return: an object that respects the gym.Env interface.
    """
    # Check
    if params["shaping"]:
        raise ValueError("Can't shape rewards in the most abstract environment.")

    # Build
    env = AbstractSapientinoOffice(
        nb_rooms=params["nb_rooms"],
        tg_reward=params["tg_reward"],
        saved_automaton=params["tg_automaton"],
        save_to=os.path.join(log_dir, "reward-dfa.dot") if log_dir else None,
        failure_probability=params["sapientino_fail_p"],
        seed=params["seed"],
    )
    return env


import os
from pathlib import Path

import yaml
from flask import request, Flask, render_template
from flask_cors import CORS
from flask_sqlalchemy import SQLAlchemy

from server import MAIN_FOLDER
from server.db_structures import UserLogs

with (MAIN_FOLDER / Path("app_config.yaml")).open("r") as stream:
    try:
        config = yaml.safe_load(stream)
    except yaml.YAMLError as exc:
        exit()

basedir = os.path.abspath(os.path.dirname(__file__))
db_name = config['database_rel_path']
host = config['host']
port = config['port']

app = Flask(__name__)
CORS(app)

app.config['SQLALCHEMY_DATABASE_URI'] = f"sqlite:///{os.path.join(basedir, db_name)}"
app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = True
db = SQLAlchemy(app)

@app.route('/db')
def render_logs_database():
    data = db.session.query(UserLogs).all()
    rows = [d.as_dict() for d in data]
    columns = (
        "ip_address", "time", "session_id", "kernel_id",
        "notebook_name", "cell_index", "cell_num",
        "event", "cell_source", "cell_output"
    )

    return render_template('db_template.html', rows=rows, columns=columns)

@app.route('/', methods=['GET', 'POST'])
def add_message():
    content = request.json
    content['cell_source'] = str(content['cell_source']) if content['cell_source'] else None
    content['cell_output'] = str(content['cell_output']) if content['cell_output'] else None

    user_log = UserLogs(**content)
    db.session.add(user_log)
    db.session.commit()
    return content

def run_server() -> None:
    app.run(host=host, port=port)

if __name__ == "__main__":
    run_server()


from . import individual
import numpy as np
import random

def create_population(nb_population: int, hidden_layers, world, x, y, rotation):
    """
    Crea una población para el algoritmo genético.

    :param hidden_layers: Forma de la red
    :param nb_population: Cantidad de individuos en la población
    :return: Arreglo de numpy con la población creada
    """
    population = np.ndarray((nb_population,), dtype=np.object)
    for i in range(nb_population):
        population[i] = individual.Individual(world, x, y, rotation, hidden_layers)

    return population

class GeneticAlgorithm:

    def __init__(self, initial_population, lambda_, mu, selection,
                 mutation, crossover, world, hidden_layers, x, y, rotation, fitness=np.inf, generations=50,
                 instances=None):
        self.initial_population = initial_population
        self.population = create_population(initial_population, hidden_layers, world, x, y, rotation)
        self.problem = world
        self.lambda_ = lambda_
        self.mu = mu
        self.mutation_rate = mutation
        self.best_agent = None
        self.best_fitness = -1 * np.inf
        self.generations = 0
        self.max_fitness = fitness
        self.max_gen = generations
        self.instances = instances if instances is not None else None

    def execute(self):
        print("1 First generation")
        while not self.stop_condition():
            self.evaluate()
            self.population = sorted(self.population,
                                     key=lambda x: x.fitness,
                                     reverse=True)

            self.best_agent = self.population[0]
            self.best_fitness = self.population[0].fitness

            offsprings = self.crossover()
            self.mutation(offsprings)

            self.selection(offsprings)

            self.generations += 1
            print(f'{self.generations+1} '
                  f'best current fitness: {self.best_fitness} '
                  f'{self.population[1].fitness} '
                  f'{self.population[2].fitness}\t '
                  f'size: {len(self.population)}')

        print(f'best agent fitness: {self.best_agent.fitness}')
        self.best_agent.save()
        return self.best_agent

    def pick_weights(self):
        self.population[0].load()

    def parent_selection(self):
        """
        Realiza la seleccion de padres de la nueva generación

        :return:
        """
        # Ordenando entidades de acuerdo a su fitness

        self.population = sorted(self.population,
                                 key=lambda x: x.fitness,
                                 reverse=True)
        # Selección de elementos padre de la siguiente generación
        fitness_list = np.array([agent.fitness for agent in self.population])

        first_parent = self.__pick_parent(random.random(),
                                          fitness_list/sum(fitness_list))
        np.delete(fitness_list, first_parent)
        second_parent = self.__pick_parent(random.random(),
                                           fitness_list/sum(fitness_list))

        return self.population[first_parent], self.population[second_parent]

    def __pick_parent(self, probability, candidates):
        """
        Selecciona un padre de acuerdo a un valor aleatorio
        :param probability: valor aleatorio entregado
        :param candidates: lista de padres candidatos
        :return: indice del padre candidato
        """
        counter = 0

        while True:
            if sum(candidates[0: counter+1]) >= probability:
                break
            counter += 1
        return counter

    def crossover(self):
        """
        Lleva a cbao la combinación genética de los padres para crear una nueva
        entidad.

        :return:
        """
        offsprings = np.empty([self.lambda_,
                               individual.Individual.size_flattened],
                              dtype=np.object)

        for i in range(0, self.lambda_, 2):
            first_parent, second_parent = self.parent_selection()
            split_point = random.randint(0,
                                         individual.Individual.size_flattened-1)

            genes_1 = first_parent.get_genes()
            genes_2 = second_parent.get_genes()

            offsprings[i] = np.concatenate((genes_1[0:split_point+1],
                                           genes_2[split_point:-1]))
            offsprings[i+1] = np.concatenate((genes_2[0:split_point+1],
                                             genes_1[split_point:-1]))

        return offsprings

    def mutation(self, offsprings):
        """
        Realiza modificaciones sobre un individuo para mejorar la
        diversificación de la búsqueda.

        :return:
        """
        for child in offsprings:
            for i in range(random.randint(1, 10)):
                random_value = random.uniform(-1, 1)
                random_choice = random.randint(0,
                                               individual.Individual.size_flattened-1)

                if random.random() <= self.mutation_rate:
                    child[random_choice] = random_value

    def evaluate(self):
        # Calcular el fitness de cada individuo
        # Propuesta: ejecutar solamente los nuevos individuos creados
        for car in self.population:
            car.calculate_fitness(self.instances)
            print('.', end='')

        print('\n')

    def selection(self, offsprings):
        """
        Selecciona los individuos que sobrevivirán para la próxima iteración.

        :return:
        """
        if self.lambda_ < self.mu:
            for i in range(self.lambda_):
                agent = self.population[-1*(i+1)]
                agent.overwrite(offsprings[i])
        elif self.lambda_ > self.mu:
            pass
        else:
            raise ValueError("Lambda no puede ser igual a mu")

    def stop_condition(self):
        """
        Evalúa la condición de termino asignada para finalizar la búsqueda.
        :return:
        """
        if self.best_fitness > self.max_fitness or self.generations > self.max_gen:
            return True
        return False


import arcade

# Constants
SCREEN_WIDTH = 1000
SCREEN_HEIGHT = 550
SCREEN_TITLE = "Platformer"

# Constants used to scale our sprites from their original size
CHARACTER_SCALING = 1
TILE_SCALING = 0.5
COIN_SCALING = 0.5
SPRITE_PIXEL_SIZE = 128
GRID_PIXEL_SIZE = (SPRITE_PIXEL_SIZE * TILE_SCALING)

# How many pixels to keep as a minimum margin between the character
# and the edge of the screen.
LEFT_VIEWPORT_MARGIN = 200
RIGHT_VIEWPORT_MARGIN = 200
BOTTOM_VIEWPORT_MARGIN = 150
TOP_VIEWPORT_MARGIN = 100

# Movement speed of player, in pixels per frame
PLAYER_MOVEMENT_SPEED = 5
GRAVITY = 1
PLAYER_JUMP_SPEED = 20

PLAYER_START_X = 128
PLAYER_START_Y = 300

# number of levels
LEVELS = 3

class MyGame(arcade.Window):
    """
    Main application class.
    """

    def __init__(self):

        # Call the parent class and set up the window
        super().__init__(SCREEN_WIDTH, SCREEN_HEIGHT, SCREEN_TITLE)

        # These are 'lists' that keep track of our sprites. Each sprite should
        # go into a list.
        self.coin_list = None
        self.wall_list = None
        self.foreground_list = None
        self.background_list = None
        self.dont_touch_list = None
        self.player_list = None
        self.enemy_list = None

        # Separate variable that holds the player sprite
        self.player_sprite = None

        # Our physics engine
        self.physics_engine = None

        # Used to keep track of our scrolling
        self.view_bottom = 0
        self.view_left = 0

        # Keep track of the score
        self.score = 0

        # Where is the right edge of the map?
        self.end_of_map = 0

        # Level
        self.level = 1

        # Game state
        self.game_over = False

        # Load sounds
        self.collect_coin_sound = arcade.load_sound("sounds/coin1.wav")
        self.jump_sound = arcade.load_sound("sounds/jump1.wav")
        self.gameover = arcade.load_sound("sounds/gameover4.wav")
        self.enemy_hit = arcade.load_sound("sounds/hurt3.wav")
        self.fall = arcade.load_sound("sounds/fall1.wav")
        self.dont_touch = arcade.load_sound("sounds/hit1.wav")

    def setup(self, level):
        # Used to keep track of our scrolling
        self.view_bottom = 0
        self.view_left = 0

        # Keep track of the score
        self.score = 0

        # Create the Sprite lists
        self.player_list = arcade.SpriteList()
        self.foreground_list = arcade.SpriteList()
        self.background_list = arcade.SpriteList()
        self.wall_list = arcade.SpriteList()
        self.coin_list = arcade.SpriteList()
        self.enemy_list = arcade.SpriteList()

        # Set up the player, specifically placing it at these coordinates.
        self.player_sprite = arcade.Sprite("images/player_1/player_stand.png", CHARACTER_SCALING)
        self.player_sprite.center_x = PLAYER_START_X
        self.player_sprite.center_y = PLAYER_START_Y
        self.player_list.append(self.player_sprite)

        # --- Load in a map from the tiled editor ---

        # Name of the layer in the file that has our platforms/walls
        platforms_layer_name = 'Platforms'
        # Name of the layer that has items for pick-up
        coins_layer_name = 'Coins'
        # Name of the layer that has items for foreground
        foreground_layer_name = 'Foreground'
        # Name of the layer that has items for background
        background_layer_name = 'Background'
        # Name of the layer that has items we shouldn't touch
        dont_touch_layer_name = "Don't Touch"

        # Map name
        map_name = f"Level{level}.tmx"

        # Read in the tiled map
        my_map = arcade.read_tiled_map(map_name, TILE_SCALING)

        # Grab the layer of items we can't move through
        map_array = my_map.layers_int_data[platforms_layer_name]

        # Calculate the right edge of the my_map in pixels
        self.end_of_map = len(map_array[0]) * GRID_PIXEL_SIZE

        # -- Background
        self.background_list = arcade.generate_sprites(my_map, background_layer_name, TILE_SCALING)

        # -- Foreground
        self.foreground_list = arcade.generate_sprites(my_map, foreground_layer_name, TILE_SCALING)

        # -- Platforms
        self.wall_list = arcade.generate_sprites(my_map, platforms_layer_name, TILE_SCALING)

        # -- Coins
        self.coin_list = arcade.generate_sprites(my_map, coins_layer_name, TILE_SCALING)

        # -- Don't Touch Layer
        self.dont_touch_list = arcade.generate_sprites(my_map, dont_touch_layer_name, TILE_SCALING)

        # this line needed to check when player gets to the end of map
        self.end_of_map = (len(map_array[0]) - 1) * GRID_PIXEL_SIZE

        # Set the background color from the map
        if my_map.backgroundcolor:
            arcade.set_background_color(my_map.backgroundcolor)

        # create enemies
        if self.level == 1:
            enemy = arcade.Sprite("images/enemies/slimePurple.png", TILE_SCALING)
            enemy.bottom = GRID_PIXEL_SIZE * 4
            enemy.left = GRID_PIXEL_SIZE * 5
            # right and left boundary limits for the enemy
            enemy.boundary_right = GRID_PIXEL_SIZE * 11
            enemy.boundary_left = GRID_PIXEL_SIZE * 1
            # set enemy speed
            enemy.change_x = 2
            self.enemy_list.append(enemy)

            # create enemies
            enemy = arcade.Sprite("images/enemies/wormGreen.png", TILE_SCALING)
            enemy.bottom = GRID_PIXEL_SIZE * 4
            enemy.left = GRID_PIXEL_SIZE * 25
            # right and left boundary limits for the enemy
            enemy.boundary_right = GRID_PIXEL_SIZE * 31
            enemy.boundary_left = GRID_PIXEL_SIZE * 21
            # set enemy speed
            enemy.change_x = 2
            self.enemy_list.append(enemy)

            # create enemies
            enemy = arcade.Sprite("images/enemies/bee.png", TILE_SCALING)
            enemy.bottom = GRID_PIXEL_SIZE * 7
            enemy.left = GRID_PIXEL_SIZE * 16
            # set enemy speed
            enemy.change_x = -4
            self.enemy_list.append(enemy)

        if self.level == 2:
            enemy = arcade.Sprite("images/enemies/slimeGreen.png", TILE_SCALING)
            enemy.bottom = GRID_PIXEL_SIZE * 3
            enemy.left = GRID_PIXEL_SIZE * 5
            # right and left boundary limits for the enemy
            enemy.boundary_right = GRID_PIXEL_SIZE * 6
            enemy.boundary_left = GRID_PIXEL_SIZE
            # set enemy speed
            enemy.change_x = 2
            self.enemy_list.append(enemy)

            # create enemies
            enemy = arcade.Sprite("images/enemies/saw.png", TILE_SCALING)
            enemy.bottom = GRID_PIXEL_SIZE * 6
            enemy.left = GRID_PIXEL_SIZE * 19
            # set enemy speed
            enemy.change_x = -3
            self.enemy_list.append(enemy)

        if self.level == 3:
            enemy = arcade.Sprite("images/enemies/fishGreen.png", TILE_SCALING)
            enemy.bottom = GRID_PIXEL_SIZE * 4
            enemy.left = GRID_PIXEL_SIZE * 5
            # set enemy speed
            enemy.change_x = -2
            self.enemy_list.append(enemy)

            # create enemies
            enemy = arcade.Sprite("images/enemies/fishPink.png", TILE_SCALING)
            enemy.bottom = GRID_PIXEL_SIZE * 4
            enemy.left = GRID_PIXEL_SIZE * 25
            # right and left boundary limits for the enemy
            enemy.boundary_right = GRID_PIXEL_SIZE * 31
            enemy.boundary_left = GRID_PIXEL_SIZE * 21
            # set enemy speeddw
            enemy.change_x = 2
            #enemy.change_y = 2
            self.enemy_list.append(enemy)

        # Create the 'physics engine'
        self.physics_engine = arcade.PhysicsEnginePlatformer(self.player_sprite,
                                                             self.wall_list,
                                                             GRAVITY)
        # less gravity map
        if self.level == 3:
            self.physics_engine = arcade.PhysicsEnginePlatformer(self.player_sprite,
                                                                 self.wall_list,
                                                                 0.1)

    def draw_game_over(self):
        output1 = "Game Over"
        output2 = "Click to restart game"
        output3 = "Congratulations!"
        output4 = "Click to start again"

        if self.level == LEVELS and self.player_sprite.right >= self.end_of_map:
            arcade.draw_text(output3, 300 + self.view_left, 300 + self.view_bottom, arcade.color.GREEN, 54)
            arcade.draw_text(output4, 320 + self.view_left, 250 + self.view_bottom, arcade.color.RED, 25)
        else:
            arcade.draw_text(output1, 300 + self.view_left, 300 + self.view_bottom, arcade.color.WHITE, 54)
            arcade.draw_text(output2, 320 + self.view_left, 250 + self.view_bottom, arcade.color.RED, 25)

    def on_mouse_press(self, x, y, button, modifiers):
        # move to the first level after successful completion of all levels
        if self.game_over is True and self.level == 3 and self.player_sprite.right >= self.end_of_map:
            self.level = 1
            self.setup(self.level)
            self.game_over = False
        # restart the existing level
        if self.game_over is True:
            self.setup(self.level)
            self.game_over = False

    def draw_game(self):
        """ Render the screen. """

        arcade.start_render()

        # Draw our sprites
        self.wall_list.draw()
        self.background_list.draw()
        self.coin_list.draw()
        self.dont_touch_list.draw()
        self.player_list.draw()
        self.enemy_list.draw()
        self.foreground_list.draw()

        # Draw our score on the screen, scrolling it with the viewport
        score_text = f"Score: {self.score}"
        arcade.draw_text(score_text, 10 + self.view_left, 10 + self.view_bottom,
                         arcade.csscolor.BLACK, 18)

    def on_draw(self):
        # render the screen
        arcade.start_render()

        if not self.game_over:
            self.draw_game()

        else:
            self.draw_game()
            self.draw_game_over()

    def on_key_press(self, key, modifiers):
        """Called whenever a key is pressed. """

        if key == arcade.key.UP or key == arcade.key.W:
            if self.physics_engine.can_jump() and self.level != 3:
                self.player_sprite.change_y = PLAYER_JUMP_SPEED
                arcade.play_sound(self.jump_sound)
            else:
                self.player_sprite.change_y = PLAYER_JUMP_SPEED / 5
                arcade.play_sound(self.jump_sound)
        elif key == arcade.key.LEFT or key == arcade.key.A:
            self.player_sprite.change_x = -PLAYER_MOVEMENT_SPEED
        elif key == arcade.key.RIGHT or key == arcade.key.D:
            self.player_sprite.change_x = PLAYER_MOVEMENT_SPEED

    def on_key_release(self, key, modifiers):
        """Called when the user releases a key. """

        if key == arcade.key.LEFT or key == arcade.key.A:
            self.player_sprite.change_x = 0
        elif key == arcade.key.RIGHT or key == arcade.key.D:
            self.player_sprite.change_x = 0

    def update(self, delta_time):

        # game logic
        if not self.game_over:

            self.physics_engine.update()

            # move enemies
            self.enemy_list.update()

            # Track if we need to change the viewport
            changed_viewport = False

            # See if we hit any coins
            coin_hit_list = arcade.check_for_collision_with_list(self.player_sprite,
                                                                 self.coin_list)
            # Loop through each coin we hit (if any) and remove it
            for coin in coin_hit_list:
                # Remove the coin
                coin.remove_from_sprite_lists()
                # Play a sound
                arcade.play_sound(self.collect_coin_sound)
                # Add one to the score
                self.score += 1

            if self.level != 3:
                for enemy in self.enemy_list:
                    # check if enemy hit a wall, if yes then reverse
                    if len(arcade.check_for_collision_with_list(enemy, self.wall_list)) > 0:
                        enemy.change_x *= -1
                    # check if enemy hit left boundary, reverse
                    elif enemy.boundary_left is not None and enemy.left < enemy.boundary_left:
                        enemy.change_x *= -1
                    # check for right boundary
                    elif enemy.boundary_right is not None and enemy.right > enemy.boundary_right:
                        enemy.change_x *= -1
            else:
                for enemy in self.enemy_list:
                    # check if enemy hit a wall, if yes then reverse
                    if len(arcade.check_for_collision_with_list(enemy, self.dont_touch_list)) > 0:
                        enemy.change_x *= -1
                    # check if enemy hit left boundary, reverse
                    elif enemy.boundary_left is not None and enemy.left < enemy.boundary_left:
                        enemy.change_x *= -1
                    # check for right boundary
                    elif enemy.boundary_right is not None and enemy.right > enemy.boundary_right:
                        enemy.change_x *= -1

            # check if player touched a enemy
            if len(arcade.check_for_collision_with_list(self.player_sprite, self.enemy_list)) > 0:
                # Set the camera to the start
                self.game_over = True
                self.view_left = 0
                self.view_bottom = 0
                arcade.play_sound(self.enemy_hit)

            # Did the player fall off the map?
            if self.player_sprite.center_y < -100 and self.level != 3:
                self.player_sprite.center_x = PLAYER_START_X
                self.player_sprite.center_y = PLAYER_START_Y

                # Set the camera to the start
                self.view_left = 0
                self.view_bottom = 0
                changed_viewport = True
                arcade.play_sound(self.fall)

            # Did the player touch something they should not?
            if arcade.check_for_collision_with_list(self.player_sprite, self.dont_touch_list):
                self.player_sprite.center_x = PLAYER_START_X
                self.player_sprite.center_y = PLAYER_START_Y

                # Set the camera to the start
                self.view_left = 0
                self.view_bottom = 0
                changed_viewport = True
                arcade.play_sound(self.dont_touch)

            # # See if the user got to the end of the level and end of game
            if self.player_sprite.right >= self.end_of_map and self.level < LEVELS:
                # if end of current map, change map
                self.level += 1
                self.setup(self.level)

                # Set the camera to the start
                self.view_left = 0
                self.view_bottom = 0
                changed_viewport = True
            elif self.player_sprite.right >= self.end_of_map and self.level == LEVELS:
                self.game_over = True
                changed_viewport = False
                arcade.play_sound(self.gameover)

            # --- Manage Scrolling ---

            # Scroll left

            left_boundary = self.view_left + LEFT_VIEWPORT_MARGIN
            if self.player_sprite.left < left_boundary:
                self.view_left -= left_boundary - self.player_sprite.left
                changed_viewport = True

            right_boundary = self.view_left + SCREEN_WIDTH - RIGHT_VIEWPORT_MARGIN
            if self.player_sprite.right > right_boundary:
                self.view_left += self.player_sprite.right - right_boundary
                changed_viewport = True

            # Scroll up
            top_boundary = self.view_bottom + SCREEN_HEIGHT - TOP_VIEWPORT_MARGIN
            if self.player_sprite.top > top_boundary:
                self.view_bottom += self.player_sprite.top - top_boundary
                changed_viewport = True

            # Scroll down
            bottom_boundary = self.view_bottom + BOTTOM_VIEWPORT_MARGIN
            if self.player_sprite.bottom < bottom_boundary:
                self.view_bottom -= bottom_boundary - self.player_sprite.bottom
                changed_viewport = True

            if changed_viewport:
                # Only scroll to integers. Otherwise we end up with pixels that
                # don't line up on the screen
                self.view_bottom = int(self.view_bottom)
                self.view_left = int(self.view_left)

            # Do the scrolling
                arcade.set_viewport(self.view_left, SCREEN_WIDTH + self.view_left,
                                self.view_bottom, SCREEN_HEIGHT + self.view_bottom)

def main():
    """ Main method """

    window = MyGame()
    window.setup(window.level)
    arcade.run()

if __name__ == "__main__":
    main()


from csv_parser import Parser
import gmplot

a = Parser('transport_data.csv')
a.open()
dots = a.get_data()

gmap = gmplot.GoogleMapPlotter(59.95, 30.3, 12)

zero = []
first = []
second = []
q = []
for dot in dots:
    if dot.label == '0':
        zero.append(dot)
    if dot.label == '1':
        first.append(dot)
    if dot.label == '2':
        second.append(dot)
    if dot.label == '?':
        q.append(dot)

lats, logs = [], []
for dot in zero:
    lats.append(dot.lat)
    logs.append(dot.log)

gmap.scatter(lats, logs, 'red', marker=False)

lats, logs = [], []
for dot in first:
    lats.append(dot.lat)
    logs.append(dot.log)

gmap.scatter(lats, logs, 'pink', marker=False)

lats, logs = [], []
for dot in second:
    lats.append(dot.lat)
    logs.append(dot.log)

gmap.scatter(lats, logs, 'cornflowerblue', marker=False)

lats, logs = [], []
for dot in q:
    lats.append(dot.lat)
    logs.append(dot.log)

gmap.scatter(lats, logs, 'black', marker=False)

gmap.draw("mymap.html")

print(len(q))

"""
label 0 - маршрут 5
label 1 - маршрут 11
label 2 - маршрут 7
"""


import numpy as np

from pynucastro.reduction.reduction_utils import (mpi_importer,
                                                  mpi_numpy_decomp, to_list)

MPI = mpi_importer()

def calc_interaction_matrix(net, rvals):
    """Calculate direct interaction coefficients."""

    N_species = len(net.unique_nuclei)

    # create index mapping
    j_map = {}
    for idx, n in enumerate(net.unique_nuclei):
        j_map[n] = idx

    # Calculate coefficients for A
    p_A = np.zeros(N_species, dtype=np.float64)
    c_A = np.zeros(N_species, dtype=np.float64)

    # A along rows, B along columns
    p_AB = np.zeros((N_species, N_species))
    c_AB = np.zeros((N_species, N_species))

    for i in range(N_species):

        n = net.unique_nuclei[i]

        for r in net.nuclei_produced[n]:

            rval = r.products.count(n) * rvals[r]
            p_A[i] += rval

            bs = set(r.products) | set(r.reactants)
            for b in bs:
                p_AB[i, j_map[b]] += rval

        for r in net.nuclei_consumed[n]:

            rval = r.reactants.count(n) * rvals[r]
            c_A[i] += rval

            bs = set(r.products) | set(r.reactants)
            for b in bs:
                c_AB[i, j_map[b]] += rval

    denom = np.maximum(p_A, c_A)[:, np.newaxis]

    # Calculate direct interaction coefficients
    r_AB = np.abs(p_AB - c_AB) / denom

    return r_AB

def calc_interaction_matrix_numpy(net, rvals_arr):
    """Calculate direct interaction coefficients using NumPy."""

    # Evaluate terms on RHS of ODE system
    prod_terms = net.nuc_prod_count * rvals_arr
    cons_terms = net.nuc_cons_count * rvals_arr

    # Calculate total production and consumption of each nucleus A
    p_A = prod_terms.sum(axis=1)
    c_A = cons_terms.sum(axis=1)

    # Calculate production / consumption of A in reactions involving B
    p_AB = prod_terms @ net.nuc_used
    c_AB = cons_terms @ net.nuc_used

    # We will normalize by maximum of production and consumption fluxes
    denom = np.maximum(p_A, c_A)[:, np.newaxis]

    # Calculate direct interaction coefficients
    r_AB = np.abs(p_AB - c_AB) / denom

    return r_AB

def get_adj_nuc(net):
    """
    Get set of adjacent nuclei for each nucleus in the net. Returns dictionary keyed by
    nuclei.
    """

    adj_nuc = {}

    for n in net.unique_nuclei:

        bs = set()

        for r in net.nuclei_produced[n]:
            bs |= set(r.products) | set(r.reactants)
        for r in net.nuclei_consumed[n]:
            bs |= set(r.products) | set(r.reactants)

        adj_nuc[n] = bs

    return adj_nuc

def drgep_dijkstras(net, r_AB, target, adj_nuc):
    """
    Modified Dijkstra's algorithm to find paths that maximize the overall interaction coefficient
    between the target and each other nucleus in the net.
    """

    # Number of species
    nspec = len(net.unique_nuclei)

    # Create data structures
    j_map = {}

    inf = float('inf')

    dist = np.zeros(nspec, dtype=np.float64)
    R_TB = np.zeros(nspec, dtype=np.float64)

    for idx, n in enumerate(net.unique_nuclei):
        j_map[n] = idx
        dist[idx] = -inf
        R_TB[idx] = -inf

    dist[j_map[target]] = 1.0
    R_TB[j_map[target]] = inf

    imax = np.argmax(dist)

    # Main loop
    while dist[imax] > -inf:

        r_TA = dist[imax]
        R_TB[imax] = r_TA
        n = net.unique_nuclei[imax]
        bs = adj_nuc[n]

        for b in bs:
            jb = j_map[b]
            if R_TB[jb] > -inf:
                continue
            dist[jb] = max(dist[jb], r_TA * r_AB[imax, jb])

        dist[imax] = -inf
        imax = np.argmax(dist)

    return R_TB

def _drgep_kernel(net, R_TB, rvals, targets, tols, adj_nuc):

    r_AB = calc_interaction_matrix(net, rvals)

    for target, tol in zip(targets, tols):

        R_TB_i = drgep_dijkstras(net, r_AB, target, adj_nuc)
        np.maximum(R_TB, R_TB_i, out=R_TB, where=R_TB_i >= tol)

def _drgep_kernel_numpy(net, R_TB, rvals, targets, tols, adj_nuc):

    r_AB = calc_interaction_matrix_numpy(net, rvals)

    for target, tol in zip(targets, tols):

        R_TB_i = drgep_dijkstras(net, r_AB, target, adj_nuc)
        np.maximum(R_TB, R_TB_i, out=R_TB, where=R_TB_i >= tol)

def _drgep(net, conds, targets, tols):

    #-----------------------------------
    # Calculate interaction coefficients
    #-----------------------------------

    R_TB = np.zeros(len(net.unique_nuclei), dtype=np.float64)
    adj_nuc = get_adj_nuc(net)

    for comp, rho, T in conds:
        rvals = net.evaluate_rates(rho=rho, T=T, composition=comp)
        _drgep_kernel(net, R_TB, rvals, targets, tols, adj_nuc)

    return R_TB

def _drgep_mpi(net, conds, targets, tols):

    #----------
    # Init. MPI
    #----------

    comm = MPI.COMM_WORLD
    MPI_N = comm.Get_size()
    MPI_rank = comm.Get_rank()

    #-----------------------------------------------
    # Calculate interaction coefficients in parallel
    #-----------------------------------------------

    R_TB_loc = np.zeros(len(net.unique_nuclei), dtype=np.float64)
    adj_nuc = get_adj_nuc(net)

    for i in range(MPI_rank, len(conds), MPI_N):
        comp, rho, T = conds[i]
        rvals = net.evaluate_rates(rho=rho, T=T, composition=comp)
        _drgep_kernel(net, R_TB_loc, rvals, targets, tols, adj_nuc)

    R_TB = np.zeros_like(R_TB_loc)
    comm.Allreduce([R_TB_loc, MPI.DOUBLE], [R_TB, MPI.DOUBLE], op=MPI.MAX)

    return R_TB

def _drgep_numpy(net, conds, targets, tols):

    #----------------------------------------
    # Unpack conditions; make precalculations
    #----------------------------------------

    comp_L, rho_L, T_L = conds

    net.calc_count_matrices()
    net.update_rate_coef_arr()
    adj_nuc = get_adj_nuc(net)

    #------------------------------------------------
    # Calculate interaction coefficients (vectorized)
    #------------------------------------------------

    R_TB = np.zeros(len(net.unique_nuclei), dtype=np.float64)

    for comp in comp_L:
        net.update_yfac_arr(comp)
        for rho in rho_L:
            net.update_prefac_arr(rho, comp)
            for T in T_L:
                rvals_arr = net.evaluate_rates_arr(T)
                _drgep_kernel_numpy(net, R_TB, rvals_arr, targets, tols, adj_nuc)

    net.clear_arrays()
    return R_TB

def _drgep_mpi_numpy(net, conds, targets, tols):

    #----------------------------------------
    # Unpack conditions; make precalculations
    #----------------------------------------

    n = tuple(map(len, conds))
    comp_L, rho_L, T_L = conds

    net.calc_count_matrices()
    net.update_rate_coef_arr()
    adj_nuc = get_adj_nuc(net)

    #-----------------------------------
    # Init. MPI and divide up conditions
    #-----------------------------------

    comm = MPI.COMM_WORLD
    MPI_N = comm.Get_size()
    MPI_rank = comm.Get_rank()

    comp_idx, comp_step, rho_idx, rho_step, T_idx, T_step = mpi_numpy_decomp(MPI_N, MPI_rank, n)

    #--------------------------------------------------------------
    # Calculate interaction coefficients (vectorized and using MPI)
    #--------------------------------------------------------------

    R_TB_loc = np.zeros(len(net.unique_nuclei), dtype=np.float64)

    for i in range(comp_idx, n[0], comp_step):
        comp = comp_L[i]
        net.update_yfac_arr(comp)
        for j in range(rho_idx, n[1], rho_step):
            rho = rho_L[j]
            net.update_prefac_arr(rho, comp)
            for k in range(T_idx, n[2], T_step):
                T = T_L[k]
                rvals_arr = net.evaluate_rates_arr(T)
                _drgep_kernel_numpy(net, R_TB_loc, rvals_arr, targets, tols, adj_nuc)

    R_TB = np.zeros_like(R_TB_loc)
    comm.Allreduce([R_TB_loc, MPI.DOUBLE], [R_TB, MPI.DOUBLE], op=MPI.MAX)

    net.clear_arrays()
    return R_TB

def drgep(net, conds, targets, tols, returnobj='net', use_mpi=False, use_numpy=False):
    """
    Implementation of Directed Relation Graph with Error Propagation (DRGEP) reduction
    method described in Pepiot-Desjardins and Pitch 2008 (doi:10.1016/j.combustflame.2007.10.020)
    and Niemeyer and Sung 2011 (doi:10.1016/j.combustflame.2010.12.010).

    :param net: The network (RateCollection) to reduce.
    :param conds: A set of conditions to reduce over. Should either be a sequence of (composition,
        density, temperature) sequences/tuples if running in standard mode, or a sequence of 3
        sequences ((composition, density, temperature) ordering) if running in NumPy mode. In the
        latter case, the sequences will be permuted to create the dataset. The compositions should
        be pynucastro Composition objects.
    :param targets: A collection of target nuclei (or a single target nucleus) to run the
        graph search algorithm from. Should be supplied as pynucastro Nucleus objects.
    :param tols: Tolerance(s) or cutoff threshold(s) to use for paths from each of the target nuclei.
        Nuclei whose interaction coefficients do not meet the specified tolerance will have their
        interaction coefficients set to 0.0. Can be a single number (will be the same for all targets)
        or a separate value for each target nucleus.
    :param returnobj: The type of object to return. Options are 'net' (a reduced network, the default
        setting), 'nuclei' (unique nuclei with nonzero interaction coefficients, ordered so the
        interaction coefficients are descending), and 'coeff' (the interaction coefficients as a
        NumPy array, with entries corresponding to nuclei in *net.unique_nuclei*).
    :param use_mpi: Whether to divide up the set of conditions across MPI processes or not. Default
        setting is *False*.
    :param use_numpy: Whether to use NumPy to vectorize the interaction coefficient calculations or
        not. This is more memory intensive and may actually hinder performance for some setups.
        Conditions should be supplied as 3 lists that will be permuted to form the dataset (see
        *conds* parameter). Default setting is *False*.
    """

    #------------------
    # Process arguments
    #------------------

    if returnobj not in {'net', 'nuclei', 'coeff'}:
        raise ValueError(f"Invalid 'returnobj' argument: '{returnobj}'.")
    targets = to_list(targets)
    tols = to_list(tols, len(targets))

    #-------------------------------------------------------
    # Determine operation mode and launch appropriate helper
    #-------------------------------------------------------

    if use_mpi:
        if use_numpy:
            R_TB = _drgep_mpi_numpy(net, conds, targets, tols)
        else:
            R_TB = _drgep_mpi(net, conds, targets, tols)
    elif use_numpy:
        R_TB = _drgep_numpy(net, conds, targets, tols)
    else:
        R_TB = _drgep(net, conds, targets, tols)

    #------------------------
    # Return requested object
    #------------------------

    if returnobj == 'net':
        nuclei = [net.unique_nuclei[i] for i in range(len(net.unique_nuclei))
                  if R_TB[i] > 0.0]
        return net.linking_nuclei(nuclei)
    if returnobj == 'nuclei':
        idx = sorted(range(len(R_TB)), key=lambda i: R_TB[i], reverse=True)
        return [net.unique_nuclei[i] for i in idx if R_TB[i] > 0.0]
    if returnobj == 'coeff':
        return R_TB
    # returnobj was checked at the start of the function
    assert False