import astropy.io.fits as pyfits
from numpy import *
from jplephem.spk import SPK
from tdb2tdt import tdb2tdt

# To run the code one needs to have access to one of the ephemeride files provided by jpl:
# https://ssd.jpl.nasa.gov/ftp/eph/planets/bsp/
# here I use the DE421 file

kernel = SPK.open('de421.bsp')

c = 299792.458 # km/s  (kilometers per second)

def barycor(date,ra,dec, orbit=False, return_correction=True, approx_einstein=10):
    """Apply barycentric correction to vector
       date (in MJD) assuming coordinates ra,dec given in degrees
       optionally use orbit of a satellinte in heasarc format
       This includes geometric and shapiro corrections, but einstein
       correction is not implemented for now. Accurate to about 3e-4s.
       If return_correction=True (default), correction in seconds is returned.
       approx einstein is used to define number of points to calculate interpolated einstein correction (0 means do for each data point).
    see also https://asd.gsfc.nasa.gov/Craig.Markwardt/bary/ for detail and https://pages.physics.wisc.edu/~craigm/idl/ephem.html for the codebase
    which at least inspired the code below.
    v0.2: 21.03.2022 by V. Doroshenko
       """
    ra = radians(double(ra))
    dec = radians(double(dec))
    jd =  array(date,dtype=float64) + 2400000.5

    msol = 0.0002959122082855911

    # get positions and velocities of the sun, earth-moon barycenter and earth
    (x_sun,y_sun,z_sun),(vx_sun,vy_sun,vz_sun) = kernel[0,10].compute_and_differentiate(jd)
    (x_em,y_em,z_em),(vx_em,vy_em,vz_em) = kernel[0,3].compute_and_differentiate(jd)
    (x_e,y_e,z_e),(vx_e,vy_e,vz_e) = kernel[3,399].compute_and_differentiate(jd)

    x_earth = x_em + x_e
    y_earth = y_em + y_e
    z_earth = z_em + z_e

    vx_earth = (vx_em + vx_e)/86400 # velocities are always distance units per day as per jplephem documentation, so need to divide to get km/s
    vy_earth = (vy_em + vy_e)/86400
    vz_earth = (vz_em + vz_e)/86400

    if orbit:
        orbit = pyfits.open(orbit)
        mjdref = orbit[1].header['mjdreff']+orbit[1].header['mjdrefi']
        try:
            minmet = (min(date)-1.-mjdref)*86400
            maxmet =(max(date)+1.-mjdref)*86400
        except:
            minmet = (date-1-mjdref)*86400
            maxmet = (date+1-mjdref)*86400
        try:
            t = orbit[1].data.field('time')
        except:
            t = orbit[1].data.field('sclk_utc')
        mi, ma = t.searchsorted([minmet,maxmet])
        t = t[mi:ma]/86400.+mjdref

        # interpolate orbit to observed time and convert to km and km/s
        try:
            x_s = interp(date, t, orbit[1].data.field('x')[mi:ma]/1000.)
            y_s = interp(date, t, orbit[1].data.field('y')[mi:ma]/1000.)
            z_s = interp(date, t, orbit[1].data.field('z')[mi:ma]/1000.)

            vx_s = interp(date, t, orbit[1].data.field('vx')[mi:ma]/1000.)
            vy_s = interp(date, t, orbit[1].data.field('vy')[mi:ma]/1000.)
            vz_s = interp(date, t, orbit[1].data.field('vz')[mi:ma]/1000.)
        except:
            x_s = interp(date, t, orbit[1].data.field('pos_x')[mi:ma]/1000.)
            y_s = interp(date, t, orbit[1].data.field('pos_y')[mi:ma]/1000.)
            z_s = interp(date, t, orbit[1].data.field('pos_z')[mi:ma]/1000.)

            vx_s = interp(date, t, orbit[1].data.field('vel_x')[mi:ma]/1000.)
            vy_s = interp(date, t, orbit[1].data.field('vel_y')[mi:ma]/1000.)
            vz_s = interp(date, t, orbit[1].data.field('vel_z')[mi:ma]/1000.)

        x_obs, y_obs, z_obs = x_earth + x_s, y_earth + y_s, z_earth + z_s
        vx_obs, vy_obs, vz_obs = vx_earth + vx_s, vy_earth + vy_s, vz_earth + vz_s
        # orbital correction
        ocor = (vx_earth*x_s+vy_earth*y_s+vz_earth*z_s)/c**2
    else:
        x_obs, y_obs, z_obs = x_earth, y_earth , z_earth
        vx_obs, vy_obs, vz_obs = vx_earth , vy_earth , vz_earth
        ocor = 0.

    # #components of the object unit vector:
    x_obj = cos(dec)*cos(ra)
    y_obj = cos(dec)*sin(ra)
    z_obj = sin(dec)

    #geometric correction
    geo_corr = (x_obs*x_obj + y_obs*y_obj + z_obs*z_obj)/c

    #einstein correction
    if approx_einstein == 0:
        einstein_corr = tdb2tdt(jd)
    else:
        xx = linspace(jd.min(),jd.max(),approx_einstein)
        einstein_corr = tdb2tdt(xx)
        einstein_corr = interp(jd,xx,einstein_corr)
    #shapiro correction ("Shapiro") = - (2 G Msun/c^3) log(1 + cos th)
    sun_dist = sqrt((x_sun-x_obs)**2+(y_sun-y_obs)**2+(z_sun-z_obs)**2)
    costh = ((x_obs-x_sun)*x_obj+(y_obs-y_sun)*y_obj + (z_obs-z_sun)*z_obj)/sun_dist
    shapiro_corr = - 9.8509819e-06*log(1.+costh)
    corr = geo_corr + ocor + einstein_corr - shapiro_corr
    if return_correction:
        return corr
    else:
        return date + corr/86400.


import queue
import sys
from collections import namedtuple

from comm_support_lib.comm_interfaces.debug_cli import DebugCLI
from utils.common.cli_command_consts import CliCommandConsts
from utils.common.cli_regex_consts import CliRegexConsts
from utils.common.dbus_func_consts import DbusFuncConsts
from utils.common.dbus_signal_consts import DbusSignalConsts
from utils.config.config import CLI_DBUS_SIGNAL_TIMEOUT, CLI_DBUS_FUNCTION_TIMEOUT, CLI_DBUS_SUBSCR_TIMEOUT, \
    CLI_DBUS_UNSUBSCR_TIMEOUT

