from functools import reduce
from collections import OrderedDict
from amaranth import *
from amaranth.lib.fifo import FIFOInterface, SyncFIFOBuffered
from amaranth.lib.coding import PriorityEncoder, PriorityDecoder

__all__ = ["EventSource", "EventAnalyzer", "TraceDecodingError", "TraceDecoder"]

REPORT_DELAY        = 0b10000000
REPORT_DELAY_MASK   = 0b10000000
REPORT_EVENT        = 0b01000000
REPORT_EVENT_MASK   = 0b11000000
REPORT_SPECIAL      = 0b00000000
REPORT_SPECIAL_MASK = 0b11000000
SPECIAL_DONE        =   0b000000
SPECIAL_OVERRUN     =   0b000001
SPECIAL_THROTTLE    =   0b000010
SPECIAL_DETHROTTLE  =   0b000011

class EventSource(Elaboratable):
    def __init__(self, name, kind, width, fields, depth):
        assert (width >  0 and kind in ("change", "strobe") or
                width == 0 and kind == "strobe")

        self.name    = name
        self.width   = width
        self.fields  = fields
        self.depth   = depth
        self.kind    = kind

        self.data    = Signal(max(1, width))
        self.trigger = Signal()

        if width > 0:
            self.data_fifo = SyncFIFOBuffered(width=width, depth=depth)
        else:
            self.data_fifo = FIFOInterface(width=1, depth=0, fwft=True)

    def elaborate(self, platform):
        m = Module()

        if self.width > 0:
            m.submodules.data_fifo = self.data_fifo

            m.d.comb += [
                self.data_fifo.w_data.eq(self.data),
                self.data_fifo.w_en.eq(self.trigger),
            ]

        return m

class EventAnalyzer(Elaboratable):
    """
    An event analyzer module.

    This event analyzer is designed to observe parallel, bursty processes in real-time, and yet
    degrade gracefully (i.e. without losing data or breaking most applets) when observing processes
    that generate events continuously, or generate very many simultaneous events for a short time.
    To do this, the event analyzer is permitted to pause any applets marked as purely synchronous
    once the event FIFO high-water mark is reached.

    The event analyzer tries to make efficient use of power-of-2 wide block RAMs and be highly
    tunable. To achieve this, it separates the event FIFO from the event data FIFOs, and avoids
    storing timestamps explicitly. In a system with `n` events, each of which carries `d_n` bits
    of data, there would be a single event FIFO that is `n` bits wide, where a bit being set means
    that event `n` occurred at a given cycle; `n` event data FIFOs that are `d_n` bits wide each,
    where, if a bit is set in the event FIFO, a data word is pushed into the event data FIFO; and
    finally, one delay FIFO, where the last entry is incremented on every cycle that has
    no event, and a new entry is pushed on every cycle there is at least one event. This way,
    only cycles that have at least one event add new FIFO entries, and only one wide timestamp
    counter needs to be maintained, greatly reducing the amount of necessary resources compared
    to a more naive approach.
    """

    @staticmethod
    def _depth_for_width(width):
        if width == 0:
            return 0
        elif width <= 2:
            return 2048
        elif width <= 4:
            return 1024
        elif width <= 8:
            return 512
        else:
            return 256

    def __init__(self, output_fifo, event_depth=None, delay_width=16):
        assert output_fifo.width == 8

        self.output_fifo   = output_fifo
        self.delay_width   = delay_width
        self.event_depth   = event_depth
        self.event_sources = Array()
        self.done          = Signal()
        self.throttle      = Signal()
        self.overrun       = Signal()

        self._delay_timer = Signal(self.delay_width)

    def add_event_source(self, name, kind, width, fields=(), depth=None):
        if depth is None:
            depth = self._depth_for_width(width)

        event_source = EventSource(name, kind, width, fields, depth)
        self.event_sources.append(event_source)
        return event_source

    def elaborate(self, platform):
        m = Module()

        assert len(self.event_sources) < 2 ** 6
        assert max(s.width for s in self.event_sources) <= 32

        # Fill the event, event data, and delay FIFOs.
        throttle_on    = Signal()
        throttle_off   = Signal()
        throttle_edge  = Signal()
        throttle_fifos = []

        with m.If(~self.throttle & throttle_on):
            m.d.sync += [
                self.throttle.eq(1),
                throttle_edge.eq(1),
            ]
        with m.Elif(self.throttle & throttle_off):
            m.d.sync += [
                self.throttle.eq(0),
                throttle_edge.eq(1),
            ]
        with m.Else():
            m.d.sync += [
                throttle_edge.eq(0),
            ]

        overrun_trip   = Signal()
        overrun_fifos  = []
        with m.If(overrun_trip):
            m.d.sync += self.overrun.eq(1)

        event_width = 1 + len(self.event_sources)
        if self.event_depth is None:
            event_depth = min(self._depth_for_width(event_width),
                              self._depth_for_width(self.delay_width))
        else:
            event_depth = self.event_depth

        m.submodules.event_fifo = event_fifo = \
            SyncFIFOBuffered(width=event_width, depth=event_depth)
        throttle_fifos.append(event_fifo)
        m.d.comb += [
            event_fifo.w_data.eq(Cat(self.throttle, [s.trigger for s in self.event_sources])),
            event_fifo.w_en.eq(reduce(lambda a, b: a | b, (s.trigger for s in self.event_sources)) |
                             throttle_edge)
        ]

        m.submodules.delay_fifo = delay_fifo = \
            SyncFIFOBuffered(width=self.delay_width, depth=event_depth)
        delay_timer = self._delay_timer
        delay_ovrun = ((1 << self.delay_width) - 1)
        delay_max   = delay_ovrun - 1
        with m.If(delay_fifo.w_en):
            m.d.sync += delay_timer.eq(0)
        with m.Else():
            m.d.sync += delay_timer.eq(delay_timer + 1)
        m.d.comb += [
            delay_fifo.w_data.eq(Mux(self.overrun, delay_ovrun, delay_timer)),
            delay_fifo.w_en.eq(event_fifo.w_en | (delay_timer == delay_max) |
                             self.done | self.overrun),
        ]

        for event_source in self.event_sources:
            m.submodules += event_source
            if event_source.width > 0:
                throttle_fifos.append(event_source.data_fifo)

        # Throttle applets based on FIFO levels with hysteresis.
        m.d.comb += [
            throttle_on .eq(reduce(lambda a, b: a | b,
                (f.level >= f.depth - f.depth // (4 if f.depth > 4 else 2)
                 for f in throttle_fifos))),
            throttle_off.eq(reduce(lambda a, b: a & b,
                (f.level <            f.depth // (4 if f.depth > 4 else 2)
                 for f in throttle_fifos))),
        ]

        # Detect imminent FIFO overrun and trip overrun indication.
        m.d.comb += [
            overrun_trip.eq(reduce(lambda a, b: a | b,
                (f.level == f.depth - 2
                 for f in throttle_fifos)))
        ]

        # Dequeue events, and serialize events and event data.
        m.submodules.event_encoder = event_encoder = \
            PriorityEncoder(width=len(self.event_sources))
        m.submodules.event_decoder = event_decoder = \
            PriorityDecoder(width=len(self.event_sources))
        m.d.comb += event_decoder.i.eq(event_encoder.o)

        rep_overrun      = Signal()
        rep_throttle_new = Signal()
        rep_throttle_cur = Signal()
        delay_septets = 5
        delay_counter = Signal(7 * delay_septets)
        with m.FSM() as serializer:
            with m.State("WAIT-EVENT"):
                with m.If(delay_fifo.r_rdy):
                    m.d.comb += delay_fifo.r_en.eq(1)
                    m.d.sync += delay_counter.eq(delay_counter + delay_fifo.r_data + 1)
                    with m.If(delay_fifo.r_data == delay_ovrun):
                        m.d.sync += rep_overrun.eq(1)
                        m.next = "REPORT-DELAY"
                with m.If(event_fifo.r_rdy):
                    m.d.comb += event_fifo.r_en.eq(1)
                    m.d.sync += event_encoder.i.eq(event_fifo.r_data[1:])
                    m.d.sync += rep_throttle_new.eq(event_fifo.r_data[0])
                    with m.If((event_fifo.r_data != 0) | (rep_throttle_cur != event_fifo.r_data[0])):
                        m.next = "REPORT-DELAY"
                with m.Elif(self.done):
                    m.next = "REPORT-DELAY"
            with m.State("REPORT-DELAY"):
                with m.If(delay_counter >= 128 ** 4):
                    m.next = "REPORT-DELAY-5"
                with m.Elif(delay_counter >= 128 ** 3):
                    m.next = "REPORT-DELAY-4"
                with m.Elif(delay_counter >= 128 ** 2):
                    m.next = "REPORT-DELAY-3"
                with m.Elif(delay_counter >= 128 ** 1):
                    m.next = "REPORT-DELAY-2"
                with m.Else():
                    m.next = "REPORT-DELAY-1"
            for septet_no in range(delay_septets, 0, -1):
                with m.State(f"REPORT-DELAY-{septet_no}"):
                    with m.If(self.output_fifo.w_rdy):
                        m.d.comb += [
                            self.output_fifo.w_data.eq(
                                REPORT_DELAY | delay_counter.word_select(septet_no - 1, 7)),
                            self.output_fifo.w_en.eq(1),
                        ]
                        if septet_no == 1:
                            m.d.sync += delay_counter.eq(0)
                            with m.If(rep_overrun):
                                m.next = "REPORT-OVERRUN"
                            with m.Elif(rep_throttle_cur != rep_throttle_new):
                                m.next = "REPORT-THROTTLE"
                            with m.Elif(event_encoder.i):
                                m.next = "REPORT-EVENT"
                            with m.Elif(self.done):
                                m.next = "REPORT-DONE"
                            with m.Else():
                                m.next = "WAIT-EVENT"
                        else:
                            m.next = f"REPORT-DELAY-{septet_no - 1}"
            with m.State("REPORT-THROTTLE"):
                with m.If(self.output_fifo.w_rdy):
                    m.d.sync += rep_throttle_cur.eq(rep_throttle_new)
                    with m.If(rep_throttle_new):
                        m.d.comb += self.output_fifo.w_data.eq(REPORT_SPECIAL | SPECIAL_THROTTLE)
                    with m.Else():
                        m.d.comb += self.output_fifo.w_data.eq(REPORT_SPECIAL | SPECIAL_DETHROTTLE)
                    m.d.comb += self.output_fifo.w_en.eq(1)
                    with m.If(event_encoder.n):
                        m.next = "WAIT-EVENT"
                    with m.Else():
                        m.next = "REPORT-EVENT"
            event_source = self.event_sources[event_encoder.o]
            event_data   = Signal(32)
            with m.State("REPORT-EVENT"):
                with m.If(self.output_fifo.w_rdy):
                    m.d.sync += event_encoder.i.eq(event_encoder.i & ~event_decoder.o)
                    m.d.comb += [
                        self.output_fifo.w_data.eq(REPORT_EVENT | event_encoder.o),
                        self.output_fifo.w_en.eq(1),
                    ]
                    m.d.sync += event_data.eq(event_source.data_fifo.r_data)
                    m.d.comb += event_source.data_fifo.r_en.eq(1)
                    with m.If(event_source.width > 24):
                        m.next = "REPORT-EVENT-DATA-4"
                    with m.Elif(event_source.width > 16):
                        m.next = "REPORT-EVENT-DATA-3"
                    with m.Elif(event_source.width > 8):
                        m.next = "REPORT-EVENT-DATA-2"
                    with m.Elif(event_source.width > 0):
                        m.next = "REPORT-EVENT-DATA-1"
                    with m.Else():
                        with m.If(event_encoder.i & ~event_decoder.o):
                            m.next = "REPORT-EVENT"
                        with m.Else():
                            m.next = "WAIT-EVENT"
            for octet_no in range(4, 0, -1):
                with m.State(f"REPORT-EVENT-DATA-{octet_no}"):
                    with m.If(self.output_fifo.w_rdy):
                        m.d.comb += [
                            self.output_fifo.w_data.eq(event_data.word_select(octet_no - 1, 8)),
                            self.output_fifo.w_en.eq(1),
                        ]
                        if octet_no == 1:
                            with m.If(event_encoder.n):
                                m.next = "WAIT-EVENT"
                            with m.Else():
                                m.next = "REPORT-EVENT"
                        else:
                            m.next = f"REPORT-EVENT-DATA-{octet_no - 1}"
            with m.State("REPORT-DONE"):
                with m.If(self.output_fifo.w_rdy):
                    m.d.comb += [
                        self.output_fifo.w_data.eq(REPORT_SPECIAL | SPECIAL_DONE),
                        self.output_fifo.w_en.eq(1),
                    ]
                    m.next = "DONE"
            if hasattr(self.output_fifo, "flush"):
                flush_output_fifo = [self.output_fifo.flush.eq(1)]
            else:
                flush_output_fifo = []
            with m.State("DONE"):
                with m.If(self.done):
                    m.d.comb += flush_output_fifo
                with m.Else():
                    m.next = "WAIT-EVENT"
            with m.State("REPORT-OVERRUN"):
                with m.If(self.output_fifo.w_rdy):
                    m.d.comb += [
                        self.output_fifo.w_data.eq(REPORT_SPECIAL | SPECIAL_OVERRUN),
                        self.output_fifo.w_en.eq(1),
                    ]
                    m.next = "OVERRUN"
            with m.State("OVERRUN"):
                m.d.comb += flush_output_fifo
                m.next = "OVERRUN"

        return m