class CliDbusUtil:
    DBusSignalResult = namedtuple("DBusSignalResult", ["signal", "result_list"])

    __RETURN_TYPE_BOOL = "boolean"
    __RETURN_TYPE_STRING = "string"
    __RETURN_BOOL_TRUE = "true"

    __SIGNALS_WITH_TWO_RESULTS = [DbusSignalConsts.NEW_PACKAGE_AVAILABLE,
                                  DbusSignalConsts.FORCED_PACKAGE_CHECKED,
                                  DbusSignalConsts.PACKAGE_UPDATE_STATE]

    __RETURN_REGEX_LIST_FUNCTIONS = {
        DbusFuncConsts.GET_CURR_PARTITION: CliRegexConsts.REGEX_DBUS_RESULT_PARTITION,
        DbusFuncConsts.GET_CURR_SW_VERSION: CliRegexConsts.REGEX_DBUS_RESULT_FW_VERSION,
        DbusFuncConsts.GET_ALT_SW_VERSION: CliRegexConsts.REGEX_DBUS_RESULT_FW_VERSION,
        DbusFuncConsts.SUSPEND_FW_UPDATE: CliRegexConsts.REGEX_DBUS_RESULT_BOOL,
        DbusFuncConsts.RESUME_FW_UPDATE: CliRegexConsts.REGEX_DBUS_RESULT_BOOL,
        DbusFuncConsts.REJECT_FW_UPDATE: CliRegexConsts.REGEX_DBUS_RESULT_BOOL,
        DbusFuncConsts.FORCE_UPDATE: CliRegexConsts.REGEX_DBUS_RESULT_BOOL,
        DbusFuncConsts.FORCE_FW_UPDATE: CliRegexConsts.REGEX_DBUS_RESULT_BOOL,
        DbusFuncConsts.SWITCH_TO_ALT_FW: CliRegexConsts.REGEX_DBUS_RESULT_BOOL,
        DbusFuncConsts.SUSPEND_PACKAGE_UPDATE: CliRegexConsts.REGEX_DBUS_RESULT_BOOL,
        DbusFuncConsts.RESUME_PACKAGE_UPDATE: CliRegexConsts.REGEX_DBUS_RESULT_BOOL,
        DbusFuncConsts.REJECT_PACKAGE_UPDATE: CliRegexConsts.REGEX_DBUS_RESULT_BOOL,
        DbusFuncConsts.FORCE_PACKAGE_UPDATE: CliRegexConsts.REGEX_DBUS_RESULT_BOOL,
        DbusFuncConsts.GET_SW_VERSION: CliRegexConsts.REGEX_DBUS_RESULT_SW_VERSION,
        DbusFuncConsts.GET_CURRENT_BOOT_DEVICE: CliRegexConsts.REGEX_DBUS_RESULT_BOOT_DEV}

    __RETURN_REGEX_LIST_SIGNALS = {
        DbusSignalConsts.NEW_FIMWARE_AVAILABLE: [CliRegexConsts.REGEX_DBUS_RESULT_FW_VERSION],
        DbusSignalConsts.FORCED_FIRMWARE_CHECKED: [CliRegexConsts.REGEX_DBUS_RESULT_BOOL],
        DbusSignalConsts.FIRMWARE_UPDATE_STATE: [CliRegexConsts.REGEX_DBUS_RESULT_FW_UPDATE_STATE],
        DbusSignalConsts.FIRMWARE_CHECK_RESULTS: [CliRegexConsts.REGEX_DBUS_CHECK_RESULTS],
        DbusSignalConsts.NEW_PACKAGE_AVAILABLE:
            [CliRegexConsts.REGEX_DBUS_RESULT_PACK_NAME, CliRegexConsts.REGEX_DBUS_RESULT_FW_VERSION],
        DbusSignalConsts.FORCED_PACKAGE_CHECKED:
            [CliRegexConsts.REGEX_DBUS_RESULT_PACK_NAME, CliRegexConsts.REGEX_DBUS_RESULT_BOOL],
        DbusSignalConsts.PACKAGE_UPDATE_STATE:
            [CliRegexConsts.REGEX_DBUS_RESULT_PACK_NAME, CliRegexConsts.REGEX_DBUS_RESULT_PACKAGE_UPDATE_STATE],
        DbusSignalConsts.PACKAGE_CHECK_RESULTS: [CliRegexConsts.REGEX_DBUS_CHECK_RESULTS]}

    __expected_result_for_signal = None
    __expected_count_of_results = 0

    def __get_result_count_for_signal(self, signal: str):
        if signal in self.__SIGNALS_WITH_TWO_RESULTS:
            return 2
        else:
            return 1

    def __incoming_data_callback(self, msg: str):
        if CliRegexConsts.REGEX_DBUS_SIGNAL_STRING.search(msg):
            signal = self.__get_signal_from_string(msg)
            if signal is None:
                # signal is unknown of wrong message has passed
                return
            if signal in self.__signal_notification_dict.keys():
                self.__expected_count_of_results = self.__get_result_count_for_signal(signal)
                self.__expected_result_for_signal = self.DBusSignalResult(signal, [])
            return
        if CliRegexConsts.REGEX_DBUS_COMMON_RESULT.search(msg) and self.__expected_result_for_signal:
            for regex in self.__RETURN_REGEX_LIST_SIGNALS[self.__expected_result_for_signal.signal]:
                if regex.search(msg):
                    break
            self.__expected_result_for_signal.result_list.append(self.__get_result_from_string(msg))
            self.__expected_count_of_results -= 1
            if self.__expected_count_of_results == 0:
                self.__signal_queue.put(self.__expected_result_for_signal)
                self.__expected_result_for_signal = None
            return

    @staticmethod
    def __get_signal_from_string(string: str):
        search_result = CliRegexConsts.REGEX_DBUS_SIGNAL_MEMBER.search(string)
        if search_result:
            return search_result.group(0).lstrip(CliCommandConsts.DBUS_SIGNAL_COMMAND_MEMBER)
        else:
            return None

    @staticmethod
    def __put_signal_to_string(signal: str, string: str):
        signal_position = string.rfind(CliCommandConsts.DBUS_SIGNAL_COMMAND_MEMBER) + len(
            CliCommandConsts.DBUS_SIGNAL_COMMAND_MEMBER)
        if signal_position:
            return string[:signal_position] + signal + string[signal_position:]
        else:
            return None

    def __get_result_from_string(self, result_string):
        method_result = result_string.lstrip()
        if self.__RETURN_TYPE_BOOL in method_result:
            if self.__RETURN_BOOL_TRUE in method_result:
                return True
            else:
                return False
        elif self.__RETURN_TYPE_STRING in method_result:
            return method_result.strip("string").replace('\"', '').strip()

    def __init__(self, cli: DebugCLI):
        """
        Class constructor. Initialize the utility.
        :param cli: Debug CLI object
        """
        self.__cli = cli
        self.__signal_notification_dict = dict()
        self.__signal_queue = queue.Queue()
        self.__cli.register_message_callback(self.__incoming_data_callback, {CliRegexConsts.REGEX_DBUS_SIGNAL_STRING,
                                                                             CliRegexConsts.REGEX_DBUS_COMMON_RESULT})

    def subscribe_signal_notification(self, signal: str, timeout: float = CLI_DBUS_SUBSCR_TIMEOUT) -> bool:
        """
        Subscribes on D-Bus signal notification. After that, all the signal notifications will be placed into the queue
        and the user will be available to check them.
        :param timeout: signal subscription result waiting timeout
        :param signal: D-Bus signal to subscribe on.
        :return: True on success, False if the signal notification has already subscribed or some error occurred.
        """
        if signal in self.__signal_notification_dict:
            print("Has already subscribed on signal:" + signal, file=sys.stderr)
            return False

        command_string = self.__put_signal_to_string(signal, CliCommandConsts.COMMAND_DBUS_SIGNAL)
        self.__cli.flush_incoming_data()
        self.__cli.send_message(command_string)
        result = self.__cli.get_message(timeout, CliRegexConsts.REGEX_SIGNAL_SUBSCR_RESULT)

        if result is not None:
            dbus_monitor_process: str = CliRegexConsts.REGEX_DBUS_MONITOR_PROCESS.search(result).group(0)
            self.__signal_notification_dict[signal] = dbus_monitor_process
            return True
        else:
            print("Something went wrong during subscription on D-Bus signal: " + signal, file=sys.stderr)
            return False

    def unsubscribe_signal_notification(self, signal: str, timeout: float = CLI_DBUS_UNSUBSCR_TIMEOUT) -> bool:
        """
        Unsubscribes from D-Bus signal notification. After that, the signal notifications will not be placed into
        the queue.
        :param signal: D-Bus signal to unsubscribe.
        :param timeout: Timeout waiting for unsubscribing on the signal
        :return: True on success, False if the signal notification is not present in subscription list.
        """
        if signal not in self.__signal_notification_dict:
            print("No subscription found for signal:" + signal, file=sys.stderr)
            return False

        self.__cli.flush_incoming_data()
        self.__cli.send_message(CliCommandConsts.COMMAND_PS_WITH_GREP + self.__signal_notification_dict[signal])

        if self.__cli.get_message(timeout, CliRegexConsts.REGEX_DBUS_MONITOR) is not None:
            command_string = CliCommandConsts.COMMAND_KILL + self.__signal_notification_dict[signal]
            self.__cli.send_message(command_string)
            self.__cli.get_message(timeout, CliRegexConsts.REGEX_LOGGED_IN)
        self.__signal_notification_dict.pop(signal)
        return True

    def clear_subscription_list(self):
        """
        Clear signal subscription list with unsubscribing on those signals
        """
        for subscription in list(self.__signal_notification_dict):
            self.unsubscribe_signal_notification(subscription)

    def run_method(self, method: str, parameter: str = None,
                   timeout: float = CLI_DBUS_FUNCTION_TIMEOUT, expected_return: bool = True) -> str or None:
        """
        Runs D-Bus function using Debug CLI. Could take string parameter to pass with the function. Returns function
        execution result.
        :param method: method to be invoked.
        :param parameter: string parameter to be passed together with the method. By default is None.
        :param timeout: result waiting timeout in seconds.
        :param expected_return: Flag which represents does the method should returns some result.
        :return: Returns None if there is no result of the function executed or timeout has occurred. Otherwise –
        returns function execution result string. If expected_return is False - always returns None without delay.
        """
        command_string = CliCommandConsts.COMMAND_DBUS_METHOD + method

        if parameter:
            command_string += " string:\"" + parameter + "\""

        self.__cli.flush_incoming_data()
        self.__cli.send_message(command_string)

        if not expected_return:
            return None

        method_result_regex = self.__RETURN_REGEX_LIST_FUNCTIONS[method]
        if method_result_regex is None:
            print("run_method. No return regex from method: " + method, file=sys.stderr)
            return None

        method_result: str = self.__cli.get_message(timeout, method_result_regex)
        if method_result:
            method_result = method_result.lstrip()
            return self.__get_result_from_string(method_result)
        else:
            print("Something went wrong. No response from method: " + method, file=sys.stderr)
            return None

    def get_signal(self, timeout: float = CLI_DBUS_SIGNAL_TIMEOUT) -> str or None:
        """
        Performs getting D-Bus signals notifications, on which has subscribed before. If timeout is non-zero,
        the method blocks for required time for waiting a signal notification from the board.
        :param timeout: timeout waiting for available signal notifications in message queue. Represents time in seconds.
        :return: None if nothing to read from message queue, or timeout occurred.
        Otherwise – returns signal notification string in format "Signal: *SIGNAL_NAME*; Result: *RESULT*".
        """
        try:
            return self.__signal_queue.get(True, timeout)
        except queue.Empty:
            return None

    def clear_signal_list(self):
        """
        Clear signal list
        """
        while True:
            try:
                self.__signal_queue.get(block=False)
            except queue.Empty:
                break


import pygame as pyg
import chessStorage

Width = Height = 512
dim = 8
KSize = Height // dim
Fps = 20
Images = {}

def load_images():
    pieces = ['WP', 'BP', 'WR', 'BR', 'WB', 'BB', 'WN', 'BN', 'WQ', 'BQ', 'WK', 'BK']
    for piece in pieces:
        Images[piece] = pyg.transform.scale(pyg.image.load('chesspieces/' + piece + '.png'), (KSize, KSize))

def main():
    pyg.init()
    screen = pyg.display.set_mode((Width, Height))
    clock = pyg.time.Clock()
    screen.fill(pyg.Color('white'))
    gameState = chessStorage.GameState()
    validMoves = gameState.getValidMove()
    moveDone = False
    GameOver = False
    load_images()
    run = True
    KSelected = ()
    PlayerClick = []
    while run:
        for e in pyg.event.get():
            if e.type == pyg.QUIT:
                run = False
            elif e.type == pyg.MOUSEBUTTONDOWN:
                location = pyg.mouse.get_pos()
                coll = location[0]//KSize
                row = location[1]//KSize
                print(coll,row)
                if KSelected == (row, coll):
                    KSelected = ()
                    PlayerClick = []
                else:
                    KSelected = (row, coll)
                    PlayerClick.append(KSelected)
                if len(PlayerClick) == 2:
                    move = chessStorage.Move(PlayerClick[0], PlayerClick[1], gameState.board)
                    for i in range(len(validMoves)):
                        if move == validMoves[i]:
                            gameState.makeMove(validMoves[i])
                            moveDone = True
                            print(move.getChessNotation())
                            KSelected =()
                            PlayerClick = []
                    if not moveDone:
                        PlayerClick = [KSelected]
            elif e.type == pyg.KEYDOWN:
                if e.key == pyg.K_u:
                    gameState.undo()
                    moveDone = True
        if moveDone:
            validMoves = gameState.getValidMove()
            moveDone = False

        if gameState.CheckMate:
            GameOver = True
            if gameState.whiteToStart:
                messages(screen,'Black wins by checkmate')
            else:
                messages(screen,'White wins by checkmate')
        elif gameState.StaleMate:
            GameOver = True
            messages(screen,'Draw by stalemate')
        clock.tick(Fps)
        pyg.display.flip()
        drawGame(screen, gameState, validMoves, KSelected)
def highlightSquare(screen, gameState, validMoves, KSelected):
    if KSelected != ():
        row, coll = KSelected
        if gameState.board[row][coll][0] == ('W' if gameState.whiteToStart else 'B'):
            s = pyg.Surface((KSize, KSize))
            s.set_alpha(100)
            s.fill(pyg.Color('blue'))
            screen.blit(s, (coll * KSize, row * KSize))
            s.fill(pyg.Color('yellow'))
            for move in validMoves:
                if move.startRow == row and move.startColl == coll:
                    screen.blit(s, (move.secondColl * KSize, move.secondRow * KSize))

def drawGame(screen, gameState, validMoves, KSelected):
    drawBoard(screen)
    drawPieces(screen, gameState.board)
    highlightSquare(screen, gameState, validMoves, KSelected)

def drawBoard(screen):
    colors = [pyg.Color('white'), pyg.Color('Dark green')]
    for i in range(dim):
        for j in range(dim):
            color = colors[(i+j)%2]
            pyg.draw.rect(screen, color, pyg.Rect(j*KSize, i*KSize, KSize, KSize))

def drawPieces(screen, board):
    for i in range(dim):
        for j in range(dim):
            piece = board[i][j]
            if piece != '--':
                screen.blit(Images[piece], pyg.Rect(j*KSize, i*KSize, KSize, KSize))

def messages(screen,text):
    font = pyg.font.SysFont('Arial', 32, True, False)
    message = font.render(text, 0, pyg.Color('Black'))
    location = pyg.Rect(0,0,Width,Height).move(Width/2 - message.get_width()/2,Height/2 - message.get_height()/2)
    screen.blit(message,location)

if __name__ == '__main__':
    main()


# -*- coding: utf-8 -*
"""novelWriter Book Tree Class

 novelWriter – Book Tree Class
===============================
 Holds the tree of items and files of the book project

 File History:
 Created: 2017-10-18 [0.4.0]

"""

import logging
import nw

from os           import path
from time         import time
from hashlib      import sha256
from itertools    import chain
from nw.file.item import BookItem
from nw.file.doc  import DocFile

logger = logging.getLogger(__name__)

class BookTree():

    ORD_UP    = 0
    ORD_DOWN  = 1
    ORD_NUP   = 2
    ORD_NDOWN = 3

    validOrder = [ORD_UP,ORD_DOWN,ORD_NUP,ORD_NDOWN]

    def __init__(self):

        self.docPath    = None
        self.theTree    = []
        self.treeLookup = {}
        self.parOfItems = {}
        self.parOfFiles = {}
        self.treeOrder  = []

        self.fixedOrder = [
            BookItem.TYP_BOOK,
            BookItem.TYP_CHAR,
            BookItem.TYP_PLOT,
            BookItem.TYP_NOTE,
        ]
        self.fixedItems = {
            BookItem.TYP_BOOK : None,
            BookItem.TYP_CHAR : None,
            BookItem.TYP_PLOT : None,
            BookItem.TYP_NOTE : None,
        }

        return

    def clearTree(self):

        self.docPath    = None
        self.theTree    = []
        self.treeLookup = {}
        self.parOfItems = {}
        self.parOfFiles = {}
        self.treeOrder  = []

        self.fixedOrder = [
            BookItem.TYP_BOOK,
            BookItem.TYP_CHAR,
            BookItem.TYP_PLOT,
            BookItem.TYP_NOTE,
        ]
        self.fixedItems = {
            BookItem.TYP_BOOK : None,
            BookItem.TYP_CHAR : None,
            BookItem.TYP_PLOT : None,
            BookItem.TYP_NOTE : None,
        }

        return

    #
    # Add Elements to Main Tree
    #

    def addFile(self, pHandle):

        parEntry = self.getItem(pHandle)["entry"]
        parClass = parEntry.itemClass
        parType  = parEntry.itemType

        if not parClass == BookItem.CLS_CONT:
            logger.debug("BookTree: Entry is not a container, getting its parent")
            parParent = self.getItem(pHandle)["parent"]
            if not parParent is None:
                pHandle  = parParent
                parEntry = self.getItem(pHandle)["entry"]
                parClass = parEntry.itemClass
                parType  = parEntry.itemType
            else:
                logger.error("BookTree: A file must be added to a folder")

        if parType == BookItem.TYP_BOOK:
            newClass   = BookItem.CLS_SCENE
            newName    = "New Scene"
            newCompile = True
        else:
            newClass   = BookItem.CLS_NOTE
            newName    = "New Note"
            newCompile = None

        newItem = BookItem()
        newItem.setClass(newClass)
        newItem.setLevel(BookItem.LEV_FILE)
        newItem.setType(parType)
        newItem.setName(newName)
        newItem.setCompile(newCompile)

        self.appendItem(None,pHandle,None,newItem)
        self.sortTree()

        return

    def addChapter(self):

        newItem = BookItem()
        newItem.setClass(BookItem.CLS_CONT)
        newItem.setLevel(BookItem.LEV_ITEM)
        newItem.setType(BookItem.TYP_BOOK)
        newItem.setSubType(BookItem.SUB_CHAP)
        newItem.setName("New Chapter")
        newItem.setCompile(True)

        self.appendItem(None,self.fixedItems[BookItem.TYP_BOOK],None,newItem)
        self.sortTree()

        return

    def addCharacter(self):

        newItem = BookItem()
        newItem.setClass(BookItem.CLS_CONT)
        newItem.setLevel(BookItem.LEV_ITEM)
        newItem.setType(BookItem.TYP_CHAR)
        newItem.setName("New Character")

        self.appendItem(None,self.fixedItems[BookItem.TYP_CHAR],None,newItem)
        self.sortTree()

        return

    def addPlot(self):

        newItem = BookItem()
        newItem.setClass(BookItem.CLS_CONT)
        newItem.setLevel(BookItem.LEV_ITEM)
        newItem.setType(BookItem.TYP_PLOT)
        newItem.setName("New Plot")

        self.appendItem(None,self.fixedItems[BookItem.TYP_PLOT],None,newItem)
        self.sortTree()

        return

    #
    # Tree Item Maintenance
    #

    def updateItem(self, itemHandle, tTag, tValue):
        self.theTree[self.treeLookup[itemHandle]]["entry"].setFromTag(tTag,tValue)
        return

    def getItem(self, itemHandle):
        return self.theTree[self.treeLookup[itemHandle]]

    def changeOrder(self, itemHandle, moveIt):

        treeItem   = self.getItem(itemHandle)
        itemParent = treeItem["parent"]
        itemEntry  = treeItem["entry"]

        if itemEntry.itemLevel == BookItem.LEV_FILE:
            currList = self.parOfFiles[itemParent]
        elif itemEntry.itemLevel == BookItem.LEV_ITEM:
            currList = self.parOfItems[itemParent]
        else:
            logger.error("BookTree: Cannot change order of ROOT elements")
            return

        if moveIt not in self.validOrder:
            logger.error("BUG: Unknown ordering requested")
            return

        if itemHandle not in currList:
            logger.error("BUG: Cannot change order of %s, as it is not where it should be" % itemHandle)
            return

        # If move withing parent, just swap the indices, and refresh the master index
        if moveIt == self.ORD_UP or moveIt == self.ORD_DOWN:

            currIndex = currList.index(itemHandle)

            if moveIt == self.ORD_UP:
                newIndex = currIndex - 1
            elif moveIt == self.ORD_DOWN:
                newIndex = currIndex + 1

            if newIndex < 0 or newIndex >= len(currList): return

            currList[currIndex], currList[newIndex] = currList[newIndex], currList[currIndex]

            if itemEntry.itemLevel == BookItem.LEV_FILE:
                self.parOfFiles[itemParent] = currList
            elif itemEntry.itemLevel == BookItem.LEV_ITEM:
                self.parOfItems[itemParent] = currList

            self.buildTreeOrder()
            self.updateEntryOrder()

        # If moving to a new node, update parent and set order to None
        elif moveIt == self.ORD_NUP or moveIt == self.ORD_NDOWN:

            # This is only allowed for FILE entries
            if itemEntry.itemLevel == BookItem.LEV_ITEM: return

            treeParent = self.getItem(itemParent)
            parParent  = treeParent["parent"]
            parEntry   = treeParent["entry"]

            if parParent is not None:
                # Move to next or previous node
                parList = self.parOfItems[parParent]
                if itemParent not in parList:
                    logger.error("BUG: Something unexpected happened while moving entry")
                    return
                parIndex = parList.index(itemParent)
                if moveIt == self.ORD_NUP:
                    newIndex = parIndex - 1
                elif moveIt == self.ORD_NDOWN:
                    newIndex = parIndex + 1

                if newIndex < 0 or newIndex >= len(parList):
                    self.theTree[self.treeLookup[itemHandle]]["parent"] = parParent
                    self.sortTree()
                else:
                    self.theTree[self.treeLookup[itemHandle]]["parent"] = parList[newIndex]
                    self.sortTree()
            else:
                # Move to last child node
                parList = self.parOfItems[itemParent]
                self.theTree[self.treeLookup[itemHandle]]["parent"] = parList[len(parList)-1]
                self.sortTree()

        return

    def createRootItem(self, rootType):

        if rootType in BookItem.validTypes:
            if rootType == BookItem.TYP_BOOK: rootName = "Book"
            if rootType == BookItem.TYP_CHAR: rootName = "Characters"
            if rootType == BookItem.TYP_PLOT: rootName = "Plots"
            if rootType == BookItem.TYP_NOTE: rootName = "Notes"
        else:
            logger.error("BookTree: Cannot create root item of type '%s'" % str(rootType))
            return

        if self.fixedItems[rootType] is not None:
            logger.warning("BookTree: Root item already exists")
            return

        rootHandle = self.makeHandle()
        rootItem   = BookItem()
        rootItem.setClass(BookItem.CLS_CONT)
        rootItem.setLevel(BookItem.LEV_ROOT)
        rootItem.setType(rootType)
        rootItem.setName(rootName)

        logger.info("BookTree: Creating root item '%s' with handle %s" % (rootName, rootHandle))

        self.appendItem(rootHandle,None,None,rootItem)
        self.fixedItems[rootType] = rootHandle

        return

    def appendItem(self, tHandle, pHandle, tOrder, bookItem):
        """
        Appends an entry to the main project tree.
        """

        tHandle = self.checkString(tHandle,self.makeHandle(),False)
        pHandle = self.checkString(pHandle,None,True)
        tOrder  = self.checkInt(tOrder,None,True)

        logger.verbose("BookTree: Adding entry %s with parent %s" % (str(tHandle),str(pHandle)))

        if bookItem.itemLevel == BookItem.LEV_FILE:
            docItem = DocFile(self.docPath,tHandle,bookItem.itemClass)
        else:
            docItem = None

        self.theTree.append({
            "handle" : tHandle,
            "parent" : pHandle,
            "order"  : tOrder,
            "entry"  : bookItem,
            "doc"    : docItem,
        })
        lastIdx = len(self.theTree)-1
        self.treeLookup[tHandle] = lastIdx

        return

    def validateTree(self):

        errCount = 0

        # Checking ROOT level
        for treeItem in self.theTree:

            itemHandle = treeItem["handle"]
            itemParent = treeItem["parent"]
            bookEntry  = treeItem["entry"]
            itemIdx    = self.treeLookup[itemHandle]

            if not bookEntry.itemLevel == BookItem.LEV_ROOT: continue
            logger.verbose("BookTree: Checking ROOT with handle %s" % itemHandle)

            if itemParent is not None:
                self.theTree[itemIdx]["parent"] = None
                logger.warning("BookTree: Parent was set for ROOT element %s" % itemHandle)
                errCount += 1

            for itemType in BookItem.validTypes:
                if bookEntry.itemType == itemType:
                    if self.fixedItems[itemType] is None:
                        self.fixedItems[itemType] = itemHandle
                        logger.debug("BookTree: Root handle for type %s set to %s" % (itemType,itemHandle))
                    else:
                        logger.warning("BookTree: Encountered a second ROOT of type %s with handle %s" % (itemType,itemHandle))
                        errCount += 1

        for rootType in self.fixedItems.keys():
            if self.fixedItems[rootType] is None:
                logger.info("BookTree: Root item missing, creating it")
                self.createRootItem(rootType)

        # Checking ITEM level
        for treeItem in self.theTree:

            itemHandle = treeItem["handle"]
            itemParent = treeItem["parent"]
            bookEntry  = treeItem["entry"]
            itemIdx    = self.treeLookup[itemHandle]
            hasError   = False

            if not bookEntry.itemLevel == BookItem.LEV_ITEM: continue
            logger.verbose("BookTree: Checking ITEM with handle %s" % itemHandle)

            for itemType in BookItem.validTypes:
                if bookEntry.itemType == itemType:
                    if itemParent is None:
                        logger.warning("BookTree: Parent was missing for ITEM of type %s with handle %s" % (itemType,itemHandle))
                        self.theTree[itemIdx]["parent"] = self.fixedItems[itemType]
                        errCount += 1

        logger.info("BookTree: Found %d error(s) while parsing the project tree" % errCount)

        return

    def sortTree(self):

        # Resetting Indices
        self.parOfItems = {}
        self.parOfFiles = {}

        treeOrder = []
        logger.debug("TreeSort: Reading previous order")
        tempOrder = [None] * len(self.theTree)
        for treeItem in self.theTree:
            itemHandle = treeItem["handle"]
            itemOrder  = treeItem["order"]
            itemName   = treeItem["entry"].itemName
            if itemOrder is not None and isinstance(itemOrder,int):
                if tempOrder[itemOrder] is None:
                    tempOrder[itemOrder] = itemHandle
                    logger.vverbose("TreeSort: Entry '%s' %s has order %s" % (
                        str(itemName), str(itemHandle), str(itemOrder)
                    ))
                else:
                    tempOrder.append(itemHandle)
                    logger.vverbose("TreeSort: Entry '%s' %s has no order, appending" % (
                        str(itemName), str(itemHandle)
                    ))
            else:
                tempOrder.append(itemHandle)
                logger.vverbose("TreeSort: Entry '%s' %s has no order, appending" % (
                    str(itemName), str(itemHandle)
                ))
        for tempItem in tempOrder:
            if tempItem is not None: treeOrder.append(tempItem)

        # Scanning ROOT level
        logger.debug("TreeSort: Sorting ROOT entries")
        for rootType in self.fixedOrder:
            itemHandle = self.fixedItems[rootType]
            self.parOfItems[itemHandle] = []
            self.parOfFiles[itemHandle] = []

        # Scanning ITEM level
        logger.debug("TreeSort: Sorting ITEM entries")
        for itemHandle in treeOrder:

            itemIdx    = self.treeLookup[itemHandle]
            treeItem   = self.theTree[itemIdx]

            itemHandle = treeItem["handle"]
            itemParent = treeItem["parent"]
            bookEntry  = treeItem["entry"]

            if not bookEntry.itemLevel == BookItem.LEV_ITEM: continue

            if itemParent in self.parOfItems.keys():
                self.parOfItems[itemParent].append(itemHandle)
                logger.vverbose("TreeSort: ITEM '%s' %s appended to %s" % (
                    bookEntry.itemName, str(itemHandle), str(itemParent)
                ))
                self.parOfFiles[itemHandle] = []
            else:
                logger.warning("BUG: itemParent %s not found in itemParent" % itemParent)

        # Scanning FILE level
        logger.verbose("TreeSort: Sorting FILE entries")
        for itemHandle in treeOrder:

            itemIdx    = self.treeLookup[itemHandle]
            treeItem   = self.theTree[itemIdx]

            itemHandle = treeItem["handle"]
            itemParent = treeItem["parent"]
            bookEntry  = treeItem["entry"]

            if not bookEntry.itemLevel == BookItem.LEV_FILE: continue

            if itemParent in self.parOfFiles.keys():
                self.parOfFiles[itemParent].append(itemHandle)
                logger.vverbose("TreeSort: FILE '%s' %s appended to %s" % (
                    bookEntry.itemName, str(itemHandle), str(itemParent)
                ))
            else:
                logger.warning("BUG: itemParent %s not found in fileParent" % itemParent)

        self.buildTreeOrder()
        self.updateEntryOrder()

        return

    def buildTreeOrder(self):

        # Resetting Indices
        self.treeOrder = []

        rootOrder = []
        itemOrder = []
        fileOrder = []

        for rootType in self.fixedOrder:
            itemHandle = self.fixedItems[rootType]
            rootOrder.append(itemHandle)

        logger.debug("TreeSort: %d ROOT entries added to index" % len(rootOrder))

        for itemParent in rootOrder:
            itemOrder += self.parOfItems[itemParent]

        logger.debug("TreeSort: %d ITEM entries added to index" % len(itemOrder))

        for itemParent in chain(rootOrder,itemOrder):
            fileOrder += self.parOfFiles[itemParent]

        logger.debug("TreeSort: %d FILE entries added to index" % len(fileOrder))

        errCount = 0
        logger.debug("TreeSort: Assempling index, and checking for consistency")
        self.treeOrder = rootOrder + itemOrder + fileOrder
        for itemHandle in self.treeLookup.keys():
            if itemHandle not in self.treeOrder:
                logger.warning("BUG: Handle %s not in index" % itemHandle)
                errCount += 1
        if errCount == 0:
            logger.debug("TreeSort: Index is consistent")
        else:
            logger.warning("BUG: %d errors found in the index" % errCount)

        uniqueSet = set(self.treeOrder)
        if len(uniqueSet) == len(self.treeOrder):
            logger.debug("TreeSort: No duplicates found in index")

        return

    def updateEntryOrder(self):

        logger.debug("TreeSort: Setting order parameter of tree entries")
        for itemOrder in range(len(self.treeOrder)):
            itemHandle = self.treeOrder[itemOrder]
            itemIdx    = self.treeLookup[itemHandle]
            self.theTree[itemIdx]["order"] = itemOrder
            logger.vverbose("TreeSort: Setting '%s' %s to order %d" % (
                str(self.theTree[itemIdx]["entry"].itemName), itemHandle, itemOrder
            ))

        return

    #
    # Setters and Getters
    #

    def setPath(self, docPath):
        self.docPath = docPath
        return

    #
    # Internal Functions
    #

    def makeHandle(self,seed=""):
        itemHandle = sha256((str(time())+seed).encode()).hexdigest()[0:13]
        if itemHandle in self.treeLookup.keys():
            logger.warning("BookTree: Duplicate handle encountered! Retrying ...")
            itemHandle = self.makeHandle(seed+"!")
        return itemHandle

    def checkString(self,checkValue,defaultValue,allowNone=False):
        if allowNone:
            if checkValue == None:   return None
            if checkValue == "None": return None
        if isinstance(checkValue,str): return str(checkValue)
        return defaultValue

    def checkInt(self,checkValue,defaultValue,allowNone=False):
        if allowNone:
            if checkValue == None:   return None
            if checkValue == "None": return None
        try:
            return int(checkValue)
        except:
            return defaultValue