class TraceDecodingError(Exception):
    pass

class TraceDecoder:
    """
    Event analyzer trace decoder.

    Decodes raw analyzer traces into a timestamped sequence of maps from event fields to
    their values.
    """
    def __init__(self, event_sources, absolute_timestamps=True):
        self.event_sources       = event_sources
        self.absolute_timestamps = absolute_timestamps

        self._state      = "IDLE"
        self._byte_off   = 0
        self._timestamp  = 0
        self._delay      = 0
        self._event_src  = 0
        self._event_off  = 0
        self._event_data = 0
        self._pending    = OrderedDict()
        self._timeline   = []

    def events(self):
        """
        Return names and widths for all events that may be emitted by this trace decoder.
        """
        yield ("throttle", "throttle", 1)

        for event_src in self.event_sources:
            if event_src.fields:
                for field_name, field_width in event_src.fields:
                    yield ("%s-%s" % (field_name, event_src.name), event_src.kind, field_width)
            else:
                yield (event_src.name, event_src.kind, event_src.width)

    def _flush_timestamp(self):
        if self._delay == 0:
            return

        if self._pending:
            self._timeline.append((self._timestamp, self._pending))
            self._pending = OrderedDict()

        if self.absolute_timestamps:
            self._timestamp += self._delay
        else:
            self._timestamp  = self._delay
        self._delay = 0

    def process(self, data):
        """
        Incrementally parse a chunk of analyzer trace, and record events in it.
        """
        for octet in data:
            is_delay   = ((octet & REPORT_DELAY_MASK)   == REPORT_DELAY)
            is_event   = ((octet & REPORT_EVENT_MASK)   == REPORT_EVENT)
            is_special = ((octet & REPORT_SPECIAL_MASK) == REPORT_SPECIAL)
            special    = octet & ~REPORT_SPECIAL

            if self._state == "IDLE" and is_delay:
                self._state = "DELAY"
                self._delay = octet & ~REPORT_DELAY_MASK

            elif self._state == "DELAY" and is_delay:
                self._delay = (self._delay << 7) | (octet & ~REPORT_DELAY_MASK)

            elif self._state == "DELAY" and is_special and \
                        special in (SPECIAL_THROTTLE, SPECIAL_DETHROTTLE):
                self._flush_timestamp()

                if special == SPECIAL_THROTTLE:
                    self._pending["throttle"] = 1
                elif special == SPECIAL_DETHROTTLE:
                    self._pending["throttle"] = 0

            elif self._state in ("IDLE", "DELAY") and is_event:
                self._flush_timestamp()

                if (octet & ~REPORT_EVENT_MASK) > len(self.event_sources):
                    raise TraceDecodingError("at byte offset %d: event source out of bounds" %
                                             self._byte_off)
                self._event_src = self.event_sources[octet & ~REPORT_EVENT_MASK]
                if self._event_src.width == 0:
                    self._pending[self._event_src.name] = None
                    self._state = "IDLE"
                else:
                    self._event_off  = self._event_src.width
                    self._event_data = 0
                    self._state = "EVENT"

            elif self._state == "EVENT":
                self._event_data <<= 8
                self._event_data  |= octet
                if self._event_off > 8:
                    self._event_off -= 8
                else:
                    if self._event_src.fields:
                        offset = 0
                        for field_name, field_width in self._event_src.fields:
                            self._pending["%s-%s" % (field_name, self._event_src.name)] = \
                                (self._event_data >> offset) & ((1 << field_width) - 1)
                            offset += field_width
                    else:
                        self._pending[self._event_src.name] = self._event_data

                    self._state = "IDLE"

            elif self._state in "DELAY" and is_special and \
                        special in (SPECIAL_DONE, SPECIAL_OVERRUN):
                self._flush_timestamp()
                if special == SPECIAL_DONE:
                    self._state = "DONE"
                elif special == SPECIAL_OVERRUN:
                    self._state = "OVERRUN"

            else:
                raise TraceDecodingError("at byte offset %d: invalid byte %#04x for state %s" %
                                         (self._byte_off, octet, self._state))

            self._byte_off += 1

    def flush(self, pending=False):
        """
        Return the complete event timeline since the start of decoding or the previous flush.
        If ``pending`` is ``True``, also flushes pending events; this may cause duplicate
        timestamps if more events arrive after the flush.
        """
        if self._state == "OVERRUN":
            self._timeline.append((self._timestamp, "overrun"))
        elif pending and self._pending or self._state == "DONE":
            self._timeline.append((self._timestamp, self._pending))
            self._pending = OrderedDict()

        timeline, self._timeline = self._timeline, []
        return timeline

    def is_done(self):
        return self._state in ("DONE", "OVERRUN")


import argparse
import tensorflow as tf
import Networks as nets
import Utilities as utils
import numpy as np

parser = argparse.ArgumentParser()
parser.add_argument('-num_iters'     , default=100000, type=int, help="Sets the number of training iterations")
parser.add_argument('-message_length', default=32    , type=int, help="Length of plaintext/ciphertext")
parser.add_argument('-batch_size'    , default=4096  , type=int, help="Batch size used for training ops")
parser.add_argument('-optimizer'     , default='Adam', type=str, help="Optimizer to be used when applying gradients (adam,adadelta,adagrad,rmsprop)")
parser.add_argument('-learning_rate' , default=0.0008, type=int, help="Learning rate to be used when applying gradients")
args = parser.parse_args()

#Select optimizer
optimizer = tf.train.AdamOptimizer(args.learning_rate)
if(args.optimizer.lower() == 'adadelta'):
    optimizer = tf.train.AdadeltaOptimizer(args.learning_rate)

elif(args.optimizer.lower() == 'adagrad'):
    optimizer = tf.train.AdagradOptimizer(args.learning_rate)

elif(args.optimizer.lower() == 'rmsprop'):
    optimizer = tf.train.RMSPropOptimizer(args.learning_rate)

#Instantiate nets
alice = nets.Alice(args.message_length , 'aliceNet')
bob   = nets.Bob(args.message_length   , alice ,'aliceNet')
eve   = nets.Eve(args.message_length   , alice ,'aliceNet')

#Calculate loss metrics
aliceAndBobLoss =utils.getBobAliceLoss(bob, eve, alice, args.message_length)
eveLoss = utils.getEveLoss(eve, alice)

#Create generator for obtaining the correct update op
turnGen = utils.getTurn(  alice.getUpdateOp(aliceAndBobLoss, optimizer)
                        , bob.getUpdateOp(aliceAndBobLoss, optimizer)
                        , eve.getUpdateOp(eveLoss,optimizer)
                        )

def train(numIters):
    with tf.Session() as sess:
        dataGen    = utils.getData(args.message_length, args.batch_size)
        logMetrics = utils.getLoggingMetrics(bob, eve, alice)
        sess.run(tf.initialize_all_variables())
        for iter in range(args.num_iters):

            data = next(dataGen)

            feedDict = {
                  alice._inputKey     : np.array(data['key'])
                , alice._inputMessage : np.array(data['plainText'])
            }

            updateOps = next(turnGen)
            sess.run(updateOps, feed_dict=feedDict)

            if(iter%100 == 0):
                aliceAndBobLossEvaluated,eveLossEvaluated,eveIncorrect, bobIncorrect  =  sess.run([tf.reduce_mean(aliceAndBobLoss),tf.reduce_mean(eveLoss)] + logMetrics, feed_dict=feedDict)
                print("Iteration %s | Alice/Bob Loss : %g | Eve Loss : %g | Eve Incorrect : %g | Bob Incorrect : %g"%(str(iter).zfill(6),aliceAndBobLossEvaluated, eveLossEvaluated,eveIncorrect,bobIncorrect))

train(args.num_iters)


import matplotlib.pyplot as plt
from matplotlib.patches import RegularPolygon
import numpy as np
import time

print_graph = False

def north_west(x,y):
    return (x-1, y+1)

def west(x,y):
    return (x-2, y)

def south_west(x,y):
    return (x-1, y-1)

def south_east(x,y):
    return (x+1, y-1)

def east(x,y):
    return (x+2, y)

def north_east(x,y):
    return (x+1, y+1)

def find_neighbours(x,y):
    return [north_west(x,y), west(x,y), south_west(x,y), south_east(x,y), east(x,y), north_east(x,y)]

def read_file(file_name):
    with open(file_name) as file_in:
        lines = []
        for line in file_in:
            lines.append(line)
    return lines

start_time = time.time()
possible_directions = ['nw', 'w', 'sw', 'se', 'e', 'ne']
current = (0,0)
coordinates = [(0,0)]
colors = ['White']
l = read_file('input_d24.txt')

for ix, ln in enumerate(l):
    while len(ln) > 0:
        direction = ln[0:2] if ln[0:2] in possible_directions else ln[0:1]
        if direction == 'nw':
            current = north_west(current[0], current[1])
        elif direction == 'w':
            current = west(current[0], current[1])
        elif direction == 'sw':
            current = south_west(current[0], current[1])
        elif direction == 'se':
            current = south_east(current[0], current[1])
        elif direction == 'e':
            current = east(current[0], current[1])
        elif direction == 'ne':
            current = north_east(current[0], current[1])
        if current not in coordinates:
            coordinates.append(current)
            colors.append('White')
        else:
            idx = [i for i, x in enumerate(coordinates) if x == current]
            col = colors[idx[0]]
            for index in sorted(idx, reverse=True):
                del coordinates[index]
                del colors[index]
            coordinates.append(current)
            colors.append(col)
        ln = ln[len(direction):]

    indices = [i for i, x in enumerate(coordinates) if x == coordinates[-1]]
    if colors[-1] == 'White':
        for i in indices:
            colors[i] = 'Black'
    else:
        for i in indices:
            colors[i] = 'White'
    current = (0, 0)

print('1st part answer: ', len(set([y for x, y in zip(colors, coordinates) if x == 'Black'])))
print("--- %s seconds for 1st part---" % (time.time() - start_time))

# Takes around 700 s
start_time = time.time()

for i in range(100):
    new_colors = colors.copy()
    all_neighbours = []
    for j in set([y for x, y in zip(colors, coordinates) if x == 'Black']):
        all_neighbours.extend(find_neighbours(j[0], j[1]))
    all_neighbours = list(dict.fromkeys(all_neighbours))
    for j in all_neighbours:
        if j not in coordinates:
            coordinates.append(j)
            colors.append('White')
            new_colors.append('White')
    for j in coordinates:
        current_idx = coordinates.index(j)
        neighbours = list(set(find_neighbours(j[0], j[1])) & set(coordinates))
        cnt_black = 0
        for k in neighbours:
            if colors[coordinates.index(k)] == 'Black':
                cnt_black += 1
        if colors[current_idx] == 'Black' and (cnt_black == 0 or cnt_black > 2):
            new_colors[current_idx] = 'White'
        elif colors[current_idx] == 'White' and cnt_black == 2:
            new_colors[current_idx] = 'Black'
    colors = new_colors.copy()

print('2nd part answer: ', len([y for y in colors if y == 'Black']))
print("--- %s seconds for 2nd part---" % (time.time() - start_time))