# End Class BookTree


import os
import sys
import inspect
import collections
import threading

class HidingLock(object):
    def __init__(self, obj, lock=None):
        self.lock = lock or threading.RLock()
        self._obj = obj

    def __enter__(self):
        self.lock.acquire()
        return self._obj

    def __exit__(self, exc_type, exc_value, traceback):
        self.lock.release()

from pimpmyclass.mixins import LogMixin

from lantz.core.driver import Driver, Base
from lantz.core.flock import initialize_many, finalize_many
from lantz.core.helpers import MISSING

from .connect import connect_setup, connect_driver, connect_feat
from .widgets import DriverTestWidget, SetupTestWidget
from .objwrapper import QDriver
from .utils.qt import QtCore, QtGui, SuperQObject, MetaQObject
from .log import get_logger, LOGGER

ICON_FEDORA = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'assets', 'fedora.png')

class ThreadLogMixin:

    _PP_THREADS = HidingLock({})

    def log_current_thread(self, msg='Current thread is'):
        thread = self.thread()
        thid = int(thread.currentThreadId())
        thname = thread.objectName()

        with self._PP_THREADS as d:
            if thid in d:
                pp = d[thid]
            else:
                pp = d[thid] = len(d)

        self.log_debug(msg + ' %s (%d-%s)' % (thname, pp, thid))

    def moveToThread(self, thread):
        super().moveToThread(thread)
        self.log_current_thread('Moved to thread:')

class InstrumentSlot:

    def __str__(self):
        return '<InstrumentSlot>'

    def initialize(self):
        LOGGER.warning('The Instrument slot {} has not been assigned '
                       'to an actual instrument', self._slot_name )

    def finalize(self):
        pass

class FlockSlot:

    def __init__(self, *required_names):
        self.require_names = required_names

    def __str__(self):
        return '<FlockSlot>'

    def initialize(self):
        LOGGER.warning('The Flock slot {} has not been assigned '
                       'to an actual flock', self._slot_name )

    def finalize(self):
        pass

class BackendSlot:

    def __str__(self):
        return '<BackendSlot>'

    def initialize(self):
        LOGGER.warning('The Backend slot {} has not been assigned '
                       'to an actual instrument', self._slot_name )

    def finalize(self):
        pass

Front2Back = collections.namedtuple('Front2Back', 'frontend_class backend_name')

class Back2Back(collections.namedtuple('Back2Back', 'backend_class local_attribute foreign_attribute')):

    def link(self, local_name, foreign_name=None):
        return Back2Back(self.backend_class, local_name, foreign_name or local_name)

class _BackendType(MetaQObject):

    def create_instrument_property(cls, key):
        def getter(self):
            return self.instruments[key]
        def setter(self, instrument):
            if not isinstance(instrument, QtCore.QObject):
                raise ValueError("The object provided for InstrumentSlot '%s' is not subclass of QObject.\n"
                                 "Maybe you forgot to wrap your driver class?\n"
                                 "  from lantz.qt import wrap_driver_cls\n"
                                 "  QMyDriver = wrap_driver_cls(QMyDriver)\n" % key)
            instrument.name = key
            self.instruments[key] = instrument

        return property(getter, setter)

    def create_flock_property(cls, key):
        def getter(self):
            return self.flocks[key]
        def setter(self, flock):
            current = self.flocks[key]
            if set(current.required_names) <= set(flock.keys()):
                raise ValueError('The flock {} must have '
                                 'the following names (given {})'.format(key,
                                                                         set(current.required_names),
                                                                         set(flock.keys())))
            self.flocks[key] = flock

        return property(getter, setter)

    def create_backend_property(cls, key):
        def getter(self):
            return self.backends[key]
        def setter(self, backend):
            self.backends[key] = backend

        return property(getter, setter)

    def __init__(cls, classname, bases, class_dict):
        super().__init__(classname, bases, class_dict)

        cls.instruments = dict()
        cls.flocks = dict()
        cls.backends = dict()
        for key, value in class_dict.items():
            if isinstance(value, InstrumentSlot) or value is InstrumentSlot:
                if value is InstrumentSlot:
                    value = value()
                cls.instruments[key] = value
                cls.instruments[key]._slot_name = key
                setattr(cls, key, cls.create_instrument_property(key))
                LOGGER.debug('In {}, adding instrument named {} of type {}'.format(cls, key, value))

            elif isinstance(value, FlockSlot) or value is FlockSlot:
                if value is FlockSlot:
                    value = value()
                cls.flocks[key] = value
                cls.flocks[key]._slot_name = key
                setattr(cls, key, cls.create_flock_property(key))
                LOGGER.debug('In {}, adding flock named {} of type {}'.format(cls, key, value))

            elif value is BackendSlot:
                cls.backends[key] = value
                cls.backends[key]._slot_name = key
                setattr(cls, key, cls.create_backend_property(key))
                LOGGER.debug('In {}, adding backend named {} of type {}'.format(cls, key, value))

            else:
                pass
                #LOGGER.debug('In {}, unhandled attribute named {} = {}'.format(cls, key, value))

    def __str__(cls):
        return cls.__name__

class _FrontendType(MetaQObject):

    def create_frontend_property(cls, key):
        def getter(self):
            return self.frontends[key]
        def setter(self, frontend):
            self.frontends[key] = frontend

        return property(getter, setter)

    def __init__(cls, classname, bases, class_dict):
        super().__init__(classname, bases, class_dict)

        cls.frontends = dict()

        for key, value in class_dict.items():
            if isinstance(value, Front2Back) or hasattr(value, 'frontends'):
                cls.frontends[key] = value
                setattr(cls, key, cls.create_frontend_property(key))
                LOGGER.debug('{}, adding frontend named {} of type {}'.format(cls, key, value))

    def __str__(cls):
        return cls.__name__

class Frontend(LogMixin, ThreadLogMixin, QtGui.QMainWindow, metaclass=_FrontendType):

    logger_name = None
    _get_logger = get_logger

    frontends = {}

     # a declarative way to indicate the user interface file to use.
    gui = None

    # connect widgets to instruments using connect_setup
    auto_connect = False

    def __init__(self, parent=None, backend=None):
        super().__init__(parent)

        self._backend = None

        if self.logger_name is None:
            self.logger_name = 'lantz.qt.frontend.' + str(self)

        if self.gui:
            for cls in self.__class__.__mro__:
                if cls is object:
                    raise ValueError('{}: loading gui file {}, reached object parent'.format(self, self.gui))

                filename = os.path.dirname(inspect.getfile(cls))
                if isinstance(self.gui, tuple):
                    filename = os.path.join(filename, *self.gui)
                else:
                    filename = os.path.join(filename, self.gui)
                if os.path.exists(filename):
                    self.log_debug('loading gui file {} in Main Window'.format(filename))
                    self.widget = QtGui.loadUi(filename)
                    if isinstance(self.widget, QtGui.QMainWindow):
                        self.log_debug('reloading gui file {} AS Main Window'.format(filename))
                        self.widget = QtGui.loadUi(filename, self)
                    else:
                        self.setCentralWidget(self.widget)
                    break
            else:
                raise ValueError('{}: loading gui file {}'.format(self, self.gui))

        # Iterate over all frontend items in the current frontend
        # and instantiate each of them.
        # Note: a frontend declares which sub backend requires
        # but instantation is delegated to lantz

        for name, frontend in self.frontends.items():

            if isinstance(frontend, Front2Back):

                # The item is a frontend that requires a backend.

                cls = frontend.frontend_class
                if backend is None:
                    # If the current backend is None, then we cannot give
                    # anything to the sub frontend item
                    widget = cls()
                    self.log_debug('{} ({}) requires a backend but no backend defined'.format(name, cls))
                else:
                    # If the current backend exists then we get which of its attribute
                    # is the backend of the subfrontend
                    sub_backend_name = frontend.backend_name
                    if sub_backend_name is None:
                        # Is the same backend
                        widget = cls(backend=backend)
                        self.log_debug('{} ({}) connected to parent backend'.format(name, cls))
                    else:
                        # Is in an attribute
                        try:
                            sub_backend=getattr(backend, sub_backend_name)
                        except AttributeError:
                            raise ValueError("{} ({}) requires a '{}' attribute which is not provided by {}".format(name, cls, sub_backend_name, backend))
                        widget = cls(backend=sub_backend)
                        self.log_debug('{} ({}) connected to backend.{}'.format(name, cls, sub_backend_name))
            else:
                # This backend does not declare a required backend
                self.log_debug('{} ({}) created'.format(name, frontend))
                widget = frontend()

            widget.setParent(self)
            setattr(self, name, widget)

        self.setupUi()
        self.backend = backend

        self.log_current_thread()

    def __str__(self):
        return self.__class__.__name__

    def setupUi(self):
        pass

    def connect_backend(self):
        pass

    @property
    def backend(self):
        return self._backend

    @backend.setter
    def backend(self, backend):
        if self._backend is not None:
            self.log_debug('disconnecting backend: {}'.format(backend))
            self.disconnect(backend)

        self._backend = backend

        if backend is not None:
            self.log_debug('connecting backend: {}'.format(backend))
            if self.auto_connect:
                connect_setup(self.widget, backend.instruments.values())

            self.connect_backend()

    @classmethod
    def using(cls, backend_name):
        return Front2Back(cls, backend_name)

    @classmethod
    def using_parent_backend(cls):
        return Front2Back(cls, backend_name=None)

    def connect_feat(self, widget, target, feat_name=None, feat_key=MISSING):
        return connect_feat(widget, target, feat_name, feat_key)

class Backend(Base, ThreadLogMixin, SuperQObject, metaclass=_BackendType):

    _observer_signal_init = lambda: QtCore.pyqtSignal(object, object)

    backends = None
    instruments = None
    flocks = None

    def __init__(self, parent=None, **instruments_and_backends):

        # First we check that the names provided do not collide with parent classes
        invalid = set(dir(Base)) | set(dir(SuperQObject)) | set(dir(_BackendType))
        for name in instruments_and_backends.keys():
            if name in invalid:
                raise ValueError('{} is an invalid instrument or backend name'
                                 ' as it collides with attribute of parent class'.format(name))

        if self.logger_name is None:
            self.logger_name = 'lantz.qt.backend.' + str(self)

        Base.__init__(self, self.logger_name)
        super().__init__(parent)

        inst_keys = set(self.instruments.keys())
        flo_keys = set(self.flocks.keys())
        be_keys = set(self.backends.keys())

        for key, value in instruments_and_backends.items():
            if key in inst_keys:
                inst_keys.remove(key)
                setattr(self, key, value)
            elif key in flo_keys:
                flo_keys.remove(key)
                self.flocks[key] = value
            elif key in be_keys:
                be_keys.remove(key)
                self.backends[key] = value
            else:
                raise ValueError('Cannot link argument. No InstrumentSlot, BackendSlot or FlockSlot named {}'.format(key))

        if inst_keys:
            raise ValueError('No value provided for InstrumentSlot {}'.format(inst_keys))

        if be_keys:
            raise ValueError('No value provided for BackendSlot {}'.format(be_keys))

        self.log_current_thread()

    def __str__(self):
        return self.__class__.__name__

    def initialize(self, register_finalizer=False):
        initialize_many(self.instruments.values(), register_finalizer=register_finalizer)

    def finalize(self):
        finalize_many(self.instruments.values())

    def __enter__(self):
        self.initialize(register_finalizer=False)
        return self

    def __exit__(self, *args):
        self.finalize()

    def invoke(self, func, *args, **kwargs):
        QtCore.QTimer.singleShot(0, lambda: func(*args, **kwargs))

    def moveToThread(self, thread):
        if self.thread().objectName().startswith('pinned-'):
            return
        super().moveToThread(thread)
        for inst in self.instruments.values():
            inst.moveToThread(thread)
        for be in self.backends.values():
            be.moveToThread(thread)

def build_qapp(qapp_or_args=None, after_func=None):
    if isinstance(qapp_or_args, QtGui.QApplication):
        qapp = qapp_or_args
    else:
        qapp = QtGui.QApplication(qapp_or_args or [''])
        qapp.setWindowIcon(QtGui.QIcon(ICON_FEDORA))

    if after_func:
        qapp = after_func(qapp)

    return qapp

def start_gui_app(backend, frontend_class, qapp_or_args=None, after_qapp_creation=None):
    qapp = build_qapp(qapp_or_args, after_qapp_creation)

    QtCore.QThread.currentThread().setObjectName('main')

    background_thread = QtCore.QThread()
    background_thread.setObjectName('background')
    backend.moveToThread(background_thread)
    background_thread.start()

    frontend = frontend_class(backend=backend)
    frontend.show()

    qapp.aboutToQuit.connect(background_thread.quit)

    if sys.platform.startswith('darwin'):
        frontend.raise_()

    sys.exit(qapp.exec_())

def start_test_app(target, width=500, qapp_or_args=None, after_qapp_creation=None):
    """Start a single window test application with a form automatically
    generated for the driver.

    Parameters
    ----------
    target :
        a driver object or a collection of drivers.
    width :
        to be used as minimum width of the window. (Default value = 500)
    qapp_or_args :
        arguments to be passed to QApplication. (Default value = None)

    Returns
    -------

    """

    qapp = build_qapp(qapp_or_args, after_qapp_creation)

    if isinstance(target, (Driver, QDriver)):
        main = DriverTestWidget(None, target)
    else:
        main = SetupTestWidget(None, target)

    main.setMinimumWidth(width)
    main.setWindowTitle('Lantz Driver Test Panel')
    main.show()
    if sys.platform.startswith('darwin'):
        main.raise_()
    qapp.exec_()

def start_gui(ui_filename, drivers, qapp_or_args=None, after_qapp_creation=None):
    """Start a single window application with a form generated from
    a designer file.

    Parameters
    ----------
    ui_filename :
        the full path of a file generated with QtDesigner.
    drivers :
        a driver object or a collection of drivers.
    qapp_or_args :
        arguments to be passed to QApplication. (Default value = None)

    Returns
    -------

    """
    qapp = build_qapp(qapp_or_args, after_qapp_creation)

    main = QtGui.loadUi(ui_filename)

    if isinstance(drivers, Driver):
        connect_driver(main, drivers)
    else:
        connect_setup(main, drivers)

    main.show()
    if sys.platform.startswith('darwin'):
        main.raise_()
    qapp.exec_()

def start_frontend(frontend_class, qapp_or_args=None, after_qapp_creation=None):
    qapp = build_qapp(qapp_or_args, after_qapp_creation)

    frontend = frontend_class()
    frontend.show()

    if sys.platform.startswith('darwin'):
        frontend.raise_()

    sys.exit(qapp.exec_())


#HTML запрос через терминал при помощи файла пайтон

import requests

url = 'https://www.google.com/'
response = requests.get(url)
print(f'Requests to {url}. Status code is {response.status_code}')

# Requests to https://www.google.com/. Status code is 200 код выдве

# Можно вывести весь html текст на этой странице
print(response.text)

# https://earthquake.usgs.gov/fdsnws/event/1/query?
# format=geojson&
# starttime=2014-01-01&
# endtime=2014-01-02


from invoke import task

@task()
def clean(ctx):
	"""
	Delete 'build' folder.
	"""
	print("Cleaning!")
	ctx.run("rm -Rf build")

@task(
	help = {
		'cclean': "Call Clean task (Delete 'build' folder) before build again."
	}
)
def build(ctx, cclean=False):
	"""
	Build Fortran95 code.
	"""
	if cclean:
		clean(ctx)

	print("Building!")
	commands = [
		'mkdir build',
		'cd build',
		'f95 -c ../src/vector_math.f95 ../src/hello.f95',
		'f95 hello.o vector_math.o -o hello',
	]
	ctx.run(' && '.join(commands))

@task(
	help = {
		'cclean': "Call Clean task (Delete 'build' folder) before build again."
	}
)
def build_fc(ctx, cclean=False):
	"""
	Build C code that calls a Fortran95 module.
	"""
	if cclean:
		clean(ctx)

	print("Building!")
	commands = [
		'mkdir build',
		'cd build',
		'gcc -c ../src/callfortran.c',
		'f95 callfortran.o ../src/modulefort.f95',
	]
	ctx.run(' && '.join(commands))


import torch
from safetensors.torch import save_file

def load_and_save(model_path: str, output_path: str):
    model = torch.load(model_path)
    save_file(dict(model), output_path)

def save_files(
    encoder_path,
    encoder_output,
    vae_path,
    vae_output,
    unet_path,
    unet_output,
):
    load_and_save(encoder_path, encoder_output)
    load_and_save(vae_path, vae_output)
    load_and_save(unet_path, unet_output)

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--encoder_path", type=str, required=True)
    parser.add_argument("--encoder_output", type=str, required=True)
    parser.add_argument("--vae_path", type=str, required=True)
    parser.add_argument("--vae_output", type=str, required=True)
    parser.add_argument("--unet_path", type=str, required=True)
    parser.add_argument("--unet_output", type=str, required=True)

    args = parser.parse_args()

    save_files(
        args.encoder_path,
        args.encoder_output,
        args.vae_path,
        args.vae_output,
        args.unet_path,
        args.unet_output,
    )


# FIXME: make exceptions nicer
import untypy
import typing
import dataclasses
import inspect
import types

_DEBUG = False
def _debug(s):
    if _DEBUG:
        print('[DEBUG] ' + s)

# Types
Any = typing.Any
Optional = typing.Optional
Union = typing.Union
Literal = typing.Literal
Iterable = typing.Iterable
Iterator = typing.Iterator
Sequence = typing.Sequence
Generator = typing.Generator
ForwardRef = typing.ForwardRef
Protocol = typing.Protocol

Mapping = typing.Mapping

Callable = typing.Callable

dataclass = dataclasses.dataclass

unchecked = untypy.unchecked

intPositive = typing.Annotated[int, lambda i: i > 0, 'intPositive']
nat = typing.Annotated[int, lambda i: i >= 0, 'nat']
intNonNegative = typing.Annotated[int, lambda i: i >= 0, 'intNonNegative']
intNonPositive = typing.Annotated[int, lambda i: i <= 0, 'intNonPositive']
intNegative = typing.Annotated[int, lambda i: i < 0, 'intNegative']

floatPositive = typing.Annotated[float, lambda x: x > 0, 'floatPositive']
floatNonNegative = typing.Annotated[float, lambda x: x >= 0, 'floatNonNegative']
floatNegative = typing.Annotated[float, lambda x: x < 0, 'floatNegative']
floatNonPositive = typing.Annotated[float, lambda x: x <= 0, 'floatNonPositive']

class Lock(Protocol):
    def acquire(self, blocking: bool = True, timeout:int = -1) -> Any:
       pass

    def release(self) -> Any:
        pass

    def locked(self) -> Any:
        pass

LockFactory = typing.Annotated[Callable[[], Lock], 'LockFactory']

T = typing.TypeVar('T')
U = typing.TypeVar('U')
V = typing.TypeVar('V')

def _literalInstanceOf(self, value):
    for p in self.__args__:
        # typing.Literal checks for type(value) == type(arg),
        # however this does not work well with untypy.
        if isinstance(value, type(p)) and p == value:
            return True
    return False

def _invalidCall(self, *args, **kwds):
    argStr = ', '.join([untypy.util.typehints.qualname(x) for x in args])
    if hasattr(self, '__name__'):
        name = self.__name__
    else:
        name = str(self)
        typingPrefix = 'typing.'
        if name.startswith(typingPrefix):
            name = name[len(typingPrefix):]
    raise untypy.error.WyppTypeError(f"Cannot instantiate {name}. Did you mean {name}[{argStr}]?")

# Dirty hack ahead: we patch some methods of internal class of the typing module.