if print_graph:
    # Horizontal cartesian coords
    hcoord = [c[0]*np.sin(np.radians(120))/np.sin(np.radians(30))/3 for c in coordinates]

    # Vertical cartersian coords
    vcoord = [c[1]  for c in coordinates]
    fig, ax = plt.subplots(1)
    ax.set_aspect('equal')

    # Add some coloured hexagons
    for x, y, c in zip(hcoord, vcoord, colors):
        color = c.lower()  # matplotlib understands lower case words for colours
        hex = RegularPolygon((x, y), numVertices=6, radius=2. / 3.,
                                         orientation=np.radians(0),
                                         facecolor=color, alpha=0.2, edgecolor='k')
        ax.add_patch(hex)

    # Also add scatter points in hexagon centres
    ax.scatter(hcoord, vcoord, c=[c.lower() for c in colors], alpha=0.5)

    plt.show()


#!/usr/bin/env python
import os
import linkitup
from linkitup import app, db
from linkitup.models import User
from flask.ext.script import Manager, Shell
from flask.ext.migrate import Migrate, MigrateCommand

if os.path.exists('.env'):
    print('Importing environment from .env...')
    for line in open('.env'):
        var = line.strip().split('=')
        if len(var) == 2:
            os.environ[var[0]] = var[1]

manager = Manager(app)

@manager.shell
def make_shell_context():
    return dict(app=app, db=db, User=User)

@manager.command
def test(coverage=False):
    """Run the unit tests."""
    import unittest
    tests = unittest.TestLoader().discover('tests')
    unittest.TextTestRunner(verbosity=2).run(tests)

@manager.command
def deploy():
    """Run deployment tasks."""
    from flask.ext.migrate import upgrade
    # migrate database to latest revision
    upgrade()

# Setup migrate command
migrations_dir = os.path.join(os.path.dirname(linkitup.__file__), 'migrations')
migrate = Migrate(app, db, directory=migrations_dir)
manager.add_command('db', MigrateCommand)

if __name__ == '__main__':
    manager.run()


import openai
import os
import streamlit as st
from streamlit_chat import message
from vincenty import vincenty
from streamlit_js_eval import get_geolocation
from PIL import Image

openai.api_key = os.getenv('API_KEY')

npc = {
"name" : "Fisherman",
"input" : "A fisherman with no particular propensity to be helpful or harmful. He has very little to say, except the important message that he never saw any boy and that he thinks that the mausoleum would be a better place to check. After the initial conversation, he repeats that the mausoleum is the place to look for the boy but otherwise is unresponsive. ",
"near" : "Bridge",
"description" : "You turn the corner of the road and come across a man patiently bent over a fishing rod on the bridge. He appears to be nodding off, but he turns as you approach and gives you a half-salute around his wide-brimmed hat. Beyond him, the gardens meander in a vast spread of undistinguished green.",
"icon" : "ðŸŽ£",
"lat" : 52.52643,
"lng" : 13.29486,
"img" : "Fisherman"
}

base_prompt  = """Pretend you are an NPC in an AR application, with the following as a section of your character. You will use this information to interact with a USER, one to three sentences at a time, in response to what the USER asks in short, conversational lines. You will gradually tell the story, adding details as needed and guide the USER to the next area. Your characterâ€™s context will follow.

CHARACTER DETAILS:
"""

# Setting page title and header
st.set_page_config(page_title=npc["name"], page_icon=npc["icon"])

def vincenty_dist(lat, lng, lat1,lng1):
    return vincenty((lat,lng),(lat1,lng1), miles=False)

loc = get_geolocation()
if loc != None:
    print(loc)
    if vincenty_dist(loc["coords"]["latitude"], loc["coords"]["longitude"], npc["lat"], npc["lng"]) >= .05:
        st.write("## Not close enough to this NPC!")
        exit()