# This patch is needed to be able to use a literal L to check whether a value x
# is an instance of this literal: instanceof(x, L)
setattr(typing._LiteralGenericAlias, '__instancecheck__', _literalInstanceOf)

# This patch is needed to provide better error messages if a student passes type arguments
# with paranthesis instead of square brackets
setattr(typing._SpecialForm, '__call__', _invalidCall)

def _collectDataClassAttributes(cls):
    result = dict()
    for c in cls.mro():
        if hasattr(c, '__kind') and c.__kind == 'record' and hasattr(c, '__annotations__'):
            result = c.__annotations__ | result
    return result

def _patchDataClass(cls, mutable):
    fieldNames = [f.name for f in dataclasses.fields(cls)]
    setattr(cls, EQ_ATTRS_ATTR, fieldNames)

    if hasattr(cls, '__annotations__'):
        # add annotions for type checked constructor.
        cls.__kind = 'record'
        cls.__init__.__annotations__ = _collectDataClassAttributes(cls)
        cls.__init__.__original = cls # mark class as source of annotation
        cls.__init__ = untypy.typechecked(cls.__init__)

    if mutable:
        # prevent new fields being added
        fields = set(fieldNames)

        checker = {}
        # Note: Partial annotations are disallowed by untypy.typechecked(cls.__init__)
        #       So no handling in this code is required.
        for name in fields:
            if name in cls.__annotations__:
                # This the type is wrapped in an lambda expression to allow for Forward Ref.
                # Would the lambda expression be called at this moment, it may cause an name error
                # untypy.checker fetches the annotation lazily.
                checker[name] = untypy.checker(\
                    lambda cls=cls,name=name: typing.get_type_hints(cls, include_extras=True)[name],
                    cls)

        oldSetattr = cls.__setattr__
        def _setattr(obj, k, v):
            # Note __setattr__ also gets called in the constructor.
            if k in checker:
                oldSetattr(obj, k, checker[k](v))
            elif k in fields:
                oldSetattr(obj, k, v)
            else:
                raise untypy.error.WyppAttributeError(f'Unknown attribute {k} for record {cls.__name__}')
        setattr(cls, "__setattr__", _setattr)
    return cls

def record(cls=None, mutable=False):
    def wrap(cls):
        newCls = dataclasses.dataclass(cls, frozen=not mutable)
        return _patchDataClass(newCls, mutable)
    # See if we're being called as @record or @record().
    if cls is None:
        # We're called with parens.
        return wrap
    else:
        # We're called as @dataclass without parens.
        return wrap(cls)

# Tests

_die = False

def setDieOnCheckFailures(b):
    _die = b

def _dieOnCheckFailures():
    return _die

_testCount = {'total': 0, 'failing': 0}
_checksEnabled = True

def initModule(enableChecks=True, quiet=False):
    global _checksEnabled
    _checksEnabled = enableChecks
    resetTestCount()

def resetTestCount():
    global _testCount
    _testCount = {'total': 0, 'failing': 0}

def printTestResults(prefix=''):
    total = _testCount['total']
    failing = _testCount['failing']
    if total == 0:
        pass
    elif failing == 0:
        print(f'{prefix}{total} Tests, alle erfolgreich :-)')
    else:
        print(f'{prefix}{total} Tests, {failing} Fehler :-(')
    return {'total': total, 'failing': failing}

def checkEq(actual, expected):
    return check(actual, expected, structuralObjEq=False)

def check(actual, expected, *, structuralObjEq=True, floatEqWithDelta=True):
    if not _checksEnabled:
        return
    global _testCount
    flags = {'structuralObjEq': structuralObjEq, 'floatEqWithDelta': floatEqWithDelta}
    matches = deepEq(actual, expected, **flags)
    _testCount = {
        'total': _testCount['total'] + 1,
        'failing': _testCount['failing'] + (0 if matches else 1)
    }
    if not matches:
        stack = inspect.stack()
        caller = stack[1] if len(stack) > 1 else None
        def fmt(x):
            if type(x) == str:
                return repr(x)
            else:
                return str(x)
        msg = f"{caller.filename}:{caller.lineno}: Erwartet wird {fmt(expected)}, aber das " \
            f"Ergebnis ist {fmt(actual)}"
        if _dieOnCheckFailures():
            raise Exception(msg)
        else:
            print("FEHLER in " + msg)

def uncoveredCase():
    stack = inspect.stack()
    caller = stack[1] if len(stack) > 1 else None
    raise Exception(f"{caller.filename}, Zeile {caller.lineno}: ein Fall ist nicht abgedeckt")

#
# Deep equality
#
def _isNumber(x):
    t = type(x)
    return (t is int or t is float)

def _seqEq(seq1, seq2, flags):
    if len(seq1) != len(seq2):
        return False
    for i, x in enumerate(seq1):
        y = seq2[i]
        if not deepEq(x, y, **flags):
            return False
    return True

def _dictEq(d1, d2, flags):
    ks1 = set(d1.keys())
    ks2 = set(d2.keys())
    if ks1 != ks2: # keys should be exactly equal
        return False
    for k in ks1:
        if not deepEq(d1[k], d2[k], **flags):
            return False
    return True

def _objToDict(o):
    d = {}
    attrs = dir(o)
    useAll = False
    if hasattr(o, EQ_ATTRS_ATTR):
        useAll = True
        attrs = getattr(o, EQ_ATTRS_ATTR)
    for n in attrs:
        if not useAll and n and n[0] == '_':
            continue
        x = getattr(o, n)
        d[n] = x
    #print(f'_objToDict: {d}')
    return d

def _objEq(o1, o2, flags):
    return _dictEq(_objToDict(o1), _objToDict(o2), flags)

_EPSILON = 0.00001
EQ_ATTRS_ATTR = '__eqAttrs__'

def _useFloatEqWithDelta(flags):
    return flags.get('floatEqWithDelta', False)

def _useStructuralObjEq(flags):
    return flags.get('structuralObjEq', False)

# Supported flags:
# - structuralObjEq: causes objects to be compared as dictionaries. The keys are by default
#   taken from dir(obj), if obj as the attribute __eqAttrs__, then the names listed there
#   are taken. Defaults to False.
# - floatEqWithDelta: compares floats by checking whether the difference is smaller than a
#   small delta. Setting this to True loses transitivity of eq.
def deepEq(v1, v2, **flags):
    """
    Computes deep equality of v1 and v2. With structuralObjEq=False, objects are compared
    by __eq__. Otherwise, objects are compared attribute-wise, only those attributes
    returned by dir that do not start with an underscore are compared.
    """
    #print(f'deepEq: v1={v1}, v2={v2}, flags={flags}')
    if v1 == v2:
        return True
    if _isNumber(v1) and _isNumber(v2) and _useFloatEqWithDelta(flags):
        diff = v1 - v2
        if abs(diff) < _EPSILON:
            return True
        else:
            return False
    if isinstance(v1, list):
        return isinstance(v2, list) and _seqEq(v1, v2, flags)
    if isinstance(v1, tuple):
        return isinstance(v2, tuple) and _seqEq(v1, v2, flags)
    if isinstance(v1, dict):
        return isinstance(v2, dict) and _dictEq(v1, v2, flags)
    if callable(v1):
        return False
    if type(v1) == str:
        return False
    if hasattr(v1, '__class__'):
        if _useStructuralObjEq(flags):
            res = _objEq(v1, v2, flags)
            return res
        else:
            return False # v1 == v2 already checked
    return False

class TodoError(Exception, untypy.error.DeliberateError):
    pass

def todo(msg=None):
    if msg is None:
        msg = 'TODO'
    raise TodoError(msg)

class ImpossibleError(Exception, untypy.error.DeliberateError):
    pass

def impossible(msg=None):
    if msg is None:
        msg = 'the impossible happened'
    raise ImpossibleError(msg)

# Additional functions and aliases

import math as moduleMath

math = moduleMath


from functools import wraps
import jwt
from flask import jsonify, request, current_app as app
from auth.daos import user_dao
from auth.services import blacklist_token_service

def token_required(f):
    """
    @param f:
    @return:
    """
    @wraps(f)
    def decorator(*args, **kwargs):

        token = None

        if 'Authorization' in request.headers:
            token = request.headers['Authorization']

        if not token:
            return jsonify({'message': 'Invalid Token'})

        if blacklist_token_service.in_blacklist(token):
            return jsonify({'message': 'Token is in blacklist'})

        try:
            data = jwt.decode(token, app.config["SECRET_KEY"])
            current_user = user_dao.get_user_by_username(data['username'])
        except Exception as e:
            return jsonify({'message': str(e)})

        return f(current_user, *args, **kwargs)

    return decorator


import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import numpy as np
import datetime
import matplotlib.pyplot as plt
import torchvision.transforms as T
import car_racing as cr
import replay_memory as rm
import data_collector as dc
import random

class VAE(nn.Module):
    # In part taken from:
    # https://github.com/pytorch/examples/blob/master/vae/main.py

    def __init__(self, n_screens, n_latent_states, lr=1e-5, device='cpu'):
        super(VAE, self).__init__()

        self.device = device

        self.n_screens = n_screens
        self.n_latent_states = n_latent_states

        # The convolutional encoder
        self.encoder = nn.Sequential(
                nn.Conv2d(self.n_screens, 32, 4, 2), # (1, 8, 42, 42) --> (1, 32, 20, 20)
                nn.BatchNorm2d(32),
                nn.ReLU(inplace=True),

                nn.Conv2d(32, 64, 4, 2), # (1, 32, 20, 20) --> (1, 64, 9, 9)
                nn.BatchNorm2d(64),
                nn.ReLU(inplace=True),

                nn.Conv2d(64, 128, 5, 2), # (1, 64, 9, 9) --> (1, 128, 3, 3)
                nn.BatchNorm2d(128),
                nn.ReLU(inplace=True),

                nn.Conv2d(128, 256, 3, 2), # (1, 128, 3, 3) --> (1, 256, 1, 1)
                nn.ReLU(inplace=True),

                ).to(self.device)

        # The size of the encoder output
        self.encoder_output_shape = (256, 1, 1)
        self.encoder_output_size = np.prod(self.encoder_output_shape)

        # The convolutional decoder
        self.decoder = nn.Sequential(
                nn.ConvTranspose2d(256, 128, 3, 2), # (1, 256, 1, 1) --> (1, 128, 3, 3)
                nn.BatchNorm2d(128),
                nn.ReLU(inplace=True),

                nn.ConvTranspose2d(128, 64, 5, 2), # (1, 128, 3, 3) --> (1, 64, 9, 9)
                nn.BatchNorm2d(64),
                nn.ReLU(inplace=True),

                nn.ConvTranspose2d(64, 32, 4, 2), # (1, 64, 9, 9) --> (1, 32, 20, 20)
                nn.BatchNorm2d(32),
                nn.ReLU(inplace=True),

                nn.ConvTranspose2d(32, self.n_screens, 4, 2), # (1, 32, 20, 20) --> (1, n_screens, 42, 42)
                nn.BatchNorm2d(self.n_screens),
                nn.ReLU(inplace=True),

                nn.Sigmoid()
                ).to(self.device)

        # Fully connected layers connected to encoder
        self.fc1 = nn.Linear(self.encoder_output_size, self.encoder_output_size // 2) # 1024 --> 512
        self.fc2_mu = nn.Linear(self.encoder_output_size // 2, self.n_latent_states) # 512 --> 128
        self.fc2_logvar = nn.Linear(self.encoder_output_size // 2, self.n_latent_states) # 512 --> 128

        # Fully connected layers connected to decoder
        self.fc3 = nn.Linear(self.n_latent_states, self.encoder_output_size // 2) # 128 --> 512
        self.fc4 = nn.Linear(self.encoder_output_size // 2, self.encoder_output_size) # 512 --> 1024

        self.optimizer = optim.Adam(self.parameters(), lr)

        self.to(self.device)

    def encode(self, x):
        # Deconstruct input x into a distribution over latent states
        conv = self.encoder(x)
        h1 = F.relu(self.fc1(conv.view(conv.size(0), -1)))
        mu, logvar = self.fc2_mu(h1), self.fc2_logvar(h1)
        return mu, logvar

    def reparameterize(self, mu, logvar):
        # Apply reparameterization trick
        std = torch.exp(0.5*logvar)
        eps = torch.randn_like(std)
        return mu + eps*std

    def decode(self, z, batch_size=1):
        # Reconstruct original input x from the (reparameterized) latent states
        h3 = F.relu(self.fc3(z))
        deconv_input = self.fc4(h3)
        deconv_input = deconv_input.view([batch_size] + [dim for dim in self.encoder_output_shape])
        y = self.decoder(deconv_input)
        return y

    def forward(self, x, batch_size=1):
        # Deconstruct and then reconstruct input x
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        recon = self.decode(z, batch_size)
        return recon, mu, logvar

    # Reconstruction + KL divergence losses summed over all elements and batch
    def loss_function(self, recon_x, x, mu, logvar, batch=True):
        if batch:
            BCE = F.binary_cross_entropy(recon_x, x, reduction='none')
            BCE = torch.sum(BCE, dim=(1, 2, 3))

            KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp(), dim=1)
        else:
            BCE = F.binary_cross_entropy(recon_x, x, reduction='sum')
            # see Appendix B from VAE paper:
            # Kingma and Welling. Auto-Encoding Variational Bayes. ICLR, 2014
            # https://arxiv.org/abs/1312.6114
            # 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
            KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())

        return BCE + KLD

class Model(nn.Module):

   def __init__(self, n_inputs, n_outputs, n_hidden=64, lr=1e-3, softmax = False, device='cpu'):
        super(Model, self).__init__()

        self.n_inputs = n_inputs   # Number of inputs
        self.n_hidden = n_hidden   # Number of hidden units
        self.n_outputs = n_outputs # Number of outputs
        self.softmax = softmax

        self.fc1 = nn.Linear(self.n_inputs, self.n_hidden)  # Hidden layer
        self.fc2 = nn.Linear(self.n_hidden, self.n_outputs) # Output layer

        self.optimizer = optim.Adam(self.parameters(), lr)  # Adam optimizer

        self.device = device
        self.to(self.device)

   def forward(self, x):
        # Define the forward pass:
        x = x.to(self.device)
        h1 = F.relu(self.fc1(x))

        if self.softmax: # If true apply a softmax function to the output
            y = F.softmax(self.fc2(h1), dim=-1).clamp(min=1e-9, max=1-1e-9) # This is used to get a proper distribution over all actions with a sum of 1.
        else:
            y = self.fc2(h1)

        return y

class Agent():

    def __init__(self, device = 'cuda'):

        self.run_id = 1
        self.device = device
        self.env = cr.CarRacing()
        self.render_view = False      # Set to True if you want to see what it is doing
        self.print_timer = 10         # Print average result of Agent every '...' episodes

        self.height = self.width = 42 # observation size (height and width)
        self.color = 1                # number of colors
        self.n_screens = 8            # number of observations stacked

        self.obs_shape = (self.height, self.width)
        self.obs_size = int(np.prod(self.obs_shape))     # The size of the observation
        self.linear = False                              # True if the input is a vector

        # Initialize last observations array
        self.obs_batch = np.array([np.zeros((self.height, self.width), dtype = 'float32') for i in range(self.n_screens)])

        # Discretization of continuous action space for CarRacing-v0
        # [0] = steering, [1] = accelerating, [2] = braking
        self.discrete_actions = {0 : np.array([0,0,0]),       # do nothing
                                 1 : np.array([-1,0,0]),      # steer sharp left
                                 2 : np.array([1,0,0]),       # steer sharp right
                                 3 : np.array([-0.5,0,0]),    # steer left
                                 4 : np.array([0.5,0,0]),     # steer right
                                 5 : np.array([0,1,0]),       # accelerate 100%
                                 6 : np.array([0,0.5,0]),     # accelerate 50%
                                 7 : np.array([0,0.25,0]),    # accelerate 25%
                                 8 : np.array([0,0,1]),       # brake 100%
                                 9 : np.array([0,0,0.5]),     # brake 50%
                                 10 : np.array([0,0,0.25])}   # brake 25%

        # The number of actions available to the agent
        self.n_actions = len(self.discrete_actions)

        self.freeze_cntr = 0           # Keeps track of when to (un)freeze the target network
        self.freeze_period = 50        # How long the network is frozen
        self.batch_size = 250
        self.freeze_vae = True

        self.memory_capacity = 100000     # memory size
        self.VAE_memory_capacity = 10000  # VAE pre-train memory size
        self.n_episodes = 1000            # number of episodes
        self.n_play_episodes = 150        # number of episodes used for average reward test
        self.max_length_episode = 1000    # max number of steps a episode of the CarRacing environment lasts

        self.gamma = 12                   # Precision parameter
        self.Beta = 0.99                  # Discount factor
        self.alpha = 18000                # VAE loss scaler

        self.n_hidden_trans = 512         # number of hidden units transition network
        self.lr_trans = 1e-3              # learning rate transition network
        self.n_hidden_pol = 512           # number of hidden units policy network
        self.lr_pol = 1e-4                # learning rate policy network
        self.n_hidden_val = 512           # number of hidden units value network
        self.lr_val = 1e-5                # learning rate value network

        self.n_latent_states = 128        # size latent space
        self.lr_vae = 5e-6                # learning rate VAE
        self.vae_data = 'pre_train_data/vae_data_10000.pt'
        self.vae_plot = False
        self.pre_train_vae = True         # if True pre-trains the VAE

        self.load_pre_trained_vae = True
        self.pt_vae_load_path = "networks/pre_trained_vae/vae_daif_CarRacing_{}_end.pth".format(self.n_latent_states)

        self.load_network = True
        self.network_load_path = "networks/daif/daif_CarRacing_{}net_r{}.pth".format("{}", self.run_id)

        # Initialize the networks:
        self.vae = VAE(self.n_screens, self.n_latent_states, lr = self.lr_vae, device=self.device)
        self.transition_net = Model(self.n_latent_states*2 + 1, self.n_latent_states, self.n_hidden_trans, lr=self.lr_trans, device=self.device)  # + 1, for 1 action
        self.policy_net = Model(self.n_latent_states*2, self.n_actions, self.n_hidden_pol, lr=self.lr_pol, softmax=True, device=self.device)
        self.value_net = Model(self.n_latent_states*2, self.n_actions, self.n_hidden_val, lr=self.lr_val, device=self.device)
        self.target_net = Model(self.n_latent_states*2, self.n_actions, self.n_hidden_val, lr=self.lr_val, device=self.device)
        self.target_net.load_state_dict(self.value_net.state_dict())

        if self.load_pre_trained_vae: # If true: load a pre-trained VAE
            self.vae.load_state_dict(torch.load(self.pt_vae_load_path, map_location=self.device))
            self.vae.eval()
            print("Succesfully loaded a pre-trained VAE")

        if self.load_network: # If true: load the networks given paths
            self.vae.load_state_dict(torch.load(self.network_load_path.format("vae"), map_location=self.device))
            self.vae.eval()
            self.transition_net.load_state_dict(torch.load(self.network_load_path.format("trans"), map_location=self.device))
            self.transition_net.eval()
            self.policy_net.load_state_dict(torch.load(self.network_load_path.format("pol"), map_location=self.device))
            self.policy_net.eval()
            self.value_net.load_state_dict(torch.load(self.network_load_path.format("val"), map_location=self.device))
            self.value_net.eval()
            print("Succesfully loaded networks")

        # Initialize the replay memory
        self.memory = rm.ReplayMemory(self.memory_capacity, self.obs_shape, self.obs_size, self.linear, device=self.device)
        if self.pre_train_vae:
            self.VAE_memory = rm.ReplayMemory(self.VAE_memory_capacity, self.obs_shape, self.obs_size, self.linear, device=self.device)

        # When sampling from memory at index i, obs_indices indicates that we want observations with indices i-obs_indices, works the same for the others
        self.obs_indices = [(self.n_screens+1)-i for i in range(self.n_screens+2)]
        self.action_indices = [2, 1]
        self.reward_indices = [1]
        self.done_indices = [0]
        self.max_n_indices = max(max(self.obs_indices, self.action_indices, self.reward_indices, self.done_indices)) + 1

        # Used to pre-process the observations (screens)
        self.preprocess = T.Compose([T.ToPILImage(),
                    T.Grayscale(num_output_channels=1),
                    T.Resize((self.height, self.width)),
                    T.ToTensor()])

        self.save_results = True
        self.save_network = True
        self.results_path = "results/daif/daif_CarRacing_results_r{}.npz".format(self.run_id)
        self.network_save_path = "networks/daif/daif_CarRacing_{}net_r{}.pth".format("{}",self.run_id)

        self.log_path = "logs/daif_CarRacing_log_r{}.txt".format(self.run_id)
        self.record = open(self.log_path, "a")
        self.record.write("\n\n-----------------------------------------------------------------\n")
        self.record.write("File opened at {}\n".format(datetime.datetime.now()))

    def get_screen(self, device='cuda'):
        # Get observation, reshape and but in right order.
        screen = self.env.render(mode='state_pixels')
        screen = screen.reshape(96, 96, 3)
        screen = screen.transpose((2, 1, 0))

        # stips of bottom part of the image which contains a black bar with the accumulated reward and control value bars, and makes sure the width is equal size as height
        screen = screen[:, 6:90, int(96*0):int(96 * 0.875)]

        # Convert to to float and normalize
        screen = np.ascontiguousarray(screen, dtype=np.float32) / 255

        # Add resize
        screen = self.preprocess(torch.from_numpy(np.flip(screen, axis=0).copy()))

        return screen

    """ stack the X latest observations into one batch """
    def get_obs_batch(self, obs):
        # add new observation to obs_batch, remove oldest.
        self.obs_batch = np.concatenate((obs.numpy(), self.obs_batch[0:self.n_screens-1]), axis = 0)

        # resize to (1, self.n_screens, 84, 84) and convert to torch
        obs_batch2 = torch.from_numpy(np.flip(self.obs_batch, axis=0).copy()).unsqueeze(0).to(self.device)

        return obs_batch2

    def select_action(self, obs):
        with torch.no_grad():
            action_index = 0

            if self.memory.push_count < self.batch_size + self.n_screens:
                action_index = random.randint(0, self.n_actions - 1)
            else:
                # Derive a distribution over states state from the last n observations (screens):
                prev_n_obs = self.get_obs_batch(obs)
                state_mu, state_logvar = self.vae.encode(prev_n_obs)
                x = torch.cat((state_mu, torch.exp(state_logvar)), dim=1) # does not work?
                policy = self.policy_net(x)
                action_index = torch.multinomial(policy, 1).item()

            return action_index

    def get_mini_batches(self):
        # Retrieve transition data in mini batches
        all_obs_batch, all_actions_batch, reward_batch_t1, done_batch_t2 = self.memory.sample(
                self.obs_indices, self.action_indices, self.reward_indices,
                self.done_indices, self.max_n_indices, self.batch_size)

        # Retrieve a batch of observations for 3 consecutive points in time
        obs_batch_t0 = all_obs_batch[:, 0:self.n_screens, :, :]
        obs_batch_t1 = all_obs_batch[:, 1:self.n_screens+1, :, :]
        obs_batch_t2 = all_obs_batch[:, 2:self.n_screens+2, :, :]

        # Retrieve a batch of distributions over states for 3 consecutive points in time
        state_mu_batch_t0, state_logvar_batch_t0 = self.vae.encode(obs_batch_t0)
        state_mu_batch_t1, state_logvar_batch_t1 = self.vae.encode(obs_batch_t1)
        state_mu_batch_t2, state_logvar_batch_t2 = self.vae.encode(obs_batch_t2)

        # Combine the sufficient statistics (mean and variance) into a single vector
        state_batch_t0 = torch.cat((state_mu_batch_t0, torch.exp(state_logvar_batch_t0)), dim=1)
        state_batch_t1 = torch.cat((state_mu_batch_t1, torch.exp(state_logvar_batch_t1)), dim=1)
        state_batch_t2 = torch.cat((state_mu_batch_t2, torch.exp(state_logvar_batch_t2)), dim=1)

        # Reparameterize the distribution over states for time t1
        z_batch_t1 = self.vae.reparameterize(state_mu_batch_t1, state_logvar_batch_t1)

        # Retrieve the agent's action history for time t0 and time t1
        action_batch_t0 = all_actions_batch[:, 0].unsqueeze(1)
        action_batch_t1 = all_actions_batch[:, 1].unsqueeze(1)

        # At time t0 predict the state at time t1:
        X = torch.cat((state_batch_t0.detach(), action_batch_t0.float()), dim=1)
        pred_batch_t0t1 = self.transition_net(X)

        # Determine the prediction error wrt time t0-t1:
        pred_error_batch_t0t1 = torch.mean(F.mse_loss(
                pred_batch_t0t1, state_mu_batch_t1, reduction='none'), dim=1).unsqueeze(1)

        return (state_batch_t1, state_batch_t2, action_batch_t1,
                reward_batch_t1, done_batch_t2, pred_error_batch_t0t1,
                obs_batch_t1, state_mu_batch_t1,
                state_logvar_batch_t1, z_batch_t1)

    def compute_value_net_loss(self, state_batch_t1, state_batch_t2,
                           action_batch_t1, reward_batch_t1,
                           done_batch_t2, pred_error_batch_t0t1):

        with torch.no_grad():
            # Determine the action distribution for time t2:
            policy_batch_t2 = self.policy_net(state_batch_t2)

            # Determine the target EFEs for time t2:
            target_EFEs_batch_t2 = self.target_net(state_batch_t2)

            # Weigh the target EFEs according to the action distribution:
            weighted_targets = ((1-done_batch_t2) * policy_batch_t2 *
                                target_EFEs_batch_t2).sum(-1).unsqueeze(1)

            # Determine the batch of bootstrapped estimates of the EFEs:
            EFE_estimate_batch = -reward_batch_t1 + pred_error_batch_t0t1 + self.Beta * weighted_targets

        # Determine the EFE at time t1 according to the value network:
        EFE_batch_t1 = self.value_net(state_batch_t1).gather(1, action_batch_t1)

        # Determine the MSE loss between the EFE estimates and the value network output:
        value_net_loss = F.mse_loss(EFE_estimate_batch, EFE_batch_t1)

        return value_net_loss

    def compute_VFE(self, vae_loss, state_batch_t1, pred_error_batch_t0t1):
        # Determine the action distribution for time t1:
        policy_batch_t1 = self.policy_net(state_batch_t1)

        # Determine the EFEs for time t1:
        EFEs_batch_t1 = self.value_net(state_batch_t1)

        # Take a gamma-weighted Boltzmann distribution over the EFEs:
        boltzmann_EFEs_batch_t1 = torch.softmax(-self.gamma * EFEs_batch_t1, dim=1).clamp(min=1e-9, max=1-1e-9)

        # Weigh them according to the action distribution:
        energy_term_batch = -(policy_batch_t1 * torch.log(boltzmann_EFEs_batch_t1)).sum(-1).unsqueeze(1)

        # Determine the entropy of the action distribution
        entropy_batch = -(policy_batch_t1 * torch.log(policy_batch_t1)).sum(-1).unsqueeze(1)

        # Determine the VFE, then take the mean over all batch samples:
        VFE_batch = vae_loss + pred_error_batch_t0t1 + (energy_term_batch - entropy_batch)
        VFE = torch.mean(VFE_batch)

        return VFE

    def learn(self, ith_episode):

        # If there are not enough transitions stored in memory, return
        if self.memory.push_count - self.max_n_indices*2 < self.batch_size:
            return

        # After every freeze_period time steps, update the target network
        if self.freeze_cntr % self.freeze_period == 0:
            self.target_net.load_state_dict(self.value_net.state_dict())
        self.freeze_cntr += 1

        # Retrieve mini-batches of data from memory
        (state_batch_t1, state_batch_t2, action_batch_t1,
        reward_batch_t1, done_batch_t2, pred_error_batch_t0t1,
        obs_batch_t1, state_mu_batch_t1,
        state_logvar_batch_t1, z_batch_t1) = self.get_mini_batches()

        # Determine the reconstruction loss for time t1
        recon_batch = self.vae.decode(z_batch_t1, self.batch_size)

        # Determine the VAE loss for time t1
        vae_loss = self.vae.loss_function(recon_batch, obs_batch_t1, state_mu_batch_t1, state_logvar_batch_t1, batch=True) / self.alpha

        # Compute the value network loss:
        value_net_loss = self.compute_value_net_loss(state_batch_t1, state_batch_t2,
                           action_batch_t1, reward_batch_t1,
                           done_batch_t2, pred_error_batch_t0t1)

        # Compute the variational free energy:
        VFE = self.compute_VFE(vae_loss, state_batch_t1.detach(), pred_error_batch_t0t1)

        # Reset the gradients:
        if not self.freeze_vae:
            self.vae.optimizer.zero_grad()
        self.policy_net.optimizer.zero_grad()
        self.transition_net.optimizer.zero_grad()
        self.value_net.optimizer.zero_grad()

        # Compute the gradients:
        VFE.backward(retain_graph=True)
        value_net_loss.backward()

        # Perform gradient descent:
        if not self.freeze_vae:
            self.vae.optimizer.step()
        self.policy_net.optimizer.step()
        self.transition_net.optimizer.step()
        self.value_net.optimizer.step()

    ''' Run a trained model without it learning. '''
    def play(self):

        rewards = []

        self.memory.push_count = self.memory_capacity - 1

        for ith_episode in range(self.n_play_episodes):

            total_reward = 0
            nr_steps = 0
            obs = self.env.reset()
            obs = self.get_screen(self.device)
            done = False

            while not done and nr_steps <= self.max_length_episode:

                # get action
                action = self.select_action(obs)

                # get actual action from discrete actions dictionary
                action_todo = self.discrete_actions.get(int(action))

                # take step
                obs, reward, done, _ = self.env.step([action_todo[0], action_todo[1], action_todo[2]])
                nr_steps = nr_steps + 1
                obs = self.get_screen(self.device)

                # render in visible window if True
                if self.render_view:
                    self.env.render('human')

                # add reward to total
                total_reward += reward

            rewards.append(total_reward)
            print("Reward for this episode:", total_reward)
            total_reward = 0

        self.env.close()

        np.savez("rewards/daif_CarRacing_rewards", np.array(rewards))

    def train_vae(self):
        """ Train the VAE using data collected via user play. """
        vae_batch_size = 256
        vae_obs_indices = [self.n_screens-i for i in range(self.n_screens)]

        self.VAE_memory.push_count = vae_batch_size + self.n_screens*2

        try:
            # Load the pre-collected data into device
            self.VAE_memory.obs_mem = torch.load(self.vae_data, map_location = torch.device(self.device))
        except:
            # Generate data to train VAE on
            print("No data found to train the vae on.")
            data_collector = dc.DataCollector(self.VAE_memory_capacity, self.n_screens, self.height, self.width, device = self.device)
            data_collector.generate_data()
            self.VAE_memory.obs_mem = torch.load(data_collector.obs_data_path, map_location = torch.device(self.device))

        losses = []

        for data_point in range(0, len(self.VAE_memory.obs_mem)):
            self.VAE_memory.push_count = self.VAE_memory.push_count + 1

            obs_batch, _, _, _ = self.VAE_memory.sample(vae_obs_indices, [], [], [], len(vae_obs_indices), vae_batch_size)
            obs_batch = obs_batch.view(vae_batch_size, self.n_screens, self.height, self.width)

            recon, mu, logvar = self.vae.forward(obs_batch, vae_batch_size)
            loss = torch.mean(self.vae.loss_function(recon, obs_batch, mu, logvar))

            self.vae.optimizer.zero_grad()
            loss.backward()
            self.vae.optimizer.step()

            losses.append(loss)
            if data_point % 50 == 0:
                print("obs: %2f vae loss=%5.2f"%(data_point, loss.item()))

            if data_point % 1000 == 0 and data_point > 0 and self.vae_plot:
                plt.plot(losses)
                plt.show()
                plt.plot(losses[-1000:])
                plt.show()

                for i in range(self.n_screens):
                    plt.imsave("vae_images/obs/vae_obs_CarRacing_ep_{}_{}.png".format(data_point, i), obs_batch[0, i, :, :].detach().cpu().squeeze(0).permute(1, 0).numpy(), cmap='gray')
                    plt.imsave("vae_images/recon/vae_recon_CarRacin_ep_{}_{}.png".format(data_point, i), recon[0, i, :, :].detach().cpu().squeeze(0).permute(1, 0).numpy(), cmap='gray')

        self.VAE_memory.push_count = 0
        torch.save(self.vae.state_dict(), "networks/pre_trained_vae/vae_daif_CarRacing_{}_test_end.pth".format(self.n_latent_states))

    def train(self):

        if self.pre_train_vae and not self.load_pre_trained_vae: # If True: pre-train the VAE
            msg = "Environment is: {}\nPre-training vae. Starting at {}".format("CarRacing-v0", datetime.datetime.now())
            print(msg)
            self.record.write(msg+"\n")
            self.train_vae()

        msg = "Environment is: {}\nTraining started at {}".format("CarRacing-v0", datetime.datetime.now())
        print(msg)
        self.record.write(msg+"\n")

        results = []
        for ith_episode in range(self.n_episodes):

            total_reward = 0
            self.env.reset()
            obs = self.get_screen(self.device)
            done = False
            reward = 0
            nr_steps = 0

            self.prev_screen = self.env.render('rgb_array')

            while not done and nr_steps <= self.max_length_episode:

                # get action
                action = self.select_action(obs)

                # push to memory
                self.memory.push(obs, action, reward, done)

                # get actual action input
                action_real = self.discrete_actions.get(int(action))

                # take step
                obs, reward, done, _ = self.env.step([action_real[0], action_real[1], action_real[2]])
                obs = self.get_screen(self.device)
                nr_steps = nr_steps + 1

                # render in visible window if True
                if self.render_view:
                    self.env.render('human')

                # add reward to total
                total_reward += reward

                # have the networks learn
                self.learn(ith_episode)

                if done or nr_steps == self.max_length_episode:
                    self.memory.push(obs, -99, -99, True)

            results.append(total_reward)

            # Print and keep a (.txt) record of stuff
            if ith_episode > 0 and ith_episode % self.print_timer == 0:
                avg_reward = np.mean(results)
                last_x = np.mean(results[-self.print_timer:])
                msg = "Episodes: {:4d}, avg score: {:3.2f}, over last {:d}: {:3.2f}".format(ith_episode, avg_reward, self.print_timer, last_x)
                print(msg)

                # write to log
                self.record.write(msg+"\n")

                # save log
                self.record.close()
                self.record = open(self.log_path, "a")

        self.env.close()

        # If enabled, save the results and the network (state_dict)
        if self.save_results:
            np.savez(self.results_path, np.array(results))
        if self.save_network:
            torch.save(self.transition_net.state_dict(), self.network_save_path.format("trans"))
            torch.save(self.policy_net.state_dict(), self.network_save_path.format("pol"))
            torch.save(self.value_net.state_dict(), self.network_save_path.format("val"))
            torch.save(self.VAE.state_dict(), self.network_save_path.format("VAE"))

        # Print and keep a (.txt) record of stuff
        msg = "Training finished at {}".format(datetime.datetime.now())
        print(msg)
        self.record.write(msg)
        self.record.close()

if __name__ == "__main__":
    agent = Agent()
    agent.train()