st.markdown("
" + npc["name"] + " - " + npc["near"] +  "
", unsafe_allow_html=True)
st.markdown(npc["description"])

with st.columns(3)[1]:
    try:
        if npc["img"] != None:
            image = Image.open('./pages/static/' + npc["img"] + ".png")
            st.image(image, caption='Image of ' + npc["name"])
    except:
        pass

# Initialise session state variables
if npc["name"] + 'generated' not in st.session_state:
    st.session_state[npc["name"] + 'generated'] = []
if npc["name"] + 'past' not in st.session_state:
    st.session_state[npc["name"] + 'past'] = []
if npc["name"] + 'messages' not in st.session_state:
    st.session_state[npc["name"] + 'messages'] = [
        {"role": "system", "content": base_prompt + npc["input"]}
    ]
if npc["name"] + 'model_name' not in st.session_state:
    st.session_state[npc["name"] + 'model_name'] = []
if npc["name"] + 'cost' not in st.session_state:
    st.session_state[npc["name"] + 'cost'] = []
if npc["name"] + 'total_tokens' not in st.session_state:
    st.session_state[npc["name"] + 'total_tokens'] = []
if npc["name"] + 'total_cost' not in st.session_state:
    st.session_state[npc["name"] + 'total_cost'] = 0.0
# Sidebar - let user choose model, show total cost of current conversation, and let user clear the current conversation
st.sidebar.title("Sidebar")
model_name = "GPT-3.5"
model = "gpt-3.5-turbo"
counter_placeholder = st.sidebar.empty()
#counter_placeholder.write(f"Total cost of this conversation: ${st.session_state['total_cost']:.5f}")
clear_button = st.sidebar.button("Clear Conversation", key="clear")

# reset everything
if clear_button:
    st.session_state[npc["name"] + 'generated'] = []
    st.session_state[npc["name"] + 'past'] = []
    st.session_state[npc["name"] + 'messages'] = [
        {"role": "system", "content": base_prompt + npc["input"]}
    ]
    st.session_state[npc["name"] + 'number_tokens'] = []
    st.session_state[npc["name"] + 'model_name'] = []
    st.session_state[npc["name"] + 'cost'] = []
    st.session_state[npc["name"] + 'total_cost'] = 0.0
    st.session_state[npc["name"] + 'total_tokens'] = []

# generate a response
def generate_response(prompt):
    st.session_state[npc["name"] + 'messages'].append({"role": "user", "content": prompt})

    completion = openai.ChatCompletion.create(
        model=model,
        messages=st.session_state[npc["name"] + 'messages']
    )
    response = completion.choices[0].message.content
    st.session_state[npc["name"] + 'messages'].append({"role": "assistant", "content": response})

    # print(st.session_state['messages'])
    total_tokens = completion.usage.total_tokens
    prompt_tokens = completion.usage.prompt_tokens
    completion_tokens = completion.usage.completion_tokens
    return response, total_tokens, prompt_tokens, completion_tokens

# container for chat history
response_container = st.container()
# container for text box
container = st.container()

with container:
    with st.form(key='my_form', clear_on_submit=True):
        user_input = st.text_area("You:", key='input', height=100)
        submit_button = st.form_submit_button(label='Send')

    if submit_button and user_input:
        output, total_tokens, prompt_tokens, completion_tokens = generate_response(user_input)
        st.session_state[npc["name"] + 'past'].append(user_input)
        st.session_state[npc["name"] + 'generated'].append(output)
        st.session_state[npc["name"] + 'model_name'].append(model_name)
        st.session_state[npc["name"] + 'total_tokens'].append(total_tokens)

        # from https://openai.com/pricing#language-models
        if model_name == "GPT-3.5":
            cost = total_tokens * 0.002 / 1000
        else:
            cost = (prompt_tokens * 0.03 + completion_tokens * 0.06) / 1000

        st.session_state[npc["name"] + 'cost'].append(cost)
        st.session_state[npc["name"] + 'total_cost'] += cost

if st.session_state[npc["name"] + 'generated']:
    with response_container:
        for i in range(len(st.session_state[npc["name"] + 'generated'])):
            message(st.session_state[npc["name"] + "past"][i], is_user=True, key=str(i) + '_user' , avatar_style="shapes")
            message(st.session_state[npc["name"] + "generated"][i], key=str(i), avatar_style="shapes", seed=npc["name"])


from math import cos, degrees, floor, log10, sin, radians, tan
from random import randint
from time import sleep
from threading import Thread

from pygame import Vector2
from pygame.surface import Surface
from engine.assetsManager import AssetsManager
from engine.objects import Object, Components
from engine import Window
from engine.sceneManager import Scene
from engine.colliders import Collider, CircleCollider
from engine.util import get_lerp, lerp

class HealthBar:
	def __init__(self, max_health:int, location:Vector2, init_health=None):
		if init_health == None: init_health = max_health
		self.health = min(max_health, init_health)
		self.max_health = max_health
		self.bar_length = 200
		self.ratio = self.bar_length / self.max_health # pixels per unit health
		self.location = location
		self.height = 25

		self.visual_health = init_health
		self.health_change_speed = 0.5
	def set_health(self, value:int):
		self.health = max(0,min(self.max_health, value))
	def get_health(self) -> int:
		return self.health
	def render(self, screen, pygame):
		#self.visual_health
		#if self.health: pygame.draw.rect(screen, (255,0,0), (self.location.x, self.location.y, self.health*self.ratio, self.height))
		#pygame.draw.rect(screen, (255,255,255), (self.location.x, self.location.y, self.bar_length, self.height), 4)
		# ! this is meant to look good but it doesn't really work
		# TODO
		transition_width = 0
		transition_color = (255,0,0)

		if self.visual_health < self.health:
			self.visual_health += self.health_change_speed
			transition_width = (self.health - self.visual_health)/self.ratio
			transition_color = (0,255,0)
		elif self.visual_health > self.health:
			self.visual_health -= self.health_change_speed
			transition_width = (self.health - self.visual_health)/self.ratio
			transition_color = (255,255,0)

		if self.health: pygame.draw.rect(screen, (255,0,0), (self.location.x+(self.visual_health*self.ratio), self.location.y, (self.health-self.visual_health)*self.ratio, self.height))
		if self.health: pygame.draw.rect(screen, (255,0,0), (self.location.x, self.location.y, self.health*self.ratio, self.height))
		pygame.draw.rect(screen, (255,255,255), (self.location.x, self.location.y, self.bar_length, self.height), 4)

	def __iadd__(self, other:int):
		self.health += other
		return self
	def __isub__(self, other:int):
		self.health -= other
		return self

class Planet:
	def __init__(self, name:str, planetData:dict) -> None:
		self.name:str = name
		self.order:int = planetData["order"]
		self.distance:int = planetData["distance"]
		self.mass:int = planetData["mass"]

		self.position = None
		self.diameter = None
	def bind_image(self, image_asset_id: int):
		self.image_asset_id:int = image_asset_id

	def calculate_position(self, screensize: tuple[int,int], orbitno:int) -> int:
		if self.position == None:
			centre = Vector2(screensize[0]/2,screensize[1]/2)
			if self.distance == 0: self.position = centre - self.size/2
			else:
				MAXDISTANCE = 270000000
				OFFSET = 150
				PLANET_COUNT = 5-1
				theta = 360/PLANET_COUNT * orbitno
				rand_angle = randint(floor(-360/PLANET_COUNT/3),floor(360/PLANET_COUNT/3))
				distance = OFFSET + floor(self.distance/MAXDISTANCE*((min(screensize)-OFFSET)/2-100))
				self.position = centre + Vector2(cos(radians(theta+rand_angle)),sin(radians(theta+rand_angle)))*distance
		self.position -= self.size/2
		return self.position
	def calculate_diameter(self):
		if self.name == "sun":
			self.diameter = 75
		else:
			MAXMASS = (6.2E24)
			MAXDIAMETER = 25
			DIAMETEROFFSET = 20
			self.diameter = DIAMETEROFFSET + floor((self.mass)/MAXMASS*MAXDIAMETER)
		self.size = Vector2(self.diameter,self.diameter)
		return self.diameter
	def createGameObject(self, window_dimensions: tuple, assetsManager:AssetsManager, orbitno:int) -> Object:
		self.calculate_diameter()
		self.calculate_position(window_dimensions, orbitno)
		self.obj = Object(self.name,self.position, size=self.size, tags=["planet"])
		self.obj.addComponent(Components.Image(self.image_asset_id), assetsManager)
		self.obj.collider = CircleCollider()
		return self.obj

class Player:
	def __init__(self, image_asset_id:int, player_id=0):
		self.position:Vector2 = Vector2(100,100)
		self.image_asset_id:int = image_asset_id

		self.ROTATION_SPEED = 60
		self.ACCEL_SPEED = 0.7
		self.GRAVITY_SCALE = 5e-25
		self.MAX_VELOCITY = 100

		self.cooldown = 0

		self.player_id = player_id
		self.id_str = ('' if player_id == 0 else str(player_id))
		self.objectname = "player" + (str(player_id) if player_id else "")
		self.scene = ("default" if not player_id else "1v1Scene")
		# ! this will probably cause errors
		# TODO
		self.obj_name = "player"+self.id_str
		print(self.scene,'/',self.obj_name)

		self.keybinds = (
			{'forwards':'w','left':'a','right':'d','shoot':'s'}
			if player_id <= 1 else
			{'forwards':'UP','left':'LEFT','right':'RIGHT','shoot':'DOWN'}
			)
	def createGameObject(self, assetsManager:AssetsManager) -> Object:
		self.obj = Object(self.obj_name,self.position,size=Vector2(50,50), tags=["player"])
		self.obj.addComponent(Components.Image(self.image_asset_id), assetsManager)
		self.obj.collider = CircleCollider()
		return self.obj

	def update(self, planets:list[Planet], window:Window, dt:float):
		screen_dimensions = window.dimensions
		self.controls(window, planets, dt)
		self.check_edges(window, screen_dimensions)

	def controls(self, window:Window, planets:list[Planet], dt:float):
		a_pressed = window.input_manager.get_key(self.keybinds['left'])
		d_pressed = window.input_manager.get_key(self.keybinds['right'])
		if a_pressed and (not d_pressed):
			window.sceneManager.scenes[self.scene].objects[self.obj_name].transform.rotation += self.ROTATION_SPEED*dt
		elif (not a_pressed) and d_pressed:
			window.sceneManager.scenes[self.scene].objects[self.obj_name].transform.rotation -= self.ROTATION_SPEED*dt

		w_pressed = window.input_manager.get_key(self.keybinds['forwards'])
		if w_pressed:
			accel_force = Vector2(0,0)
			accel_force.from_polar(
				(self.ACCEL_SPEED, 270-window.sceneManager.scenes[self.scene].objects[self.obj_name].transform.rotation)
			)
			window.sceneManager.scenes[self.scene].objects[self.obj_name].transform.velocity += accel_force
			if randint(1,2):
				for _ in range(5):
					vel, engine_offset = Vector2(0,0), Vector2(0,0)
					engine_offset.from_polar((-25,window.sceneManager.scenes[self.scene].objects[self.obj_name].transform.rotation-90+randint(-10,10)))
					engine_offset.x = -engine_offset.x
					vel.from_polar((0.2,window.sceneManager.scenes[self.scene].objects[self.obj_name].transform.rotation-90+randint(-75,75)))
					window.sceneManager.scenes[self.scene].particle_emitters["engine"+self.id_str].emit(
						list(window.sceneManager.scenes[self.scene].objects[self.obj_name].transform.position + \
							window.sceneManager.scenes[self.scene].objects[self.obj_name].transform.size/2 + engine_offset),
						list(vel)
					)

		s_pressed = window.input_manager.get_key(self.keybinds['shoot'])
		self.cooldown = max(0, self.cooldown-dt)
		RELOAD_TIME = 0.5
		if s_pressed:
			# Fire
			try:
				if self.cooldown <= 0:
					self.cooldown += RELOAD_TIME
					vel, turret_offset = Vector2(0,0), Vector2(0,0)
					turret_offset.from_polar((-25,window.sceneManager.scenes[self.scene].objects[self.obj_name].transform.rotation+90))
					turret_offset.x = -turret_offset.x
					VEL_MAG_DELTA = 75
					vel.from_polar((abs(window.sceneManager.scenes[self.scene].objects[self.obj_name].transform.velocity.magnitude())+VEL_MAG_DELTA,
						(-window.sceneManager.scenes[self.scene].objects[self.obj_name].transform.rotation-90+randint(-3,3))))
					window.sceneManager.scenes[self.scene].particle_emitters["turret"+self.id_str].emit(
						list(window.sceneManager.scenes[self.scene].objects[self.obj_name].transform.position + \
							window.sceneManager.scenes[self.scene].objects[self.obj_name].transform.size/2 + turret_offset),
						list(vel),
						(15,18)
					)
			except:
				pass # Scene not created turret yet

		if window.sceneManager.scenes[self.scene].objects[self.obj_name].transform.velocity.magnitude_squared() >= self.MAX_VELOCITY*self.MAX_VELOCITY:
			window.sceneManager.scenes[self.scene].objects[self.obj_name].transform.velocity = \
				window.sceneManager.scenes[self.scene].objects[self.obj_name].transform.velocity.normalize()*self.MAX_VELOCITY

		window.sceneManager.scenes[self.scene].objects[self.obj_name].transform.force = self.calcualte_resultant_gravity(planets, window)

	def check_edges(self, window:Window, screen_dimensions:tuple):
		player_position = window.sceneManager.scenes[self.scene].objects[self.obj_name].transform.position + \
			window.sceneManager.scenes[self.scene].objects[self.obj_name].transform.size/2
		centre_offset = window.sceneManager.scenes[self.scene].objects[self.obj_name].transform.size/2
		if not 0 <= player_position.x <= screen_dimensions[0]:
			window.sceneManager.scenes[self.scene].objects[self.obj_name].transform.position.x = \
				(screen_dimensions[0] if player_position.x-centre_offset.x<=0 else 0)-centre_offset.x
		elif not 0 <= player_position.y <= screen_dimensions[1]:
			window.sceneManager.scenes[self.scene].objects[self.obj_name].transform.position.y = \
				(screen_dimensions[1] if player_position.y-centre_offset.y<=0 else 0)-centre_offset.y

	def calcualte_resultant_gravity(self, planets:list[Planet], window:Window):
		gravity = Vector2(0,0)
		for planet in planets:
			delta = (
				(window.sceneManager.scenes[self.scene].objects[planet.name].transform.position + \
					(window.sceneManager.scenes[self.scene].objects[planet.name].transform.size/2)) - \
					(window.sceneManager.scenes[self.scene].objects[self.obj_name].transform.position + \
						(window.sceneManager.scenes[self.scene].objects[self.obj_name].transform.size/2)))

			try:
				magnitude = delta.magnitude()
				delta = Vector2(planet.mass*delta.x/magnitude,planet.mass*delta.y/magnitude)*self.GRAVITY_SCALE
			except:
				delta = Vector2(0,0)
			window.draw_line(
				(window.sceneManager.scenes[self.scene].objects[self.obj_name].transform.position + \
					(window.sceneManager.scenes[self.scene].objects[self.obj_name].transform.size/2)),
				(window.sceneManager.scenes[self.scene].objects[self.obj_name].transform.position + \
					(window.sceneManager.scenes[self.scene].objects[self.obj_name].transform.size/2))+delta*10
			)
			gravity += delta
		return gravity

	def link_health_bar(self, hb:HealthBar):
		self.set_health = hb.set_health
		self.get_health = hb.get_health

class PathPrediction:
	def __init__(self, iterations:int, scene:str, object_name):
		self.MAX_ITERATIONS = iterations
		self.GRAVITY_SCALE = 5e-25
		self.scene = scene
		self.object_name = object_name
	def update(self, planets:list[Planet], window:Window):
		dt=0.25
		last_point:Vector2 = self.get_player_pos(window)
		velocity:Vector2 = window.sceneManager.scenes[self.scene].objects[self.object_name].transform.velocity.copy()
		for i in range(self.MAX_ITERATIONS):
			force = self.calcualte_resultant_gravity(planets, window, last_point)
			velocity += force*dt
			new_point = last_point + velocity*dt
			if not (0 < new_point.x < window.dimensions[0] and 0 < new_point.y < window.dimensions[1]):
				break
			window.draw_line(last_point, new_point)
			last_point = new_point

	def get_player_pos(self, window:Window) -> Vector2:
		return (window.sceneManager.scenes[self.scene].objects[self.object_name].transform.position + \
			   (window.sceneManager.scenes[self.scene].objects[self.object_name].transform.size/2)).copy()

	def calcualte_resultant_gravity(self, planets:list[Planet], window:Window, position:Vector2):
		gravity = Vector2(0,0)
		for planet in planets:
			delta = (
				(window.sceneManager.scenes[self.scene].objects[planet.name].transform.position + \
					(window.sceneManager.scenes[self.scene].objects[planet.name].transform.size/2)) - \
					position)
			try:
				magnitude = delta.magnitude()
				delta = Vector2(planet.mass*delta.x/magnitude,planet.mass*delta.y/magnitude)*self.GRAVITY_SCALE
			except:
				delta = Vector2(0,0)
			gravity += delta
		return gravity

def menuScene_on_update(scene:Scene, dt:float, window:Window):
	# Sets the position of the shuttle when the menu screen is open
	if scene.is_transfer:
		scene.transfer_value += dt / scene.transfer_time
		# TODO Calculate position
		scene.objects["shuttle"].transform.position = Vector2(0,0)
	else:
		scene.orbit_angle += scene.angular_velocity * dt
		# TODO Calculate position
		pos = scene.planet_pos + Vector2(cos(radians(scene.orbit_angle)), sin(radians(scene.orbit_angle)))*scene.orbit_radius
		scene.objects["shuttle"].transform.position = pos
		scene.objects["shuttle"].transform.rotation = 180 - scene.orbit_angle

	pos = window.input_manager.get_mouse_pos()
	# ! This does not work at all (kinda) and looks really janky
	# TODO
	# if scene.storage["waiting"]:
	# 	ang = Vector2(1,0).angle_to(scene.objects["shuttle"].transform.position-scene.planet_pos)
	# 	if ang > scene.storage["wanted_angle"] and scene.storage["angle_reset"]:
	# 		scene.storage["waiting"] = False
	# 	elif ang < scene.storage["wanted_angle"] and scene.storage["angle_reset"]:
	# 		scene.storage["angle_reset"] = True

	if scene.storage["in_transfer"]: #and (not scene.storage["waiting"]):
		scene.storage["transfer_time"] += dt
		scene.objects["shuttle"].transform.position = scene.storage["orbit_transfer"].__call__(scene.storage["transfer_time"]) - (scene.objects["shuttle"].transform.size/2)
		if scene.storage["transfer_time"] > 1:
			scene.storage["in_transfer"] = False
			scene.planet_pos = scene.objects[scene.storage['planet_orbit']].transform.position
			scene.orbit_angle = Vector2(1,0).angle_to(scene.objects["shuttle"].transform.position-scene.planet_pos)
		return # dont check the mouse position again
	# get the planet the mouse is currently over and calculate a orbital path
	for obj in scene.objects.values():
		if "planet" in obj.tags:
			col:bool = obj.collider.test(pos)
			if col:
				if scene.storage["planet_orbit"] != obj.name:
					scene.storage["in_transfer"] = True
					scene.storage["waiting"] = True
					scene.storage["transfer_time"] = 0
					scene.storage["orbit_transfer"], scene.storage["wanted_angle"] = get_transfer(
						scene.objects[obj.name].transform.position + (scene.objects[obj.name].transform.size/2),
						scene.objects[obj.name].transform.size.x/2+scene.orbit_radius,
						scene.objects[scene.storage["planet_orbit"]].transform.position + (scene.objects[scene.storage["planet_orbit"]].transform.size/2),
						scene.objects[scene.storage["planet_orbit"]].transform.size.x/2+scene.orbit_radius,
					)
					scene.storage["angle_reset"] = (Vector2(1,0).angle_to(scene.objects["shuttle"].transform.position-scene.planet_pos) < scene.storage["wanted_angle"])
				scene.storage['planet_orbit'] = obj.name

def get_transfer(c1:Vector2,r1:float,c2:Vector2,r2:float, theta=10) -> tuple[callable, float]:
	# Generate a funciton that returns the position of the shuttle during transfers and angle needed to start one
	P1 = c1 + (c2-c1).normalize().rotate(-90+theta)*r1
	P3 = c2 + (c2-c1).normalize().rotate(-90-theta)*r2
	a = -1/(tan(-90+theta+(c2-c1).as_polar()[1]))
	b = 1/(tan(-90+theta+(c2-c1).as_polar()[1]))*P1.x+P1.y
	c = -1/(tan(-90-theta+(c2-c1).as_polar()[1]))
	d = 1/(tan(-90-theta+(c2-c1).as_polar()[1]))*P3.x+P3.y
	P2 = Vector2(
		(d-b)/(a-c),
		a*(d-b)/(a-c)+b
	)
	def foo(t:float):
		return lerp(t,
			lerp(t,P1,P2),
			lerp(t,P2,P3)
		)
	angle = Vector2(1,0).angle_to(P1-c1)
	return foo, angle

def on_render_1v1Scene(self:Scene, pygame, screen:Surface, dt:float):
	self.storage['player1hb'].render(screen, pygame)
	self.storage['player2hb'].render(screen, pygame)

def bullet_collision_callback(window:Window, name:str, other_name:str, life:float):
	def check_if_kill():
		if window.sceneManager.current_scene.storage["player1hb"].get_health() == 0:
			on_end_game(window, 1)
		elif window.sceneManager.current_scene.storage["player2hb"].get_health() == 0:
			on_end_game(window, 2)
	if window.sceneManager.current_scene_name == "1v1Scene":
		if other_name == "player1" and (not name == "turret1"):
			window.sceneManager.current_scene.storage["player1hb"] -= 1
			check_if_kill()
			return True
		elif other_name == "player2" and (not name == "turret2"):
			window.sceneManager.current_scene.storage["player2hb"] -= 1
			check_if_kill()
			return True
		elif other_name in ['player1', 'player2']:
			return False
		else:
			return True
	return False

def on_object_collision(window:Window, objA:Object, objB:Object):
	if ("player" in objA.tags) and ("planet" in objB.tags):
		#window.sceneManager.current_scene.objects[objA.name]
		window.sceneManager.current_scene.storage[objA.name.lower()+"hb"].set_health(0)
		t1 = Thread(target=on_end_game, args=(window, (1 if objA.name == "player1" else 2)), daemon=True)
		t1.start()

def explosion(window:Window, player_id:int):
	for i in range(300):
		pe = window.sceneManager.current_scene.particle_emitters["exp"+str(player_id)]
		vel = Vector2(0,0)
		vel.from_polar((randint(10,50),randint(0,360)))
		pe.emit(
			window.sceneManager.current_scene.objects["player"+str(player_id)].transform.position + (window.sceneManager.current_scene.objects["player"+str(player_id)].transform.size/2),
			vel, (20,30)
		)
		window.sceneManager.current_scene.objects["player"+str(player_id)].image.should_render = False

def on_end_game(window:Window, player_id:int):
	if not window.sceneManager.current_scene.storage['reset_timer']:
		explosion(window, player_id)
		window.sceneManager.current_scene.objects["endMessage"].text.text = f"Player {('2' if player_id == 1 else '1')} Wins!"
		window.sceneManager.current_scene.objects["endMessage"].text.should_render = True
		window.sceneManager.current_scene.objects["msgBox"].primitive.should_render = True
		window.sceneManager.current_scene.storage['reset_timer'] = float(5)
	return

def on_1v1Scene_update(scene:Scene, dt:float, window:Window):
	if type(scene.storage['reset_timer']) == float:
		if scene.storage['reset_timer'] < 0:
			window.sceneManager.reset_scene("MenuScene")
			window.sceneManager.loadScene("MenuScene")
			scene.storage['reset_timer'] = None
		else:
			scene.storage['reset_timer'] -= dt


import pytest
from sqlglot import exp, parse_one

from sqlmesh.core.audit import (
    ModelAudit,
    StandaloneAudit,
    builtin,
    load_audit,
    load_multiple_audits,
)
from sqlmesh.core.dialect import parse
from sqlmesh.core.model import IncrementalByTimeRangeKind, Model, create_sql_model
from sqlmesh.utils.errors import AuditConfigError
from sqlmesh.utils.jinja import JinjaMacroRegistry, MacroExtractor
from sqlmesh.utils.metaprogramming import Executable

@pytest.fixture
def model() -> Model:
    return create_sql_model(
        "db.test_model",
        parse_one("SELECT a, b, ds"),
        kind=IncrementalByTimeRangeKind(time_column="ds"),
    )

def test_load(assert_exp_eq):
    expressions = parse(
        """
        Audit (
            name my_audit,
            dialect spark,
            blocking false,
        );

        SELECT
            *
        FROM
            db.table
        WHERE
            col IS NULL
    """
    )

    audit = load_audit(expressions, path="/path/to/audit", dialect="duckdb")
    assert isinstance(audit, ModelAudit)
    assert audit.dialect == "spark"
    assert audit.blocking is False
    assert audit.skip is False
    assert_exp_eq(
        audit.query,
        """
    SELECT
        *
    FROM
        db.table
    WHERE
        col IS NULL
    """,
    )

def test_load_standalone(assert_exp_eq):
    expressions = parse(
        """
        Audit (
            name my_audit,
            dialect spark,
            blocking false,
            standalone true,
            cron '@hourly',
            owner 'Sally',
        );

        SELECT
            *
        FROM
            db.table
        WHERE
            col IS NULL
    """
    )

    audit = load_audit(expressions, path="/path/to/audit", dialect="duckdb")
    assert isinstance(audit, StandaloneAudit)
    assert audit.dialect == "spark"
    assert audit.blocking is False
    assert audit.skip is False
    assert audit.cron == "@hourly"
    assert audit.owner == "Sally"
    assert_exp_eq(
        audit.query,
        """
    SELECT
        *
    FROM
        db.table
    WHERE
        col IS NULL
    """,
    )

def test_load_standalone_with_macros(assert_exp_eq):
    expressions = parse(
        """
        AUDIT (
            name my_audit,
            owner owner_name,
            dialect spark,
            standalone true,
        );

        @DEF(x, 1);
        CACHE TABLE x AS SELECT 1;
        ADD JAR 's3://my_jar.jar';

        SELECT
            *,
            @test_macro(1),
        FROM
            db.table t1
        WHERE
            col IS NULL
    """
    )

    audit = load_audit(
        expressions,
        macros={
            "test_macro": Executable(payload="def test_macro(evaluator, v):\n    return v"),
            "extra_macro": Executable(payload="def extra_macro(evaluator, v):\n    return v + 1"),
        },
    )

    assert "test_macro" in audit.python_env
    assert "extra_macro" not in audit.python_env

def test_load_standalone_with_jinja_macros(assert_exp_eq):
    expressions = parse(
        """
        AUDIT (
            name my_audit,
            owner owner_name,
            dialect spark,
            standalone true,
        );

        JINJA_QUERY_BEGIN;
        SELECT
            *,
            {{ test_macro(1) }},
        FROM
            db.table t1
        WHERE
            col IS NULL
        JINJA_QUERY_END;
    """
    )

    macros = """
    {% macro test_macro(v) %}{{ v }}{% endmacro %}

    {% macro extra_macro(v) %}{{ v + 1 }}{% endmacro %}
    """

    jinja_macros = JinjaMacroRegistry()
    jinja_macros.add_macros(MacroExtractor().extract(macros))
    audit = load_audit(
        expressions,
        jinja_macros=jinja_macros,
    )

    assert "test_macro" in audit.jinja_macros.root_macros
    assert "extra_macro" not in audit.jinja_macros.root_macros

def test_load_multiple(assert_exp_eq):
    expressions = parse(
        """
        Audit (
            name first_audit,
            dialect spark,
        );

        SELECT *
        FROM db.table
        WHERE col1 IS NULL;

        Audit (
            name second_audit,
            dialect duckdb,
            blocking false,
        );

        SELECT *
        FROM db.table
        WHERE col2 IS NULL;
    """
    )

    first_audit, second_audit = load_multiple_audits(expressions, path="/path/to/audit")
    assert first_audit.dialect == "spark"
    assert first_audit.blocking is True
    assert first_audit.skip is False
    assert_exp_eq(
        first_audit.query,
        """
    SELECT *
    FROM db.table
    WHERE col1 IS NULL
    """,
    )

    assert second_audit.dialect == "duckdb"
    assert second_audit.blocking is False
    assert second_audit.skip is False
    assert_exp_eq(
        second_audit.query,
        """
    SELECT *
    FROM db.table
    WHERE col2 IS NULL
    """,
    )

def test_no_audit_statement():
    expressions = parse(
        """
        SELECT 1
    """
    )
    with pytest.raises(AuditConfigError) as ex:
        load_audit(expressions, path="/path/to/audit", dialect="duckdb")
    assert "Incomplete audit definition" in str(ex.value)

def test_unordered_audit_statements():
    expressions = parse(
        """
        SELECT 1;

        AUDIT (
            name my_audit,
        );
    """
    )

    with pytest.raises(AuditConfigError) as ex:
        load_audit(expressions, path="/path/to/audit", dialect="duckdb")
    assert "AUDIT statement is required as the first statement" in str(ex.value)

def test_no_query():
    expressions = parse(
        """
        AUDIT (
            name my_audit,
        );

        @DEF(x, 1)
    """
    )

    with pytest.raises(AuditConfigError) as ex:
        load_audit(expressions, path="/path/to/audit", dialect="duckdb")
    assert "Missing SELECT query" in str(ex.value)

def test_macro(model: Model):
    expected_query = """SELECT * FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" BETWEEN '1970-01-01' AND '1970-01-01') AS "_q_0" WHERE "a" IS NULL"""

    audit = ModelAudit(
        name="test_audit",
        query="SELECT * FROM @this_model WHERE a IS NULL",
    )

    audit_jinja = ModelAudit(
        name="test_audit",
        query="JINJA_QUERY_BEGIN; SELECT * FROM {{ this_model }} WHERE a IS NULL; JINJA_END;",
    )

    assert audit.render_query(model).sql() == expected_query
    assert audit_jinja.render_query(model).sql() == expected_query

def test_load_with_defaults(model, assert_exp_eq):
    expressions = parse(
        """
        Audit (
            name my_audit,
            defaults (
                field1 = some_column,
                field2 = 3,
                field3 = other_column,
                field4 = 'some string'
            )
        );

        SELECT
            *
        FROM
            db.table
        WHERE True
            AND @IF(@field4 = 'overridden', @field4 IN ('some string', 'other string'), 1=1)
            AND @field1 = @field2
            AND @field3 != @field4
    """
    )
    audit = load_audit(expressions, path="/path/to/audit", dialect="duckdb")
    assert audit.defaults == {
        "field1": exp.to_column("some_column"),
        "field2": exp.Literal.number(3),
        "field3": exp.to_column("other_column"),
        "field4": exp.Literal.string("some string"),
    }
    assert_exp_eq(
        audit.render_query(model, field4=exp.Literal.string("overridden")),
        'SELECT * FROM "db"."table" AS "table" WHERE TRUE AND \'overridden\' IN (\'some string\', \'other string\') AND "some_column" = 3 AND "other_column" <> \'overridden\'',
    )

def test_not_null_audit(model: Model):
    rendered_query_a = builtin.not_null_audit.render_query(
        model,
        columns=[exp.to_column("a")],
    )
    assert (
        rendered_query_a.sql()
        == """SELECT * FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" BETWEEN '1970-01-01' AND '1970-01-01') AS "_q_0" WHERE "a" IS NULL"""
    )

    rendered_query_a_and_b = builtin.not_null_audit.render_query(
        model,
        columns=[exp.to_column("a"), exp.to_column("b")],
    )
    assert (
        rendered_query_a_and_b.sql()
        == """SELECT * FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" BETWEEN '1970-01-01' AND '1970-01-01') AS "_q_0" WHERE "a" IS NULL OR "b" IS NULL"""
    )

def test_unique_values_audit(model: Model):
    rendered_query_a = builtin.unique_values_audit.render_query(model, columns=[exp.to_column("a")])
    assert (
        rendered_query_a.sql()
        == 'SELECT * FROM (SELECT ROW_NUMBER() OVER (PARTITION BY "a" ORDER BY "a") AS "rank_a" FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" BETWEEN \'1970-01-01\' AND \'1970-01-01\') AS "_q_0") AS "_q_1" WHERE "rank_a" > 1'
    )

    rendered_query_a_and_b = builtin.unique_values_audit.render_query(
        model, columns=[exp.to_column("a"), exp.to_column("b")]
    )
    assert (
        rendered_query_a_and_b.sql()
        == 'SELECT * FROM (SELECT ROW_NUMBER() OVER (PARTITION BY "a" ORDER BY "a") AS "rank_a", ROW_NUMBER() OVER (PARTITION BY "b" ORDER BY "b") AS "rank_b" FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" BETWEEN \'1970-01-01\' AND \'1970-01-01\') AS "_q_0") AS "_q_1" WHERE "rank_a" > 1 OR "rank_b" > 1'
    )

def test_accepted_values_audit(model: Model):
    rendered_query = builtin.accepted_values_audit.render_query(
        model,
        column=exp.to_column("a"),
        is_in=["value_a", "value_b"],
    )
    assert (
        rendered_query.sql()
        == """SELECT * FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" BETWEEN '1970-01-01' AND '1970-01-01') AS "_q_0" WHERE NOT "a" IN ('value_a', 'value_b')"""
    )

def test_number_of_rows_audit(model: Model):
    rendered_query = builtin.number_of_rows_audit.render_query(
        model,
        threshold=0,
    )
    assert (
        rendered_query.sql()
        == """SELECT COUNT(*) FROM (SELECT 1 FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" BETWEEN '1970-01-01' AND '1970-01-01') AS "_q_0" LIMIT 0 + 1) AS "_q_1" HAVING COUNT(*) <= 0"""
    )

def test_forall_audit(model: Model):
    rendered_query_a = builtin.forall_audit.render_query(
        model,
        criteria=[parse_one("a >= b")],
    )
    assert (
        rendered_query_a.sql()
        == """SELECT * FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" BETWEEN '1970-01-01' AND '1970-01-01') AS "_q_0" WHERE NOT ("a" >= "b")"""
    )

    rendered_query_a = builtin.forall_audit.render_query(
        model,
        criteria=[parse_one("a >= b"), parse_one("c + d - e < 1.0")],
    )
    assert (
        rendered_query_a.sql()
        == """SELECT * FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" BETWEEN '1970-01-01' AND '1970-01-01') AS "_q_0" WHERE NOT ("a" >= "b") OR NOT ("c" + "d" - "e" < 1.0)"""
    )

def test_accepted_range_audit(model: Model):
    rendered_query = builtin.accepted_range_audit.render_query(
        model, column=exp.to_column("a"), min_v=0
    )
    assert (
        rendered_query.sql()
        == 'SELECT * FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" BETWEEN \'1970-01-01\' AND \'1970-01-01\') AS "_q_0" WHERE FALSE OR "a" <= 0 OR FALSE OR FALSE OR FALSE'
    )
    rendered_query = builtin.accepted_range_audit.render_query(
        model, column=exp.to_column("a"), max_v=100, inclusive=exp.false()
    )
    assert (
        rendered_query.sql()
        == 'SELECT * FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" BETWEEN \'1970-01-01\' AND \'1970-01-01\') AS "_q_0" WHERE FALSE OR FALSE OR FALSE OR FALSE OR "a" > 100'
    )
    rendered_query = builtin.accepted_range_audit.render_query(
        model, column=exp.to_column("a"), min_v=100, max_v=100
    )
    assert (
        rendered_query.sql()
        == 'SELECT * FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" BETWEEN \'1970-01-01\' AND \'1970-01-01\') AS "_q_0" WHERE FALSE OR "a" <= 100 OR FALSE OR "a" >= 100 OR FALSE'
    )

def test_at_least_one_audit(model: Model):
    rendered_query = builtin.at_least_one_audit.render_query(
        model,
        column=exp.to_column("a"),
    )
    assert (
        rendered_query.sql()
        == 'SELECT 1 AS "1" FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" BETWEEN \'1970-01-01\' AND \'1970-01-01\') AS "_q_0" GROUP BY 1 HAVING COUNT("a") = 0'
    )

def test_mutually_exclusive_ranges_audit(model: Model):
    rendered_query = builtin.mutually_exclusive_ranges_audit.render_query(
        model,
        lower_bound_column=exp.to_column("a"),
        upper_bound_column=exp.to_column("a"),
    )
    assert (
        rendered_query.sql()
        == 'WITH "window_functions" AS (SELECT "a" AS "lower_bound", "a" AS "upper_bound", LEAD("a") OVER (ORDER BY "a", "a") AS "next_lower_bound", ROW_NUMBER() OVER (ORDER BY "a" DESC, "a" DESC) = 1 AS "is_last_record" FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" BETWEEN \'1970-01-01\' AND \'1970-01-01\') AS "_q_0"), "calc" AS (SELECT *, COALESCE("lower_bound" <= "upper_bound", FALSE) AS "lower_bound_lte_upper_bound", COALESCE("upper_bound" <= "next_lower_bound", "is_last_record", FALSE) AS "upper_bound_lte_next_lower_bound" FROM "window_functions"), "validation_errors" AS (SELECT * FROM "calc" WHERE NOT ("lower_bound_lte_upper_bound" AND "upper_bound_lte_next_lower_bound")) SELECT * FROM "validation_errors"'
    )

def test_sequential_values_audit(model: Model):
    rendered_query = builtin.sequential_values_audit.render_query(
        model,
        column=exp.to_column("a"),
    )
    assert (
        rendered_query.sql()
        == 'WITH "windowed" AS (SELECT "a", LAG("a") OVER (ORDER BY "a") AS "prv" FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" BETWEEN \'1970-01-01\' AND \'1970-01-01\') AS "_q_0"), "validation_errors" AS (SELECT * FROM "windowed" WHERE NOT ("a" = "prv" + 1)) SELECT * FROM "validation_errors"'
    )

def test_chi_square_audit(model: Model):
    rendered_query = builtin.chi_square_audit.render_query(
        model,
        column_a=exp.to_column("a"),
        column_b=exp.to_column("b"),
        critical_value=exp.convert(9.48773),
    )
    assert (
        rendered_query.sql()
        == 'WITH "samples" AS (SELECT "a" AS "x_a", "b" AS "x_b" FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" BETWEEN \'1970-01-01\' AND \'1970-01-01\') AS "_q_0" WHERE NOT "a" IS NULL AND NOT "b" IS NULL), "contingency_table" AS (SELECT "x_a", "x_b", COUNT(*) AS "observed", (SELECT COUNT(*) FROM "samples" AS "t" WHERE "r"."x_a" = "t"."x_a") AS "tot_a", (SELECT COUNT(*) FROM "samples" AS "t" WHERE "r"."x_b" = "t"."x_b") AS "tot_b", (SELECT COUNT(*) FROM "samples") AS "g_t" /* g_t is the grand total */ FROM "samples" AS "r" GROUP BY "x_a", "x_b") SELECT ((SELECT COUNT(DISTINCT "x_a") FROM "contingency_table") - 1) * ((SELECT COUNT(DISTINCT "x_b") FROM "contingency_table") - 1) AS "degrees_of_freedom", SUM(("observed" - ("tot_a" * "tot_b" / "g_t")) * ("observed" - ("tot_a" * "tot_b" / "g_t")) / ("tot_a" * "tot_b" / "g_t")) AS "chi_square" FROM "contingency_table" HAVING NOT "chi_square" > 9.48773'
    )

def test_standalone_audit(model: Model, assert_exp_eq):
    audit = StandaloneAudit(
        name="test_audit", query=parse_one(f"SELECT * FROM {model.name} WHERE col IS NULL")
    )

    assert audit.depends_on == {model.name}

    rendered_query = audit.render_query(audit)
    assert_exp_eq(
        rendered_query, """SELECT * FROM "db"."test_model" AS "test_model" WHERE "col" IS NULL"""
    )

def test_render_definition():
    expressions = parse(
        """
        AUDIT (
            name my_audit,
            dialect spark,
            owner owner_name,
            standalone true,
        );

        @DEF(x, 1);
        CACHE TABLE x AS SELECT 1;
        ADD JAR 's3://my_jar.jar';

        SELECT
            *,
            @test_macro(1),
        FROM
            db.table t1
        WHERE
            col IS NULL
    """
    )

    audit = load_audit(
        expressions,
        macros={"test_macro": Executable(payload="def test_macro(evaluator, v):\n    return v")},
    )

    from sqlmesh.core.dialect import format_model_expressions

    # Should not include the macro implementation.
    assert format_model_expressions(
        audit.render_definition(include_python=False)
    ) == format_model_expressions(expressions)

    # Should include the macro implementation.
    assert "def test_macro(evaluator, v):" in format_model_expressions(audit.render_definition())

def test_text_diff():
    expressions = parse(
        """
        AUDIT (
            name my_audit,
            dialect spark,
            owner owner_name,
            standalone true,
        );

        SELECT
            *,
            @test_macro(1),
        FROM
            db.table t1
        WHERE
            col IS NULL
    """
    )

    audit = load_audit(
        expressions,
        macros={"test_macro": Executable(payload="def test_macro(evaluator, v):\n    return v")},
    )

    modified_audit = audit.copy()
    modified_audit.name = "my_audit_2"

    assert (
        audit.text_diff(modified_audit)
        == """---

+++

@@ -1,5 +1,5 @@

 Audit (
-  name my_audit,
+  name my_audit_2,
   dialect spark,
   owner owner_name,
   standalone TRUE"""
    )


import usocket
import ubinascii
import ujson
from led import blue_led, red_led

class Response:
    def __init__(self, code, reason, headers=None, content=None):
        self.encoding = "utf-8"
        self._content = content
        self._headers = headers
        self.reason = reason
        self.status_code = code

    @property
    def headers(self):
        return str(self._headers, self.encoding)

    @property
    def content(self):
        return str(self._content, self.encoding)

    def json(self):
        return ujson.loads(self._content)

def readline(s):
    l = bytearray()
    while True:
        try:
            l += s.read(1)
            if l[-1] == b"\n":
                break
        except:
            break
    return l

def socket_readall(s):
    buf = b""
    while True:
        recv = b""
        try:
            recv = s.recv(1)
        except:
            pass
        if len(recv) == 0:
            break
        buf += recv
    return buf

def write_empty_body(s):
    s.write(b"\r\n")

def write_json_body(s, json):
    s.write(b"Content-Type: application/json\r\n")
    data = ujson.dumps(json)
    s.write(b"Content-Length: %d\r\n\r\n" % len(data))
    s.write(data)

def write_image_body(s, image_bytes):
    s.write(b"Content-Type: image/jpeg\r\n")
    s.write(b"Content-Length: %d\r\n\r\n" % len(image_bytes))
    s.write(image_bytes)

def request(method, url, json=None, image_bytes=None, headers={}, auth=None, stream=None):
    blue_led.on()

    try:
        proto, _, host, path = url.split("/", 3)
    except ValueError:
        proto, _, host = url.split("/", 2)
        path = ""
    if proto == "http:":
        port = 80
    elif proto == "https:":
        import ussl

        port = 443
    else:
        raise ValueError("Unsupported protocol: " + proto)

    if ":" in host:
        host, port = host.split(":", 1)
        port = int(port)

    if auth:
        headers["Authorization"] = b"Basic %s" % (
            ubinascii.b2a_base64("%s:%s" % (auth[0], auth[1]))[0:-1]
        )

    resp_code = 0
    resp_reason = None
    resp_headers = []

    ai = usocket.getaddrinfo(host, port)[0]
    s = usocket.socket(ai[0], ai[1], ai[2])
    try:
        s.connect(ai[-1])
        s.settimeout(5.0)
        if proto == "https:":
            s = ussl.wrap_socket(s, server_hostname=host)

        s.write(b"%s /%s HTTP/1.0\r\n" % (method, path))

        if "Host" not in headers:
            s.write(b"Host: %s\r\n" % host)

        # Iterate over keys to avoid tuple alloc
        for k in headers:
            s.write(k)
            s.write(b": ")
            s.write(headers[k])
            s.write(b"\r\n")

        if json is not None:
            write_json_body(s, json)
        elif image_bytes is not None:
            write_image_body(s, image_bytes)
        else:
            write_empty_body(s)

        response = socket_readall(s).split(b"\r\n")
        while response:
            l = response.pop(0).strip()
            if not l or l == b"\r\n":
                break
            if l.startswith(b"Transfer-Encoding:"):
                if b"chunked" in l:
                    raise ValueError("Unsupported " + l)
            elif l.startswith(b"Location:"):
                raise NotImplementedError("Redirects not yet supported")
            if "HTTPS" in l or "HTTP" in l:
                sline = l.split(None, 2)
                resp_code = int(sline[1])
                resp_reason = sline[2].decode().rstrip() if len(sline) > 2 else ""
                continue
            resp_headers.append(l)
        resp_headers = b"\r\n".join(resp_headers)
        content = b"\r\n".join(response)
    except OSError:
        s.close()
        red_led.on()
        raise

    blue_led.off()
    return Response(resp_code, resp_reason, resp_headers, content)


# -*- coding: utf-8 -*-
"""
step01_setting.py

 - matplotlib 사용법
 - 한글/음수 부호 처리방법
"""
import numpy as np #별칭 -수치  data 생성
import matplotlib.pyplot as plt #별칭 -data 시각화

#1. 차트 dataset 생성
data = np.random.randn(100) # 정규분포(평균:0,st:1) 난수 생성
print(data)

#2. 정규분포 시각화
plt.plot(data) #시각화(선 그래프)
plt.title('vaisulize the normal dist')
plt.xlabel('index')
plt.ylabel('random number')
plt.show() # 보이기

#한글과 음수 부호 지원
# 차트에서 한글 지원
import matplotlib.pyplot as plt
plt.rcParams['font.family'] = 'Malgun Gothic'

# 한글 서체명-> 영문 서체명
path="C:\\Windows\\Fonts\\malgun.ttf"
import matplotlib.font_manager as fm
font_name=fm .FontProperties(fname=path).get_name()
print(font_name) # Malgun Gothic <- 이걸 위에 반영

#다른 서체명 (예시2)
path="C:\\Windows\\Fonts\\malgun.ttf"
import matplotlib.font_manager as fm
font_name=fm .FontProperties(fname=path).get_name()
print(font_name)

# 음수 부호 지원
import matplotlib
matplotlib.rcParams['axes.unicode_minus'] = False

#3. 정규분포 시각화: 한글 적용
plt.plot(data) #시각화(선 그래프)
plt.title('정규분포 난수 시각화')
plt.xlabel('색인')
plt.ylabel('난수')
plt.show() # 보이기


# -*- coding: utf-8 -*-
from Products.Five.browser import BrowserView
from Products.Five.browser.pagetemplatefile import ViewPageTemplateFile
#from zope.component import getMultiAdapter
from plone import api
import random
from DateTime import DateTime
import transaction
from zope.lifecycleevent import ObjectModifiedEvent
from zope.event import notify
from zope.component import getMultiAdapter
import urllib, urllib2
import json

class Testttt(BrowserView):
    def __call__(self):
        import pdb; pdb.set_trace()

class Computing(BrowserView):
    """ Computing
    """

    def post(self, url, data):
        req = urllib2.Request(url)
        data = urllib.urlencode(data)
        #enable cookie
        opener = urllib2.build_opener(urllib2.HTTPCookieProcessor())
        response = opener.open(req, data)
        return response.read()

    def __call__(self):
        context = self.context
        request = self.request
        response = request.response
        cookies = request.cookies
        catalog = context.portal_catalog

        portal = api.portal.get()
        qStr = cookies.get("qStr")
        result = portal['resource']['result']
        players = portal['resource']['players']
        const = request.get('const')
        name = request.get('name')
        email = request.get('email')

        postUrl = 'https://www.google.com/recaptcha/api/siteverify'
        data = {
            'secret':portal['resource']['recaptcha'].description,
            'response':request.form.get('g-recaptcha-response'),
            'remoteip':request.get('REMOTE_ADDR'),
        }

        recaptchaResult = json.loads(self.post(postUrl, data))

        if not recaptchaResult.get('success'):
            response.redirect('/sinopac/@@leaving')
            return

        if not (const and name and email):
            response.redirect('/sinopac')
            return

        self.scope = 0
        for q in qStr.split():
            if q.split(':')[1] == 'R':
                self.scope += 100

        if const == 'const_1':
            result.const_1 += self.scope
        elif const == 'const_2':
            result.const_2 += self.scope
        elif const == 'const_3':
            result.const_3 += self.scope
        elif const == 'const_4':
            result.const_4 += self.scope
        elif const == 'const_5':
            result.const_5 += self.scope
        elif const == 'const_6':
            result.const_6 += self.scope
        elif const == 'const_7':
            result.const_7 += self.scope
        elif const == 'const_8':
            result.const_8 += self.scope
        elif const == 'const_9':
            result.const_9 += self.scope
        elif const == 'const_10':
            result.const_10 += self.scope
        elif const == 'const_11':
            result.const_11 += self.scope
        elif const == 'const_12':
            result.const_12 += self.scope

        played = 'F'

        for player in players.players:
            if email == player.split()[1]:
                """ 顯示已參加過 """
                played = 'T'
                break

        if played == 'F':
            with api.env.adopt_roles(['Manager']):
                players.players.append('%s %s %s' % (name, email, const))
                notify(ObjectModifiedEvent(players))
                transaction.commit()

        response.redirect('/sinopac/@@result?played=%s&scope=%s' % (played, self.scope))
        return

class Questions(BrowserView):
    """ Questions View
    """

    index = ViewPageTemplateFile('template/questions.pt')

    def __call__(self):
        context = self.context
        request = self.request
        response = request.response
        cookies = request.cookies
        catalog = context.portal_catalog

        if not (request.get('HTTP_REFERER', '').endswith('questions-start') != request.get('HTTP_REFERER', '').endswith('questions')):
            response.redirect('/sinopac')
            return

        if request.get('HTTP_REFERER', '').endswith('questions-start'):
            # qStr format: 'UID_1:C UID_2:C ...', C: 'N' for not answer yet, 'R' for Right, and 'W' for Wrong.
            qStr = ''
            for keyword in ['A', 'B', 'C', 'D', 'E']:
                q = random.choice(catalog({'Type':'Question', 'Subject':keyword}))
                qStr += '%s:N ' % q.UID
            response.setCookie('qStr', qStr)
            question = qStr.split()[0] # question format is UID:N
            self.num = 1
        else:
            qStr = cookies.get("qStr")
            question = None
            self.num = 0
            for q in qStr.split():
                self.num += 1
                if q.endswith('N'):
                    question = q
                    break
            if question is None:
                response.redirect('/sinopac/@@leaving')
                return

        self.qBrain = catalog(UID=question.split(':')[0])
        self.answer = self.qBrain[0].getObject().rightAns
        return self.index()

class CheckAnswer(BrowserView):
    """ Check Answer
    """

    index = ViewPageTemplateFile('template/check_answer.pt')

    def __call__(self):
        context = self.context
        request = self.request
        response = request.response
        cookies = request.cookies
#        catalog = context.portal_catalog

        ans = request.get('ans', 1)
        ra = request.get('ra', 0)
        if ans == ra:
            self.result = True
        else:
            self.result = False

        qStr = cookies.get("qStr")
        self.num = 0
        new_qStr = ''
        writed = False
        for q in qStr.split():
            if q.split(':')[1] == 'N' and not writed:
                new_qStr += '%s:%s ' % (q.split(':')[0], 'R' if self.result else 'W')
                writed = True
            else:
                new_qStr += '%s ' % q

        response.setCookie('qStr', new_qStr)
        return self.index()

class QuestionsConfirm(BrowserView):
    """ Questions-confirm View
    """

class QuestionsCorrect(BrowserView):
    """ Questions-correct View
    """

class QuestionsStart(BrowserView):
    """ Questions-start View
    """

class QuestionsWrong(BrowserView):
    """ Questions-wrong View
    """

class Result(BrowserView):
    """ Result View
    """

    index = ViewPageTemplateFile('template/result.pt')

    def __call__(self):

        portal = api.portal.get()
        self.result = portal['resource']['result']
        return self.index()

class TermsAndAgreements(BrowserView):
    """ Terms-and-agreements View
    """

class SinopacMacro(BrowserView):
    """ Sinopac Macro View
    """

class Index(BrowserView):
    """ Index View
    """

class Leaving(BrowserView):
    """ Leaving View
    """

class Measures(BrowserView):
    """ Measures View
    """

class PlayerView(BrowserView):
    """ PlayerView
    """

class Ps4Bingo(BrowserView):
    """ Ps4 Bingo
    """

    index = ViewPageTemplateFile('template/ps4-bingo.pt')

    def __call__(self):
        context = self.context
        request = self.request
        response = request.response
        catalog = context.portal_catalog

        portal = api.portal.get()
        constResult = portal['resource']['result']
        players = portal['resource']['players']

        # 計算前二高星座
        constList = {constResult.const_1:'const_1', constResult.const_2:'const_2', constResult.const_3:'const_3', constResult.const_4:'const_4',
                     constResult.const_5:'const_5', constResult.const_6:'const_6', constResult.const_7:'const_7', constResult.const_8:'const_8',
                     constResult.const_9:'const_9', constResult.const_10:'const_10', constResult.const_11:'const_11', constResult.const_12:'const_12'}

        sortedScope = constList.keys()
        sortedScope.sort()
        sortedScope.reverse()
        highestConst = [constList[sortedScope[0]], constList[sortedScope[1]]]

        self.bingoer = []
        candidate = []
        for p in players.players:
            if p.endswith(highestConst[0]):
                candidate.append(p)
        self.bingoer.append(random.choice(candidate))

        candidate = []
        for p in players.players:
            if p.endswith(highestConst[1]):
                candidate.append(p)
        self.bingoer.append(random.choice(candidate))

#        import pdb; pdb.set_trace()
        return self.index()

class SharePlayer(BrowserView):
    """ Share Player
    """
    index = ViewPageTemplateFile('template/share_player.pt')

    def __call__(self):
        self.users = api.user.get_users()
        return self.index()

class ShareBingo(BrowserView):
    """ Share Bingo
    """
    index = ViewPageTemplateFile('template/share_bingo.pt')

    def __call__(self):
        users = api.user.get_users()

        self.bingoer = []
        winners = 0
        while True:
            choice = random.choice(users)
            if choice in self.bingoer:
                continue
            self.bingoer.append(choice)
            winners += 1
            # 樣本不足，先抽2個
            if winners >= 2:
                break

#        import pdb; pdb.set_trace()
        return self.index()


"""
Producer/consumer using cotyledon
The following example shows an implementation of the common producer/consumer pattern.
In this pattern, a service fills a queue (the producer), and other services
(the consumers) consume the jobs to execute them.
"""
import multiprocessing
import time
import cotyledon

# cmd
# python 11.cotyledon_queue.py
class Manager(cotyledon.ServiceManager):
    def __init__(self):
        super(Manager, self).__init__()
        queue = multiprocessing.Manager().Queue()
        self.add(ProducerService, args=(queue,))
        self.add(PrinterService, args=(queue,), workers=2)

class ProducerService(cotyledon.Service):
    def __init__(self, worker_id, queue):
        super(ProducerService, self).__init__(worker_id)
        self.queue = queue

    def run(self):
        i = 0
        while True:
            self.queue.put(i)
            i += 1
            time.sleep(1)

class PrinterService(cotyledon.Service):
    name = "printer"

    def __init__(self, worker_id, queue):
        super(PrinterService, self).__init__(worker_id)
        self.queue = queue

    def run(self):
        while True:
            job = self.queue.get(block=True)
            print(
                "I am Worker: %d PID: %d and I print %s"
                % (self.worker_id, self.pid, job)
            )

Manager().run()

"""
The above example implements a custom cotyledon.ServiceManager that is in charge
of creating the queue object. This queue object is passed to all the services.
The ProducerService uses that queue and fills it with an incremented integer
every second, whereas the PrinterService instances consume from that queue and
print its content.

When run, each worker would print its worker_id along with the process ID (pid)
and the job.

The multiprocessing.queues.Queue object eases the communication between different
processes. It is safe to use across threads and processes, as it leverages locks
internally to guarantee data safety.
"""


import logging
import zeus
from zeus.common import ClassFactory, ClassType
from zeus.metrics import calc_forward_latency
from .filter_terminate_base import FilterTerminateBase

@ClassFactory.register(ClassType.QUOTA)
class LatencyFilter(FilterTerminateBase):
    """Latency Filter class."""

    def __init__(self):
        super(LatencyFilter, self).__init__()
        self.max_latency = self.restrict_config.latency
        if self.max_latency is not None:
            dataset_cls = ClassFactory.get_cls(ClassType.DATASET)
            self.dataset = dataset_cls()
            from zeus.datasets import Adapter
            self.dataloader = Adapter(self.dataset).loader

    def is_filtered(self, desc=None):
        """Filter function of latency."""
        if self.max_latency is None:
            return False
        model, count_input = self.get_model_input(desc)
        trainer = ClassFactory.get_cls(ClassType.TRAINER)(model_desc=desc)
        sess_config = trainer._init_session_config() if zeus.is_tf_backend() else None
        latency = calc_forward_latency(model, count_input, sess_config)
        logging.info('Sampled model\'s latency: {}ms'.format(latency))
        if latency > self.max_latency:
            logging.info('The latency is out of range. Skip this network.')
            return True
        else:
            return False


import json
import threading
import time

import requests

from config.settings import *
from util.redis_util import RedisUtil
from util.updater_util import UpdateUtil

class ProxyUpdater:

    def __init__(self) -> None:
        self.api_url = UPDATE_API_URL
        self.update_interval = UPDATE_INTERVAL
        self.redis_client = RedisUtil.get_redis_client(REDIS_HOST, REDIS_PORT, REDIS_DB)

    def start_update_thread(self):
        t = threading.Thread(target=self.update_proxy)
        t.start()

    def update_proxy(self):
        while True:
            self._update_proxy()
            time.sleep(self.update_interval)

    def _update_proxy(self) -> int:
        count = 0
        data = None
        try:
            res = requests.get(self.api_url)
            if UPDATE_PROXY_API_DATA_FORMAT == 'json':
                data = res.json()
            else:
                data = res.text
            proxy_data = eval(UPDATE_PROXY_DATA_KEY)
            redis_pipeline = self.redis_client.pipeline()
            now_ms = int(time.time() * 1000)
            record_key = UpdateUtil.generate_record_key()
            record_proxies = record_key + ':proxies'
            record_count = record_key + ':count'
            proxy_expire_prefix = UpdateUtil.generate_expire_prefix()
            proxy_track_key = UpdateUtil.generate_track_key()
            for proxy_info in proxy_data:
                try:
                    if ProxyUpdater.filter_proxy(proxy_info):
                        proxy = ProxyUpdater.parse_proxy(proxy_info)
                        for key in UPDATE_QUEUE_KEY:
                            redis_pipeline.zadd(REST_QUEUE_PREFIX + key, {proxy: 3})

                        proxy_expire_key = proxy_expire_prefix + proxy
                        redis_pipeline.set(proxy_expire_key, json.dumps(proxy_info))
                        redis_pipeline.expire(proxy_expire_key, UPDATE_EXPIRE_TIME_DEFAULT)
                        redis_pipeline.hset(proxy_track_key, proxy, str(now_ms))
                        # record
                        redis_pipeline.hset(record_proxies, proxy, time.time_ns())
                        redis_pipeline.incrby(record_count)
                        count = count + 1
                    else:
                        print('error proxy: ' + str(proxy_info))
                except Exception as e:
                    print('update proxy_info[{}] error: {}'.format(str(proxy_info), str(e)))
            redis_pipeline.expire(record_proxies, UPDATE_PROXY_RECORD_EXPIRE)
            redis_pipeline.expire(record_count, UPDATE_PROXY_RECORD_EXPIRE)
            redis_pipeline.execute()
            print('TestProxyUpdater get {} proxy from TEST API'.format(count))
            return count
        except Exception as e:
            print('update TEST proxy_data[{}] error: {}'.format(str(data), str(e)))
            return count

    @staticmethod
    def parse_proxy(proxy_info):
        return eval(UPDATE_PROXY_PARSE_METHOD)

    @staticmethod
    def filter_proxy(proxy_info):
        return eval(UPDATE_PROXY_FILTER_METHOD)


import heapq
import re
from enum import Enum
from Parser import EarleyParser as Parser
from Fuzzer import LimitFuzzer as LimitFuzzer
import random
MAX_TRIES_FOR_ABSTRACTION = 100

class PRes(str, Enum):
    success = 'SUCCESS'
    failed = 'FAILED'
    invalid = 'INVALID'
    timeout = 'TIMEOUT'

def is_nt(symbol):
    return symbol and (symbol[0], symbol[-1]) == ('<', '>')

def tree_to_str(tree):
    expanded = []
    to_expand = [tree]
    while to_expand:
        (key, children, *rest), *to_expand = to_expand
        if is_nt(key):
            #assert children # not necessary
            to_expand = list(children) + list(to_expand)
        else:
            assert not children
            expanded.append(key)
    return ''.join(expanded)

def get_child(tree, path):
    if not path: return tree
    cur, *path = path
    return get_child(tree[1][cur], path)

def replace_tree_node(node, path, newnode):
    if not path:
        assert node[0] == newnode[0], ("%s != %s") % (node[0], newnode[0])
        return newnode
    name, children, *rest = node
    hd, *subpath = path
    assert hd < len(children)
    new_children = []
    for i,c in enumerate(children):
        if i == hd:
            c_ = replace_tree_node(c, subpath, newnode)
        else:
            c_ = c
        new_children.append(c_)
    return (name, new_children, *rest)

def count_leaves(node):
    name, children, *_ = node
    if not children:
        return 1
    return sum(count_leaves(i) for i in children)

#=======================REDUCER=================================================

ECOUNT = 0
def add_to_pq(tup, q):
    global ECOUNT
    dtree, F_path = tup
    stree = get_child(dtree, F_path)
    n =  count_leaves(dtree)
    m =  count_leaves(stree)
    # heap smallest first
    heapq.heappush(q, (n, m, -ECOUNT, tup))
    ECOUNT += 1

def reset_pq():
    global ECOUNT
    ECOUNT = 0

def nt_group(tree, all_nodes=None):
    if all_nodes is None: all_nodes = {}
    name, children, *_ = tree
    if not is_nt(name): return
    all_nodes.setdefault(name, []).append(tree)
    for c in children:
        nt_group(c, all_nodes)
    return all_nodes

def compatible_nodes(tree, grammar):
    key, children, *_ = tree
    # Here is the first choice. Do we restrict ourselves to only children of the tree
    # or do we allow all nodes in the original tree? given in all_nodes?
    lst = nt_group(tree)
    node_lst = [(i, n) for i,n in enumerate(lst[key])]

    # insert empty if the grammar allows it as the first element
    if [] in grammar[key]: node_lst.insert(0, (-1, (key, [])))
    return node_lst

def is_token(val):
    assert val != '<>'
    assert (val[0], val[-1]) == ('<', '>')
    if val[1].isupper(): return True
    #if val[1] == '_': return val[2].isupper() # token derived.
    return False

def reduction(tree, grammar, predicate):
    first_tuple = (tree, [])
    p_q = []
    add_to_pq(first_tuple, p_q)

    ostr = tree_to_str(tree)
    assert predicate(ostr) == PRes.success
    failed_set = {ostr: True}

    min_tree, min_tree_size = tree, count_leaves(tree)
    while p_q:
        # extract the tuple
        _n, _m, _ec, (dtree, F_path) = heapq.heappop(p_q)
        stree = get_child(dtree, F_path)
        skey, schildren = stree
        found = False
        # we now want to replace stree with alternate nodes.
        for i, node in compatible_nodes(stree, grammar):
            # replace with current (copy).
            ctree = replace_tree_node(dtree, F_path, node)
            if ctree is None: continue # same node
            v = tree_to_str(ctree)
            if v in failed_set: continue
            failed_set[v] = predicate(v) # we ignore PRes.invalid results
            if failed_set[v] == PRes.success:
                found = True
                ctree_size = count_leaves(ctree)
                if ctree_size < min_tree_size: min_tree, min_tree_size = ctree, ctree_size

                if v not in failed_set:
                    print(v)
                t = (ctree, F_path)
                assert get_child(ctree, F_path) is not None
                add_to_pq(t, p_q)

        # The CHOICE here is that we explore the children if and only if we fail
        # to find a node that can replace the current
        if found: continue
        if is_token(skey): continue # do not follow children TOKEN optimization
        for i, child in enumerate(schildren):
            if not is_nt(child[0]): continue
            assert get_child(tree=dtree, path=F_path + [i]) is not None
            t = (dtree, F_path + [i])
            add_to_pq(t, p_q)
    return min_tree

#=======================GENERALIZER=============================================

def generalize(tree, path, known_paths, grammar, predicate):
    node = get_child(tree, path)
    if not is_nt(node[0]): return known_paths
    if can_abstract(tree, path, known_paths, grammar, predicate):
        known_paths.append(path)
        return known_paths
    for i,child in enumerate(node[1]):
        ps = generalize(tree, path + [i], known_paths, grammar, predicate)
    return known_paths

def can_abstract(tree, path, known_paths, grammar, predicate):
    i = 0
    while (i < MAX_TRIES_FOR_ABSTRACTION):
        t = replace_all_paths_with_generated_values(tree, known_paths + [path], grammar)
        s = tree_to_str(t)
        if predicate(s) == PRes.failed:
            return False
        elif predicate(s) == PRes.invalid:
            continue
        i += 1
    return True

def replace_all_paths_with_generated_values(tree, paths, grammar):
    my_tree = tree
    for p in paths:
        my_tree = replace_path_with_generated_value(my_tree, p, grammar)
    return my_tree

def replace_path_with_generated_value(tree, path, grammar):
    node = get_child(tree, path)
    s, gnode = generate_random_value(grammar, node[0])
    t = replace_tree_node(tree, path, gnode)
    return t

# Given a key, generate a random value for that key using the grammar.
def generate_random_value(grammar, key):
    fuzzer = LimitFuzzer(grammar)
    s = fuzzer.fuzz(key)
    return (s, fuzzer._s)

def tree_to_str_a(tree):
    name, children, *general_ = tree
    if not is_nt(name): return name
    if is_node_abstract(tree):
        return name
    return ''.join([tree_to_str_a(c) for c in children])

def mark_concrete_r(tree):
    name, children, *abstract_a = tree
    abstract = {'abstract': False} if not abstract_a else abstract_a[0]
    return (name, [mark_concrete_r(c) for c in children], abstract)

def mark_path_abstract(tree, path):
    name, children = get_child(tree, path)
    new_tree = replace_tree_node(tree, path, (name, children, {'abstract': True}))
    return new_tree

def get_abstract_tree(tree, paths):
    for path in paths:
        tree = mark_path_abstract(tree, path)
    return mark_concrete_r(tree)

def is_node_abstract(node):
    name, children, *abstract_a = node
    if not abstract_a:
        return True
    else:
        return abstract_a[0]['abstract']

#=======================GENERATOR===============================================

def abstract_to_concrete(abstract_tree, grammar):
    if is_node_abstract(abstract_tree):
        s, t = generate_random_value(grammar, abstract_tree[0])
        return t
    res = []
    for child in abstract_tree[1]:
        t = abstract_to_concrete(child, grammar)
        res.append(t)
    return (abstract_tree[0], res)

def generate_testcase(abstract_tree, grammar):
    tv = abstract_to_concrete(abstract_tree, grammar)
    s = tree_to_str(tv)
    return s

#=======================DRIVER===============================================
if __name__ == '__main__':
    random.seed(0)
    EXPR_GRAMMAR = {'<start>': [['<expr>']],
     '<expr>': [['<term>', ' + ', '<expr>'],
      ['<term>', ' - ', '<expr>'],
      ['<term>']],
     '<term>': [['<factor>', ' * ', '<term>'],
      ['<factor>', ' / ', '<term>'],
      ['<factor>']],
     '<factor>': [['+', '<factor>'],
      ['-', '<factor>'],
      ['(', '<expr>', ')'],
      ['<integer>', '.', '<integer>'],
      ['<integer>']],
     '<integer>': [['<digit>', '<integer>'], ['<digit>']],
     '<digit>': [['0'], ['1'], ['2'], ['3'], ['4'], ['5'], ['6'], ['7'], ['8'], ['9']]}

    def expr_double_paren(inp):
        if re.match(r'.*[(][(].*[)][)].*', inp):
            return PRes.success
        return PRes.failed

    my_input = '1 + ((2 * 3 / 4))'
    expr_parser = Parser(EXPR_GRAMMAR, start_symbol='<start>', canonical=True)
    parsed_expr = list(expr_parser.parse(my_input))[0]
    er = reduction(parsed_expr, EXPR_GRAMMAR, expr_double_paren)

    vals = generalize(er, [], [], EXPR_GRAMMAR, expr_double_paren)
    ta = get_abstract_tree(er, vals)
    s = tree_to_str_a(ta)
    print(s)
    for i in range(10):
        v = generate_testcase(ta, EXPR_GRAMMAR)
        print(v)
    print()
    print()
    def expr_multiply(inp):
        if re.match(r'.*[*].*', inp):
            return PRes.success
        return PRes.failed

    er = reduction(parsed_expr, EXPR_GRAMMAR, expr_multiply)

    vals = generalize(er, [], [], EXPR_GRAMMAR, expr_multiply)
    ta = get_abstract_tree(er, vals)
    s = tree_to_str_a(ta)
    print(s)
    for i in range(10):
        v = generate_testcase(ta, EXPR_GRAMMAR)
        print(v)