from ctypes import *
from ctypes.wintypes import *
from manual_score import LiveScore
import psutil
import pymem
import sys
import time
import keyboard

OpenProcess = windll.kernel32.OpenProcess
ReadProcessMemory = windll.kernel32.ReadProcessMemory
ReadProcessMemory.argtypes = HANDLE, LPCVOID, LPVOID, c_size_t, POINTER(c_size_t)
ReadProcessMemory.restype = BOOL
CloseHandle = windll.kernel32.CloseHandle

PROCESS_ALL_ACCESS = 0x1F0FFF

def forwardPtr(processHandle, origin_addr, offsetList, res_buffer, res_size):
    ptr_addr = origin_addr
    buffer = c_void_p()
    bufferSize = sizeof(buffer)
    bytesRead = c_ulonglong()
    ReadProcessMemory(processHandle, ptr_addr, byref(buffer), bufferSize, byref(bytesRead))
    for offset in offsetList:
        ptr_addr = buffer.value + offset
        if offset == offsetList[-1]:
            ReadProcessMemory(processHandle, ptr_addr, byref(res_buffer), res_size, byref(bytesRead))
        else:
            ReadProcessMemory(processHandle, ptr_addr, byref(buffer), bufferSize, byref(bytesRead))
    return res_buffer.value

class AutoScore:

    def __init__(self, init_p1=0, init_p2=0):
        self.pid = None
        self.base = None
        self.live_score = LiveScore(init_p1, init_p2)
        self.rounds_age = 0
        self.time_ref = time.time()
        self.prev_rounds = (0, 0)
        self.prev_side = None
        self.side_t_ref = None
        self.history = {}
        self.most_recent_opponent = None
        self.updateProcess()
        self.updateScore()
        if not self.most_recent_opponent:
            self.live_score.writeToFile(side=self.getPlayerSide())
            print(self.live_score.__repr__(side=self.getPlayerSide()))

    def loadHistory(self, opponentName):
        if opponentName in self.history.keys():
            print("Loading existing record with {}".format(opponentName))
            self.live_score = self.history[opponentName]
        else:
            print("Creating new record with {}".format(opponentName))
            self.live_score = LiveScore()
            self.history[opponentName] = self.live_score
        self.live_score.writeToFile(side=self.getPlayerSide())
        print(self.live_score.__repr__(side=self.getPlayerSide()))

    def saveHistory(self):
        op_name = self.readOpponentName()
        if op_name:
            self.history[op_name] = self.live_score
        else:
            self.history[self.most_recent_opponent] = self.live_score

    def checkNewOpponent(self):
        op_name = self.readOpponentName()
        if not op_name:
            return False
        if not op_name in self.history.keys():
            return True
        if self.history[op_name] == self.live_score:
            return False
        else:
            return True

    def resetScores(self):
        self.most_recent_opponent = None
        self.history.pop(self.readOpponentName(), None)
        self.live_score = LiveScore()

    def updateProcess(self):
        while self.pid == None:
            pids = psutil.pids()
            for p in pids:
                try:
                    ps = psutil.Process(p)
                except psutil.NoSuchProcess:
                    continue
                if "TekkenGame" in ps.name():
                    pm = pymem.Pymem("TekkenGame")
                    self.pid = int(ps.pid)
                    self.base = pm.base_address
                    self.processHandle = OpenProcess(PROCESS_ALL_ACCESS, False, self.pid)
                    pm.close_process()
                    print("Tekken 7 PID acquired:", self.pid)
                    return self.pid
            time.sleep(10)

    def closeProcess(self):
        CloseHandle(self.processHandle)

    def readPlayerSideFlag(self):
        while True:
            try:
                self.updatePID()
                buffer = c_ubyte()
                bufferSize = 1
                ptr = self.base + 0x034CE030
                ptr_chain = [0x498, 0x14]
                return forwardPtr(self.processHandle, ptr, ptr_chain, buffer, bufferSize)
            except TypeError:
                continue

    def readViewSideFlag(self):
        self.updatePID()
        bytesRead = c_ulonglong()
        buffer = c_ubyte()
        ptr = self.base + 0x034DF554
        ReadProcessMemory(self.processHandle, ptr, byref(buffer), 1, byref(bytesRead))
        return buffer.value

    def updateScoreOrientation(self, writing=True, protected=False):
        p_side = self.getPlayerSide()
        if self.prev_side == None:
            self.prev_side = p_side
            return
        if self.prev_side != p_side:
            if protected:
                if self.side_t_ref == None:
                    self.side_t_ref = time.time()
                    print("Detected side switch. Confirming in 30 seconds")
                    return
                else:
                    if time.time() - self.side_t_ref < 30: # 30 second protection
                        return
            self.prev_side = p_side
            if writing:
                if protected:
                    print("Confirm succeeded")
                print("Flipping score to match playing side")
                self.live_score.writeToFile(side=p_side)
                print(self.live_score.__repr__(side=p_side))
            self.side_t_ref = None
            return
        else:
            if self.side_t_ref != None:
                print("Confirm failed")
            self.side_t_ref = None

    def readOpponentName(self):
        confirm_name = None
        ptr = self.base + 0x34D55A0
        ptr_chain = [0x0, 0x8, 0x11C]
        while True:
            name = self.readString(ptr, ptr_chain, max_len=256)
            if name == "NOT_LOGGED_IN" or name == "":
                time.sleep(0.01)
                return None
            else:
                if name == confirm_name:
                    self.most_recent_opponent = name
                    return name
                else:
                    confirm_name = name
                    time.sleep(0.01)
                    continue

    def readString(self, ptr, ptr_chain, max_len=256):
        result = ''
        ptr_chain_copy = ptr_chain.copy()
        for i in range(max_len):
            while True:
                try:
                    self.updatePID()
                    buffer = c_ubyte()
                    bufferSize = 1
                    char_val = forwardPtr(self.processHandle, ptr, ptr_chain_copy, buffer, bufferSize)
                    if char_val == 0:
                        return result
                    result += chr(char_val)
                    break
                except TypeError:
                    continue
            ptr_chain_copy[-1] += 1
        return result

    def getPlayerSide(self):
        return int(self.readPlayerSideFlag() != self.readViewSideFlag())

    def readRounds(self):
        while True:
            try:
                self.updatePID()
                bytesRead = c_ulonglong()
                p1_rounds = c_ubyte()
                p2_rounds = c_ubyte()
                p1_ptr = self.base + 0x34CD500
                p2_ptr = self.base + 0x34CD5F0
                ReadProcessMemory(self.processHandle, p1_ptr, byref(p1_rounds), 1, byref(bytesRead))
                ReadProcessMemory(self.processHandle, p2_ptr, byref(p2_rounds), 1, byref(bytesRead))
            except TypeError:
                continue
            if self.readPlayerSideFlag() == 0:
                return p1_rounds.value, p2_rounds.value
            else:
                return p2_rounds.value, p1_rounds.value

    def updateScore(self):
        if self.checkNewOpponent():
            self.loadHistory(self.readOpponentName())
            self.updateScoreOrientation(writing=False, protected=False)
        self.updateScoreOrientation(protected=True)
        min_age = 3 # seconds
        o_p1, o_p2 = self.prev_rounds
        n_p1, n_p2 = self.prev_rounds = self.readRounds()
        if (o_p1, o_p2) != (n_p1, n_p2):
            if (n_p1 == 3) ^ (n_p2 == 3):
                if time.time() - self.time_ref >= min_age:
                    if n_p1 == 3 and o_p1 == 2 and n_p2 == o_p2:
                        self.live_score.incrementScore(player=0)
                        self.saveHistory()
                        self.live_score.writeToFile(side=self.getPlayerSide())
                        print(self.live_score.__repr__(side=self.getPlayerSide()))
                    elif n_p2 == 3 and o_p2 == 2 and n_p1 == o_p1:
                        self.live_score.incrementScore(player=1)
                        self.saveHistory()
                        self.live_score.writeToFile(side=self.getPlayerSide())
                        print(self.live_score.__repr__(side=self.getPlayerSide()))
            self.time_ref = time.time()

    def updatePID(self):
        if not psutil.pid_exists(self.pid):
            print("Game process has died. Waiting for restart..")
            self.closeProcess()
            self.pid = None
            self.updateProcess()

    def getLiveScore(self):
        return self.live_score

def event_reset(event):
    global auto_score
    print("Resetting score..")
    auto_score.resetScores()
    auto_score.getLiveScore().writeToFile()
    print(auto_score.getLiveScore().__repr__(side=auto_score.getPlayerSide()))

init_p1 = 0
init_p2 = 0
if len(sys.argv) == 3:
    init_p1 = int(sys.argv[1])
    init_p2 = int(sys.argv[2])
try:
    auto_score = AutoScore(init_p1, init_p2)
except KeyboardInterrupt:
    print("Exiting..")
    sys.exit()

if __name__ == "__main__":
    keyboard.on_press_key('f9', event_reset, suppress=False)
    try:
        while True:
            auto_score.updateScore()
            time.sleep(0.1)
    except KeyboardInterrupt:
        print("Exiting..")
        with open("live_score", 'w') as f:
            f.write("")


from arelle import (XmlUtil, XbrlConst)
from arelle.ModelDtsObject import ModelResource
from arelle.ModelValue import qname

class ModelRenderingResource(ModelResource):
    def __init__(self, modelDocument):
        super(ModelRenderingResource, self).__init__(modelDocument)

class ModelTable(ModelRenderingResource):
    def __init__(self, modelDocument):
        super(ModelTable, self).__init__(modelDocument)

    @property
    def propertyView(self):
        return (("id", self.id),
                ("label", self.xlinkLabel))

    def __repr__(self):
        return ("table[{0}]{1})".format(self.objectId(),self.propertyView))

# 2010 EU Table linkbase
class ModelAxisCoord(ModelRenderingResource):
    def __init__(self, modelDocument):
        super(ModelAxisCoord, self).__init__(modelDocument)

    @property
    def abstract(self):
        return self.get("abstract") if self.get("abstract") else 'false'

    @property
    def parentChildOrder(self):
        return None

    @property
    def primaryItemQname(self):
        priItem = XmlUtil.childAttr(self, XbrlConst.euRend, "primaryItem", "name")
        if priItem is not None:
            return self.prefixedNameQname(priItem)
        return None

    @property
    def explicitDims(self):
        return {(self.prefixedNameQname(e.get("dimension")),
                 self.prefixedNameQname(e.get("value")))
                for e in XmlUtil.children(self, XbrlConst.euRend, "explicitDimCoord")}

    @property
    def instant(self):
        return XmlUtil.childAttr(self, XbrlConst.euRend, "timeReference","instant")

    @property
    def propertyView(self):
        explicitDims = self.explicitDims
        return (("id", self.id),
                ("xlink:label", self.xlinkLabel),
                ("header label", self.genLabel()),
                ("header doc", self.genLabel(role="http://www.xbrl.org/2008/role/documentation")),
                ("header code", self.genLabel(role="http://www.eurofiling.info/role/2010/coordinate-code")),
                ("primary item", self.primaryItemQname),
                ("dimensions", "({0})".format(len(explicitDims)),
                  tuple((str(dim),str(mem)) for dim,mem in sorted(explicitDims)))
                  if explicitDims else (),
                ("abstract", self.abstract),
                ("instant", self.instant))

    def __repr__(self):
        return ("axisCoord[{0}]{1})".format(self.objectId(),self.propertyView))

# 2011 Table linkbase
class ModelExplicitAxisMember(ModelRenderingResource):
    def __init__(self, modelDocument):
        super(ModelExplicitAxisMember, self).__init__(modelDocument)

    @property
    def abstract(self):
        if self.localName in ("explicitAxis", "aspectRuleAxis"):
            return 'false'
        elif self.get("abstract") == 'true':
            return 'true'
        return 'false'

    @property
    def parentChildOrder(self):
        return self.get("parentChildOrder")

    @property
    def primaryItemQname(self):
        conceptRule = XmlUtil.child(self, XbrlConst.formula, "concept")
        if conceptRule is not None:
            qnameElt = XmlUtil.child(conceptRule, XbrlConst.formula, "qname")
            if qnameElt is not None:
                return qname(qnameElt, qnameElt.text)
        return None

    @property
    def explicitDims(self):
        return {(self.prefixedNameQname(e.get("dimension")),
                 self.prefixedNameQname(XmlUtil.text(qn)))
                for e in XmlUtil.children(self, XbrlConst.formula, "explicitDimension")
                for m in XmlUtil.children(e, XbrlConst.formula, "member")
                for qn in XmlUtil.children(m, XbrlConst.formula, "qname")}

    @property
    def instant(self):
        return XmlUtil.childAttr(self, XbrlConst.euRend, "timeReference","instant")

    @property
    def propertyView(self):
        explicitDims = self.explicitDims
        return (("id", self.id),
                ("xlink:label", self.xlinkLabel),
                ("header label", self.genLabel()),
                ("header doc", self.genLabel(role="http://www.xbrl.org/2008/role/documentation")),
                ("header code", self.genLabel(role="http://www.eurofiling.info/role/2010/coordinate-code")),
                ("primary item", self.primaryItemQname),
                ("dimensions", "({0})".format(len(explicitDims)),
                  tuple((str(dim),str(mem)) for dim,mem in sorted(explicitDims)))
                  if explicitDims else (),
                ("abstract", self.abstract),
                ("instant", self.instant))

    def __repr__(self):
        return ("explicitAxisMember[{0}]{1})".format(self.objectId(),self.propertyView))

from arelle.ModelObjectFactory import elementSubstitutionModelClass
elementSubstitutionModelClass.update((
    (XbrlConst.qnEuTable, ModelTable),
    (XbrlConst.qnTableTable, ModelTable),
    (XbrlConst.qnEuAxisCoord, ModelAxisCoord),
    (XbrlConst.qnTableAspectRuleAxis, ModelExplicitAxisMember),
    (XbrlConst.qnTableAspectRuleAxisMember, ModelExplicitAxisMember),
    (XbrlConst.qnTableFilterAxis, ModelExplicitAxisMember),
     ))


"""
Implemention of a ListEditor demo plugin for Traits UI demo program

This demo shows each of the four styles of ListEditor.
"""

from traits.api import HasTraits, List, Str
from traitsui.api import Item, Group, View

#-------------------------------------------------------------------------------
#  Demo Class
#-------------------------------------------------------------------------------

class ListEditorDemo ( HasTraits ):
    """ This class specifies the details of the BooleanEditor demo.
    """

    # The Trait to be displayed in the editor
    play_list = List( Str, ["The Merchant of Venice", "Hamlet", "MacBeth"])

    # Items are used to define display; one per editor style.
    list_group = Group( Item('play_list', style='simple', label='Simple'),
                        Item('_'),
                        Item('play_list', style='custom', label='Custom'),
                        Item('_'),
                        Item('play_list', style='text', label='Text'),
                        Item('_'),
                        Item('play_list',
                              style='readonly',
                              label='ReadOnly'))

    # Demo view
    view1 = View( list_group,
                  title = 'ListEditor',
                  buttons = ['OK'],
                  height=400,
                  width=400 )

# Create the demo:
popup = ListEditorDemo()

# Run the demo (if invoked from the command line):
if __name__ == '__main__':
    popup.configure_traits()


######## This script is intended to be run within the YOLO directory ########
######## It will take a directory of images and run them through the YOLO ########
######## model to detect objects in the images. ########
######## For further information visit the github in https://github.com/ultralytics/yolov5 ########

import argparse
import os
import sys
from pathlib import Path

import cv2
import torch
import torch.backends.cudnn as cudnn

FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]  # YOLOv5 root directory
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))  # add ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative

from models.common import DetectMultiBackend
from utils.datasets import IMG_FORMATS, VID_FORMATS, LoadImages, LoadStreams
from utils.general import (LOGGER, check_file, check_img_size, check_imshow, check_requirements, colorstr,
                            increment_path, non_max_suppression, print_args, scale_coords, strip_optimizer, xyxy2xywh)
from utils.plots import Annotator, colors, save_one_box
from utils.torch_utils import select_device, time_sync

from utils.augmentations import letterbox
import numpy as np

class detector:
    def __init__(self,
        weights=ROOT / 'yolov5s.pt',  # model.pt path(s)
        source=ROOT / 'data/images',  # file/dir/URL/glob, 0 for webcam
        data=ROOT / 'data/coco128.yaml',  # dataset.yaml path
        imgsz=(640, 640),  # inference size (height, width)
        conf_thres=0.25,  # confidence threshold
        iou_thres=0.45,  # NMS IOU threshold
        max_det=1000,  # maximum detections per image
        device='',  # cuda device, i.e. 0 or 0,1,2,3 or cpu
        view_img=False,  # show results
        save_txt=False,  # save results to *.txt
        save_conf=False,  # save confidences in --save-txt labels
        save_crop=False,  # save cropped prediction boxes
        nosave=False,  # do not save images/videos
        classes=None,  # filter by class: --class 0, or --class 0 2 3
        agnostic_nms=False,  # class-agnostic NMS
        augment=False,  # augmented inference
        visualize=False,  # visualize features
        update=False,  # update all models
        project=ROOT / 'runs/detect',  # save results to project/name
        name='exp',  # save results to project/name
        exist_ok=False,  # existing project/name ok, do not increment
        line_thickness=3,  # bounding box thickness (pixels)
        hide_labels=False,  # hide labels
        hide_conf=False,  # hide confidences
        half=False,  # use FP16 half-precision inference
        dnn=False,  # use OpenCV DNN for ONNX inference
        ):

        self.weights=weights
        self.source=source
        self.data=data
        self.imgsz=imgsz
        self.conf_thres=conf_thres
        self.iou_thres=iou_thres
        self.max_det=max_det
        self.device=device
        self.view_img=view_img
        self.save_txt=save_txt
        self.save_conf=save_conf
        self.save_crop=save_crop
        self.nosave=nosave
        self.classes=classes
        self.agnostic_nms=agnostic_nms
        self.augment=augment
        self.visualize=visualize
        self.update=update
        self.project=project
        self.name=name
        self.exist_ok=exist_ok
        self.line_thickness=line_thickness
        self.hide_labels=hide_labels
        self.hide_conf=hide_conf
        self.half=half
        self.dnn=dnn

        """source = str(source)
        self.save_img = not nosave and not source.endswith('.txt')  # save inference images
        is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS)
        is_url = source.lower().startswith(('rtsp://', 'rtmp://', 'http://', 'https://'))
        webcam = source.isnumeric() or source.endswith('.txt') or (is_url and not is_file)
        if is_url and is_file:
            source = check_file(source)  # download"""

        # Load model
        self.device = select_device(device)
        self.model = DetectMultiBackend(self.weights, device=self.device, dnn=self.dnn, data=self.data)
        self.stride, self.names, self.pt, jit, onnx, engine = self.model.stride, self.model.names, self.model.pt, self.model.jit, self.model.onnx, self.model.engine
        self.imgsz = check_img_size(imgsz, s=self.stride)  # check image size

        # Half
        half &= (self.pt or jit or onnx or engine) and self.device.type != 'cpu'  # FP16 supported on limited backends with CUDA
        if self.pt or jit:
            self.model.model.half() if half else self.model.model.float()

    @torch.no_grad()
    def detect(self,data):

        """# Dataloader
        if self.webcam:
            self.view_img = check_imshow()
            cudnn.benchmark = True  # set True to speed up constant image size inference
            dataset = LoadStreams(self.source, img_size=self.imgsz, stride=self.stride, auto=self.pt)
            bs = len(dataset)  # batch_size
        else:
            dataset = LoadImages(self.source, img_size=self.imgsz, stride=self.stride, auto=self.pt)
            bs = 1  # batch_size"""

        #####
        im = data

        img = letterbox(im, self.imgsz, stride=self.stride, auto=self.pt)[0]

        # Convert
        img = img.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
        im = np.ascontiguousarray(img)

        bs = 1  # batch_size

        self.vid_path, self.vid_writer = [None] * bs, [None] * bs

        # Run inference
        self.model.warmup(imgsz=(1 if self.pt else self.bs, 3, *self.imgsz), half=self.half)  # warmup
        dt, seen = [0.0, 0.0, 0.0], 0
        #for path, im, im0s, vid_cap, s in dataset:
        t1 = time_sync()
        im = torch.from_numpy(im).to(self.device)
        im = im.half() if self.half else im.float()  # uint8 to fp16/32
        im /= 255  # 0 - 255 to 0.0 - 1.0
        if len(im.shape) == 3:
            im = im[None]  # expand for batch dim
        t2 = time_sync()
        dt[0] += t2 - t1

        # Inference
        visualize = False #increment_path(self.save_dir / Path(path).stem, mkdir=True) if visualize else False
        pred = self.model(im, augment=self.augment, visualize=visualize)
        t3 = time_sync()
        dt[1] += t3 - t2

        # NMS
        pred = non_max_suppression(pred, self.conf_thres, self.iou_thres, self.classes, self.agnostic_nms, max_det=self.max_det)
        dt[2] += time_sync() - t3

        # Second-stage classifier (optional)
        # pred = utils.general.apply_classifier(pred, classifier_model, im, im0s)

        # Process predictions
        for i, det in enumerate(pred):  # per image
            seen += 1
            #im0, frame = im0s.copy()

            im0 = data

            #p = Path(p)  # to Path
            #save_path = str(self.save_dir / p.name)  # im.jpg
            #txt_path = str(self.save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}')  # im.txt
            #s += '%gx%g ' % im.shape[2:]  # print string
            #gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
            #imc = im0.copy() if self.save_crop else im0  # for save_crop
            annotator = Annotator(im0, line_width=self.line_thickness, example=str(self.names))
            if len(det):
                # Rescale boxes from img_size to im0 size
                det[:, :4] = scale_coords(im.shape[2:], det[:, :4], im0.shape).round()

                # Print results
                #for c in det[:, -1].unique():
                #    n = (det[:, -1] == c).sum()  # detections per class
                #    s += f"{n} {self.names[int(c)]}{'s' * (n > 1)}, "  # add to string

                # Write results
                for *xyxy, conf, cls in reversed(det):
                    """if self.save_txt:  # Write to file
                        xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
                        line = (cls, *xywh, conf) if self.save_conf else (cls, *xywh)  # label format
                        with open(txt_path + '.txt', 'a') as f:
                            f.write(('%g ' * len(line)).rstrip() % line + '\n')"""

                    # Add bbox to image
                    c = int(cls)  # integer class
                    label = None if self.hide_labels else (self.names[c] if self.hide_conf else f'{self.names[c]} {conf:.2f}')
                    annotator.box_label(xyxy, label, color=colors(c, True))
                    #if self.save_crop:
                    #    save_one_box(xyxy, imc, file=self.save_dir / 'crops' / self.names[c] / f'{p.stem}.jpg', BGR=True)

            # Stream results
            im0 = annotator.result()
        return im0


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

from django.db import models, migrations
import datetime

class Migration(migrations.Migration):

    dependencies = [
        ('yearbook', '0002_auto_20141212_1528'),
    ]

    operations = [
        migrations.AddField(
            model_name='pessoa',
            name='ferias_fim',
            field=models.DateTimeField(default=datetime.date(2014, 12, 12), verbose_name=b'ferias_fim', blank=True),
            preserve_default=False,
        ),
        migrations.AddField(
            model_name='pessoa',
            name='ferias_inicio',
            field=models.DateTimeField(default=datetime.date(2014, 12, 12), verbose_name=b'ferias_inicio', blank=True),
            preserve_default=False,
        ),
    ]


import pathlib
from datetime import date

import pytest
from pytest_mock.plugin import MockerFixture
from sqlglot import parse_one

import sqlmesh.core.constants
from sqlmesh.core.config import Config, ModelDefaultsConfig
from sqlmesh.core.context import Context
from sqlmesh.core.dialect import parse
from sqlmesh.core.model import load_model
from sqlmesh.core.plan import BuiltInPlanEvaluator, Plan
from sqlmesh.utils.date import yesterday_ds
from sqlmesh.utils.errors import ConfigError
from tests.utils.test_filesystem import create_temp_file

def test_global_config():
    context = Context(path="examples/sushi")
    assert context.dialect == "duckdb"
    assert context.physical_schema == "sqlmesh"

def test_named_config():
    context = Context(path="examples/sushi", config="local_config")
    assert context.dialect == "duckdb"

def test_invalid_named_config():
    with pytest.raises(
        ConfigError,
        match="Config 'blah' was not found.",
    ):
        Context(path="examples/sushi", config="blah")

def test_missing_named_config():
    with pytest.raises(ConfigError, match=r"Config 'imaginary_config' was not found."):
        Context(path="examples/sushi", config="imaginary_config")

def test_config_precedence():
    context = Context(path="examples/sushi", dialect="spark", physical_schema="test")
    assert context.dialect == "spark"
    assert context.physical_schema == "test"

    # Context parameters take precedence over config
    config = Config(model_defaults=ModelDefaultsConfig(dialect="presto"), physical_schema="dev")
    context = Context(path="examples/sushi", dialect="spark", physical_schema="test", config=config)
    assert context.dialect == "spark"
    assert context.physical_schema == "test"

def test_config_parameter():
    config = Config(model_defaults=ModelDefaultsConfig(dialect="presto"), physical_schema="dev")
    context = Context(path="examples/sushi", config=config)
    assert context.dialect == "presto"
    assert context.physical_schema == "dev"

def test_config_not_found():
    with pytest.raises(
        ConfigError,
        match=r"^nonexistent/directory is not a directory.*",
    ):
        Context(path="nonexistent/directory", config="local_config")

def test_custom_macros(sushi_context):
    assert "add_one" in sushi_context.macros

def test_dag(sushi_context):
    assert set(sushi_context.dag.upstream("sushi.customer_revenue_by_day")) == {
        "sushi.items",
        "sushi.orders",
        "sushi.order_items",
    }

def test_render(sushi_context, assert_exp_eq):
    snapshot = sushi_context.snapshots["sushi.waiter_revenue_by_day"]

    assert_exp_eq(
        sushi_context.render(
            snapshot.model,
            start=date(2021, 1, 1),
            end=date(2021, 1, 1),
            latest=date(2021, 1, 1),
            expand=True,
        ),
        f"""
        SELECT
          CAST(o.waiter_id AS INT) AS waiter_id, /* Waiter id */
          CAST(SUM(oi.quantity * i.price) AS DOUBLE) AS revenue, /* Revenue from orders taken by this waiter */
          CAST(o.ds AS TEXT) AS ds /* Date */
        FROM (
          SELECT
            CAST(NULL AS INT) AS id,
            CAST(NULL AS INT) AS customer_id,
            CAST(NULL AS INT) AS waiter_id,
            CAST(NULL AS INT) AS start_ts,
            CAST(NULL AS INT) AS end_ts,
            CAST(NULL AS TEXT) AS ds
          FROM (VALUES
            (1)) AS t(dummy)
        ) AS o
        LEFT JOIN (
          SELECT
            CAST(NULL AS INT) AS id,
            CAST(NULL AS INT) AS order_id,
            CAST(NULL AS INT) AS item_id,
            CAST(NULL AS INT) AS quantity,
            CAST(NULL AS TEXT) AS ds
          FROM (VALUES
            (1)) AS t(dummy)
        ) AS oi
          ON o.ds = oi.ds AND o.id = oi.order_id
        LEFT JOIN (
          SELECT
            CAST(NULL AS INT) AS id,
            CAST(NULL AS TEXT) AS name,
            CAST(NULL AS DOUBLE) AS price,
            CAST(NULL AS TEXT) AS ds
          FROM (VALUES
            (1)) AS t(dummy)
        ) AS i
          ON oi.ds = i.ds AND oi.item_id = i.id
        WHERE
          o.ds <= '2021-01-01' AND o.ds >= '2021-01-01'
        GROUP BY
          o.waiter_id,
          o.ds
        HAVING
          CAST(o.ds AS TEXT) <= '2021-01-01' AND CAST(o.ds AS TEXT) >= '2021-01-01'
        """,
    )
    assert_exp_eq(
        sushi_context.render(
            snapshot.name,
            start=date(2021, 1, 1),
            end=date(2021, 1, 1),
            latest=date(2021, 1, 1),
        ),
        f"""
        SELECT
          CAST(o.waiter_id AS INT) AS waiter_id, -- Waiter id
          CAST(SUM(oi.quantity * i.price) AS DOUBLE) AS revenue, -- Revenue from orders taken by this waiter
          CAST(o.ds AS TEXT) AS ds, -- Date
        FROM {sushi_context.snapshots['sushi.orders'].table_name()} AS o
        LEFT JOIN {sushi_context.snapshots['sushi.order_items'].table_name()} AS oi
          ON o.ds = oi.ds AND o.id = oi.order_id
        LEFT JOIN {sushi_context.snapshots['sushi.items'].table_name()} AS i
          ON oi.ds = i.ds AND oi.item_id = i.id
        WHERE
          o.ds <= '2021-01-01' AND o.ds >= '2021-01-01'
        GROUP BY
          o.waiter_id,
          o.ds
        HAVING
          CAST(o.ds AS TEXT) <= '2021-01-01' AND CAST(o.ds AS TEXT) >= '2021-01-01'
        """,
    )

    # unpushed render still works
    unpushed = Context(path="examples/sushi")
    assert_exp_eq(
        unpushed.render(snapshot.name),
        f"""
        SELECT
          CAST(o.waiter_id AS INT) AS waiter_id, /* Waiter id */
          CAST(SUM(oi.quantity * i.price) AS DOUBLE) AS revenue, /* Revenue from orders taken by this waiter */
          CAST(o.ds AS TEXT) AS ds /* Date */
        FROM (
          SELECT
            CAST(NULL AS INT) AS id,
            CAST(NULL AS INT) AS customer_id,
            CAST(NULL AS INT) AS waiter_id,
            CAST(NULL AS INT) AS start_ts,
            CAST(NULL AS INT) AS end_ts,
            CAST(NULL AS TEXT) AS ds
          FROM (VALUES
            (1)) AS t(dummy)
        ) AS o
        LEFT JOIN (
          SELECT
            CAST(NULL AS INT) AS id,
            CAST(NULL AS INT) AS order_id,
            CAST(NULL AS INT) AS item_id,
            CAST(NULL AS INT) AS quantity,
            CAST(NULL AS TEXT) AS ds
          FROM (VALUES
            (1)) AS t(dummy)
        ) AS oi
          ON o.ds = oi.ds AND o.id = oi.order_id
        LEFT JOIN (
          SELECT
            CAST(NULL AS INT) AS id,
            CAST(NULL AS TEXT) AS name,
            CAST(NULL AS DOUBLE) AS price,
            CAST(NULL AS TEXT) AS ds
          FROM (VALUES
            (1)) AS t(dummy)
        ) AS i
          ON oi.ds = i.ds AND oi.item_id = i.id
        WHERE
          o.ds <= '1970-01-01' AND o.ds >= '1970-01-01'
        GROUP BY
          o.waiter_id,
          o.ds
        HAVING
          CAST(o.ds AS TEXT) <= '1970-01-01' AND CAST(o.ds AS TEXT) >= '1970-01-01'
        """,
    ),

def test_diff(sushi_context: Context, mocker: MockerFixture):
    mock_console = mocker.Mock()
    sushi_context.console = mock_console
    yesterday = yesterday_ds()
    sushi_context.run(start=yesterday, end=yesterday, latest=yesterday)

    plan_evaluator = BuiltInPlanEvaluator(
        sushi_context.state_sync, sushi_context.snapshot_evaluator
    )
    plan_evaluator._promote(
        Plan(
            context_diff=sushi_context._context_diff("prod"),
            dag=sushi_context.dag,
            state_reader=sushi_context.state_reader,
        )
    )

    sushi_context.upsert_model("sushi.customers", query=parse_one("select 1"))
    sushi_context.diff("test")
    assert mock_console.show_model_difference_summary.called

def test_evaluate_limit():
    context = Context(config=Config())

    context.upsert_model(
        load_model(
            parse(
                """
        MODEL(name with_limit, kind FULL);
        SELECT t.v as v FROM (VALUES (1), (2), (3), (4), (5)) AS t(v) LIMIT 1 + 2"""
            )
        )
    )

    assert context.evaluate("with_limit", "2020-01-01", "2020-01-02", "2020-01-02").size == 3
    assert context.evaluate("with_limit", "2020-01-01", "2020-01-02", "2020-01-02", 4).size == 3
    assert context.evaluate("with_limit", "2020-01-01", "2020-01-02", "2020-01-02", 2).size == 2

    context.upsert_model(
        load_model(
            parse(
                """
        MODEL(name without_limit, kind FULL);
        SELECT t.v as v FROM (VALUES (1), (2), (3), (4), (5)) AS t(v)"""
            )
        )
    )

    assert context.evaluate("without_limit", "2020-01-01", "2020-01-02", "2020-01-02").size == 5
    assert context.evaluate("without_limit", "2020-01-01", "2020-01-02", "2020-01-02", 4).size == 4
    assert context.evaluate("without_limit", "2020-01-01", "2020-01-02", "2020-01-02", 2).size == 2

def test_ignore_files(mocker: MockerFixture, tmpdir):
    mocker.patch.object(
        sqlmesh.core.constants,
        "IGNORE_PATTERNS",
        [".ipynb_checkpoints/*.sql"],
    )

    models_dir = pathlib.Path("models")
    macros_dir = pathlib.Path("macros")

    ignore_model_file = create_temp_file(
        tmpdir,
        pathlib.Path(models_dir, "ignore", "ignore_model.sql"),
        "MODEL(name ignore.ignore_model); SELECT cola::bigint AS cola FROM db.other_table",
    )
    ignore_macro_file = create_temp_file(
        tmpdir,
        pathlib.Path(macros_dir, "macro_ignore.py"),
        """
from sqlmesh.core.macros import macro

@macro()
def test():
    return "test"
""",
    )
    constant_ignore_model_file = create_temp_file(
        tmpdir,
        pathlib.Path(models_dir, ".ipynb_checkpoints", "ignore_model2.sql"),
        "MODEL(name ignore_model2); SELECT cola::bigint AS cola FROM db.other_table",
    )
    create_temp_file(
        tmpdir,
        pathlib.Path(models_dir, "db", "actual_test.sql"),
        "MODEL(name db.actual_test, kind full); SELECT cola::bigint AS cola FROM db.other_table",
    )
    create_temp_file(
        tmpdir,
        pathlib.Path(macros_dir, "macro_file.py"),
        """
from sqlmesh.core.macros import macro

@macro()
def test():
    return "test"
""",
    )
    config = Config(ignore_patterns=["models/ignore/*.sql", "macro_ignore.py"])
    context = Context(path=str(tmpdir), config=config)

    assert ["db.actual_test"] == list(context.models)
    assert "test" == list(context.macros)[-1]

def test_plan_apply(sushi_context) -> None:
    plan = sushi_context.plan(
        "dev",
        start=yesterday_ds(),
        end=yesterday_ds(),
    )
    sushi_context.apply(plan)
    assert sushi_context.state_reader.get_environment("dev")


from django.urls import path
from . import views as v

app_name = 'user'
urlpatterns = [
    path('', v.index, name = "index"),
    path('<int:user_id>/', v.detail, name = "details"),
    path('create/', v.create, name = 'create'),
    path('getUser/', v.getUser, name = 'getUser'),
    path('<int:user_id>/update/',v.update, name = 'update'),
    path('<int:user_id>/delete/', v.delete, name = 'delete'),

    #Role path
    path('role/', v.role_index, name='role_index'),
    path('role/create/', v.create_role, name='create_role'),
    path('role/<int:role_id>', v.detail_role, name='role_detail'),
    path('getrole/', v.get_choice_field_role),
    path('search/', v.search),
    path('add_user_to_role/', v.add_user_to_role),
]


"""
Title: Augmented Warfighter Awareness System Display
Author: Bryan Ritchie
Team: UCCS Senior Design "Heads Up" 2020

Purpose: This file is intended as a driver for the 2.8" TFT display,
         to display the "real-time" simulated sensor data.

"""
import time
from multiprocessing import Value
import os
#from main import sensorData

RUN = 0
DEMO = 1

# Change this runmode to run as a standalone demo
# This is useful when modifying the layout of the display
RUNMODE = RUN

### <<< SCREEN LAYOUT >>> ###
#< COLUMNS >#
TITLE_START             =   10
SAMPLE_START            =   4
WARRIOR_NAME_START      =   14
WARRIOR_NAME_SPACING    =   13
DATA_NAME_START         =   4
DATA_START              =   WARRIOR_NAME_START + 3
DATA_SPACING            =   WARRIOR_NAME_SPACING
CONNECTION_STATUS_START =   4
CLIENT_MESSAGE_START    =   3

#< ROWS >#
TITLE_ROW               =   3
SAMPLE_ROW              =   TITLE_ROW + 1
WARRIOR_NAME_ROW        =   TITLE_ROW + 2
AMMO_ROW                =   TITLE_ROW + 3
WATER_ROW               =   TITLE_ROW + 4
HR_ROW                  =   TITLE_ROW + 5
SPO2_ROW                =   TITLE_ROW + 6
TEMP_ROW                =   TITLE_ROW + 7
RESP_ROW                =   TITLE_ROW + 8
WEAPON_ROW              =   TITLE_ROW + 9
CONNECTION_STATUS_ROW   =   3
CLIENT_MESSAGE_ROW      =   WEAPON_ROW + 2

#< CLIENT MESSAGE >#
MAX_MESSAGE_LENGTH      =   49
MAX_MESSAGE_LINES       =   5

### <<< SENSOR DATA >>> ###
#< SENSOR DATA ELEMENTS>#
WARRIOR_TOTAL           =   3
HIGH                    =   1
LOW                     =   0

ID_DATA                 =   0
AMMO_DATA               =   1
WATER_DATA              =   2
VITALS_DATA             =   3
HR_DATA                 =   0
SPO2_DATA               =   1
TEMP_DATA               =   2
RESP_DATA               =   3
WEAPON_DATA             =   4

# AMMO COLOR STATUS >#
# Units are in Rounds of Ammunition
AMMO_GREEN              =   270
AMMO_AMBER              =   202
AMMO_RED                =   135
AMMO_BLACK              =   68   # Black doesn't show up well, so it will be blue

# WATER COLOR STATUS >#
# Units are in Quarts
WATER_GREEN             =   3
WATER_AMBER             =   2
WATER_RED               =   1
WATER_BLACK             =   0

#< VITALS COLOR STATUS >#
#<< HR >>#
# Units: Beats per Minute   [LOW, HIGH]
HR_GREEN                =   [60, 100]
HR_AMBER                =   [50, 120]
HR_RED                  =   [40, 160]
HR_BLACK                =   [30, 200]

#<< SPO2 >>#
SPO2_GREEN              =   100
SPO2_AMBER              =   92
SPO2_RED                =   85
SPO2_BLACK              =   80

#<< RESPIRATIONS >>#
# Units: Breaths per min    [LOW, HIGH]
RESP_GREEN              =   [12, 15]
RESP_AMBER              =   [10, 18]
RESP_RED                =   [8, 21]
RESP_BLACK              =   [0, 25]

#<< TEMP >>#
# Units: Degrees F          [LOW, HIGH]
TEMP_GREEN              =   [96, 99]
TEMP_AMBER              =   [93, 101]
TEMP_RED                =   [85, 103]
TEMP_BLACK              =   [80, 104]

#< WEAPON COLOR STATUS >#
# WEAPON ON SAFE: Even Number
# WEAPON ON FIRE: Odd Number
# [SECONDS SINCE FIRING + SAFETY]
WEAPON_GREEN            = 0         # Has not fired, weapon on safe
WEAPON_AMBER            = 1         # Has not fired, weapon on fire
WEAPON_RED              = 11        # Been at least 10 seconds since last firing (weapon on safe or fire)
WEAPON_BLACK            = [2, 11]   # Been 2 - 10 just a short time since firing (weapon on safe or fire)

### <<< TFT NAVIGATION >>> ###
def moveCursor(row, col):
    print(chr(27) + "[" + str(row) + ";" + str(col) + "f")

def clearScreen():
    moveCursor(1,1)
    print(chr(27) + "[2J")

def hideCursor():
    print("\033[?25h")

# Print at a certain location on the TFT Screen
def print_at(row, col, message):
    print(chr(27) + "[" + str(row) + ";" + str(col) + "f" + message)

### <<< HANDLE TFT MESSAGE PRINTING >>> ###
class MessagePrintService(object):
    def __init__(self):
        self.messageLines = [None] * MAX_MESSAGE_LINES

        self.colorText_Mutex = Value("i", 0)

        self.messageColor   = "red"
        self.warriorColor   = "white"
        self.dataNameColor  = "white"
        self.dataColor      = "white"
        self.titleColor     = "green"
        self.borderColor    = "white"

        self.sample_counter = 0
        self.delay_counter   = 0

    def displayStructure(self):
        #< BORDER >#
        moveCursor(1,1)
        self.colorText(self.borderColor)
        print("*****************************************************")
        print("*                                                   *")
        print("*                                                   *")
        print("*                                                   *")
        print("*                                                   *")
        print("*                                                   *")
        print("*                                                   *")
        print("*                                                   *")
        print("*                                                   *")
        print("*                                                   *")
        print("*                                                   *")
        print("*                                                   *")
        print("*                                                   *")
        print("*                                                   *")
        print("*                                                   *")
        print("*                                                   *")
        print("*                                                   *")
        print("*                                                   *")
        print("*****************************************************")

        #< TITLE >#
        self.colorText(self.titleColor)
        print_at(TITLE_ROW, TITLE_START, "AUGMENTED WARFIGHTER AWARENESS SYSTEM")

        #< WARRIOR LABELS >#
        self.colorText(self.warriorColor)
        print_at(WARRIOR_NAME_ROW, WARRIOR_NAME_START,  "Warrior 1")
        print_at(WARRIOR_NAME_ROW, WARRIOR_NAME_START + WARRIOR_NAME_SPACING,  "Warrior 2")
        print_at(WARRIOR_NAME_ROW, WARRIOR_NAME_START + WARRIOR_NAME_SPACING * 2,  "Warrior 3")

        #< SENSOR LABELS >#
        self.colorText(self.dataNameColor)
        print_at(AMMO_ROW, DATA_NAME_START,     "Ammo   :")
        print_at(WATER_ROW, DATA_NAME_START,    "Water  :")
        print_at(HR_ROW, DATA_NAME_START,       "HR     :")
        print_at(SPO2_ROW, DATA_NAME_START,     "SPO2   :")
        print_at(TEMP_ROW, DATA_NAME_START,     "Temp   :")
        print_at(RESP_ROW, DATA_NAME_START,     "Resp   :")
        print_at(WEAPON_ROW, DATA_NAME_START,   "Weapon :")

    def clearMessagePane(self):
        self.colorText(self.borderColor)
        print_at(CLIENT_MESSAGE_ROW, CLIENT_MESSAGE_START, "                                                   ")
        self.colorText(self.borderColor)
        print_at(CLIENT_MESSAGE_ROW + 1, 1,                "*                                                   *")
        self.colorText(self.borderColor)
        print_at(CLIENT_MESSAGE_ROW + 2, 1,                "*                                                   *")
        self.colorText(self.borderColor)
        print_at(CLIENT_MESSAGE_ROW + 3, 1,                "*                                                   *")
        self.colorText(self.borderColor)
        print_at(CLIENT_MESSAGE_ROW + 4, 1,                "*                                                   *")

    def shiftMessagePane(self, numLines):

        # Shift Messages
        for i in range(numLines):
            for j in range(MAX_MESSAGE_LINES - 1, 0, -1):
                self.messageLines[j] = self.messageLines[j - 1]

        # Re-print messages
        for k in range(MAX_MESSAGE_LINES):
            if(self.messageLines[k] is not None):
                self.colorText(self.borderColor)
                print_at(CLIENT_MESSAGE_ROW + k, CLIENT_MESSAGE_START, "                                                  *")
                self.colorText(self.messageColor)
                print_at(CLIENT_MESSAGE_ROW + k, CLIENT_MESSAGE_START, self.messageLines[k])

    def printClientMessage(self, message):
        # Add concatenated timestamp to message
        message = "[" + str(int(time.time()) & 0xFFF) + "]: " + str(message)

        # Check length of message
        messageLength = len(message)
        messagesRequired = 1

        # If message is too long, count how many messages will be needed?
        while(messageLength > MAX_MESSAGE_LENGTH):
            messagesRequired += 1
            messageLength -= MAX_MESSAGE_LENGTH

        # Shift the message pane to make room for new messages
        self.shiftMessagePane(messagesRequired)

        # Clear the first message row (Bugfix)
        self.colorText(self.borderColor)
        print_at(CLIENT_MESSAGE_ROW, CLIENT_MESSAGE_START, "                                                  *")

        # Change the pen color to messageColor
        self.colorText(self.messageColor)

        # Create a messages list to split the messages into
        if(messagesRequired > 1):
            # Re-discover message length
            messageLength = len(message)

            # Track how many messages have been parsed
            messagesParsed = 0

            while(messagesParsed < messagesRequired and messageLength > MAX_MESSAGE_LENGTH) and messagesParsed < MAX_MESSAGE_LINES:

                # Store message to messageLines (for shifting purposes)
                self.messageLines[messagesParsed] = str(message[ messagesParsed*MAX_MESSAGE_LENGTH : (messagesParsed + 1)*MAX_MESSAGE_LENGTH ])

                # Print the message that was stored
                print_at(CLIENT_MESSAGE_ROW + messagesParsed, CLIENT_MESSAGE_START, self.messageLines[messagesParsed])

                # Iterate how many messages have been parsed
                messagesParsed += 1

                # Reduce remaining message length to be parsed
                messageLength -= MAX_MESSAGE_LENGTH

            # Check if there is any leftover message
            if(messageLength > 0 and messagesParsed < MAX_MESSAGE_LINES):
                # Store the message to messageLines (for shifting purposes)
                self.messageLines[messagesParsed] = str(message[ messagesParsed*MAX_MESSAGE_LENGTH : len(message) ])

                # Print the message that was stored
                print_at(CLIENT_MESSAGE_ROW + messagesParsed, CLIENT_MESSAGE_START, self.messageLines[messagesParsed])

        else:
            # Just print the single message
            self.messageLines[0] = str(message)
            print_at(CLIENT_MESSAGE_ROW, CLIENT_MESSAGE_START, self.messageLines[0])

    def displayData(self, connectionStatus, sensorData_packages):
        while(1):
            # Iterate through the sample data
            if(self.sample_counter < 3):

                # Apply a delay to make the changing data observable
                if(self.delay_counter >= 2000):
                    self.sample_counter += 1
                    self.delay_counter = 0 # cyclic reset
                else:
                    self.delay_counter += 1
            else:
                self.sample_counter = 0 # cyclic reset

            # redundant check (Bugfix)
            if(self.sample_counter == 3):
                self.sample_counter = 0

            # Select the current sample data
            sensorData = sensorData_packages[self.sample_counter]

            #< AMMO >#
            # Determine Color Status
            for i in range(WARRIOR_TOTAL):
                if(sensorData[i][AMMO_DATA] > AMMO_AMBER and sensorData[i][AMMO_DATA] <= AMMO_GREEN):
                    self.colorText("green")
                elif(sensorData[i][AMMO_DATA] > AMMO_RED and sensorData[i][AMMO_DATA] <= AMMO_AMBER):
                    self.colorText("yellow")
                elif(sensorData[i][AMMO_DATA] > AMMO_BLACK and sensorData[i][AMMO_DATA] <= AMMO_RED):
                    self.colorText("red")
                elif(sensorData[i][AMMO_DATA] <= AMMO_BLACK):
                    self.colorText("blue")

                # Print Data
                print_at(AMMO_ROW, DATA_START + (DATA_SPACING * i), str(sensorData[i][AMMO_DATA]) + "   ")

            #< WATER >#
            for i in range(WARRIOR_TOTAL):
                # Determine Color Status
                if(sensorData[i][WATER_DATA] > WATER_AMBER):
                    self.colorText("green")
                elif(sensorData[i][WATER_DATA] > WATER_RED and sensorData[i][WATER_DATA] <= WATER_AMBER):
                    self.colorText("yellow")
                elif(sensorData[i][WATER_DATA] > WATER_BLACK and sensorData[i][WATER_DATA] <= WATER_RED):
                    self.colorText("red")
                elif(sensorData[i][WATER_DATA] <= WATER_BLACK):
                    self.colorText("blue")

                # Print Data
                print_at(WATER_ROW, DATA_START + (DATA_SPACING * i), str(sensorData[i][WATER_DATA]) + "   ")

            #< HR >#
            for i in range(WARRIOR_TOTAL):
                # Determine Color Status
                if(sensorData[i][VITALS_DATA][HR_DATA] >= HR_GREEN[LOW] and sensorData[i][VITALS_DATA][HR_DATA] <= HR_GREEN[HIGH]):
                    self.colorText("green")
                elif((sensorData[i][VITALS_DATA][HR_DATA] >= HR_AMBER[LOW] and sensorData[i][VITALS_DATA][HR_DATA] < HR_GREEN[LOW]) or \
                    (sensorData[i][VITALS_DATA][HR_DATA] > HR_GREEN[HIGH] and sensorData[i][VITALS_DATA][HR_DATA] <= HR_AMBER[HIGH])):
                    self.colorText("yellow")
                elif((sensorData[i][VITALS_DATA][HR_DATA] >= HR_RED[LOW] and sensorData[i][VITALS_DATA][HR_DATA] < HR_AMBER[LOW]) or \
                    (sensorData[i][VITALS_DATA][HR_DATA] > HR_AMBER[HIGH] and sensorData[i][VITALS_DATA][HR_DATA] <= HR_RED[HIGH])):
                    self.colorText("red")
                elif(sensorData[i][VITALS_DATA][HR_DATA] <= HR_BLACK[LOW] or sensorData[i][VITALS_DATA][HR_DATA] >= HR_BLACK[HIGH]):
                    self.colorText("blue")

                # Print Data
                print_at(HR_ROW, DATA_START + (DATA_SPACING * i), str(sensorData[i][VITALS_DATA][HR_DATA]) + "   ")

            #< SPO2 >#
            for i in range(WARRIOR_TOTAL):
                # Determine Color Status
                if(sensorData[i][VITALS_DATA][SPO2_DATA] > SPO2_AMBER and sensorData[i][VITALS_DATA][SPO2_DATA] <= SPO2_GREEN):
                    self.colorText("green")
                elif(sensorData[i][VITALS_DATA][SPO2_DATA] > SPO2_RED and sensorData[i][VITALS_DATA][SPO2_DATA] <= SPO2_AMBER):
                    self.colorText("yellow")
                elif(sensorData[i][VITALS_DATA][SPO2_DATA] > SPO2_BLACK and sensorData[i][VITALS_DATA][SPO2_DATA] <= SPO2_RED):
                    self.colorText("red")
                elif(sensorData[i][VITALS_DATA][SPO2_DATA] <= SPO2_BLACK):
                    self.colorText("blue")

                # Print Data
                print_at(SPO2_ROW, DATA_START + (DATA_SPACING * i), str(sensorData[i][VITALS_DATA][SPO2_DATA]) + "   ")

            #< RESP >#
            for i in range(WARRIOR_TOTAL):
                # Determine Color Status
                if(sensorData[i][VITALS_DATA][RESP_DATA] >= RESP_GREEN[LOW] and sensorData[i][VITALS_DATA][RESP_DATA] <= RESP_GREEN[HIGH]):
                    self.colorText("green")
                elif((sensorData[i][VITALS_DATA][RESP_DATA] >= RESP_AMBER[LOW] and sensorData[i][VITALS_DATA][RESP_DATA] < RESP_GREEN[LOW]) or \
                    (sensorData[i][VITALS_DATA][RESP_DATA] > RESP_GREEN[HIGH] and sensorData[i][VITALS_DATA][RESP_DATA] <= RESP_AMBER[HIGH])):
                    self.colorText("yellow")
                elif((sensorData[i][VITALS_DATA][RESP_DATA] >= RESP_RED[LOW] and sensorData[i][VITALS_DATA][RESP_DATA] < RESP_AMBER[LOW]) or \
                    (sensorData[i][VITALS_DATA][RESP_DATA] > RESP_AMBER[HIGH] and sensorData[i][VITALS_DATA][RESP_DATA] <= RESP_RED[HIGH])):
                    self.colorText("red")
                elif(sensorData[i][VITALS_DATA][RESP_DATA] <= RESP_BLACK[LOW] or sensorData[i][VITALS_DATA][RESP_DATA] >= RESP_BLACK[HIGH]):
                    self.colorText("blue")

                # Print Data
                print_at(RESP_ROW, DATA_START + (DATA_SPACING * i), str(sensorData[i][VITALS_DATA][RESP_DATA]) + "   ")

            #< TEMP >#
            for i in range(WARRIOR_TOTAL):
                # Determine Color Status
                if(sensorData[i][VITALS_DATA][TEMP_DATA] >= TEMP_GREEN[LOW] and sensorData[i][VITALS_DATA][TEMP_DATA] <= TEMP_GREEN[HIGH]):
                    self.colorText("green")
                elif((sensorData[i][VITALS_DATA][TEMP_DATA] >= TEMP_AMBER[LOW] and sensorData[i][VITALS_DATA][TEMP_DATA] < TEMP_GREEN[LOW]) or \
                    (sensorData[i][VITALS_DATA][TEMP_DATA] > TEMP_GREEN[HIGH] and sensorData[i][VITALS_DATA][TEMP_DATA] <= TEMP_AMBER[HIGH])):
                    self.colorText("yellow")
                elif((sensorData[i][VITALS_DATA][TEMP_DATA] >= TEMP_RED[LOW] and sensorData[i][VITALS_DATA][TEMP_DATA] < TEMP_AMBER[LOW]) or \
                    (sensorData[i][VITALS_DATA][TEMP_DATA] > TEMP_AMBER[HIGH] and sensorData[i][VITALS_DATA][TEMP_DATA] <= TEMP_RED[HIGH])):
                    self.colorText("red")
                elif(sensorData[i][VITALS_DATA][TEMP_DATA] <= TEMP_BLACK[LOW] or sensorData[i][VITALS_DATA][TEMP_DATA] >= TEMP_BLACK[HIGH]):
                    self.colorText("blue")

                # Print Data
                print_at(TEMP_ROW, DATA_START + (DATA_SPACING * i), str(sensorData[i][VITALS_DATA][TEMP_DATA]) + "   ")

            #< WEAPON >#
            for i in range(WARRIOR_TOTAL):
                if(sensorData[i][WEAPON_DATA] == WEAPON_GREEN):
                    self.colorText("green")
                elif(sensorData[i][WEAPON_DATA] == WEAPON_AMBER):
                    self.colorText("yellow")
                elif(sensorData[i][WEAPON_DATA] > WEAPON_RED):
                    self.colorText("red")
                elif(sensorData[i][WEAPON_DATA] >= WEAPON_BLACK[LOW] or sensorData[i][WEAPON_DATA] <= WEAPON_BLACK[HIGH]):
                    self.colorText("blue")

                # Print Data
                print_at(WEAPON_ROW, DATA_START + (DATA_SPACING * i), str(sensorData[i][WEAPON_DATA]) + "   ")

    def colorText(self, color):
        # lock the colorText if it is in use by another process
        while(self.colorText_Mutex.value == 1):
            pass

        # lock mutex
        self.colorText_Mutex.value = 1
        if(color == "red"):
            print(chr(27) + "[31m" + chr(27) + "[033F")
        elif(color == "green"):
            print(chr(27) + "[32m" + chr(27) + "[033F")
        elif(color == "yellow"):
            print(chr(27) + "[33m" + chr(27) + "[033F")
        elif(color == "blue"):
            print(chr(27) + "[34m" + chr(27) + "[033F")
        elif(color == "magenta"):
            print(chr(27) + "[35m" + chr(27) + "[033F")
        elif(color == "cyan"):
            print(chr(27) + "[36m" + chr(27) + "[033F")
        elif(color == "white"):
            print(chr(27) + "[37m" + chr(27) + "[033F")

        #unlock mutex
        self.colorText_Mutex.value = 0

    def hideCursor(self):
        print("\033[?25h")

class ConnectionStatus(object):

    def __init__(self):
        self.stateIndicatorTable = ["|", "/", "-", "\\"]
        self.errorIndicator = "0"
        self.errorState = 0
        self.stateIndicator = 0

    """ Cyclic Incrementing """
    def incrementStateIndicator(self):
        if(self.stateIndicator < len(self.stateIndicatorTable) - 1):
            self.stateIndicator += 1

        else:
            self.stateIndicator = 0

    def getStateIndicator(self):
        return self.stateIndicatorTable[self.stateIndicator]

    def displayConnectionStatus(self):
        while(1):
            time.sleep(1)
            if(not self.errorState):
                print_at(CONNECTION_STATUS_ROW, CONNECTION_STATUS_START, str(self.getStateIndicator()))
            else:
                print_at(CONNECTION_STATUS_ROW, CONNECTION_STATUS_START, str(self.errorIndicator))

# Depricated, but still might be useful
if(RUNMODE is DEMO):
    # Setup
    connectionStatus = ConnectionStatus()
    messagePrintService = MessagePrintService()

    clearScreen()
    messagePrintService.displayStructure()

    counter = 0

    # Loop
    while(1):
        #messagePrintService.displayData() #< to run this, pass in sample data
        if(not(counter % 200)):
            connectionStatus.incrementStateIndicator()
            counter += 1

        elif(counter is 1000):
            clearScreen()
            messagePrintService.displayStructure()
            counter = 0

        else:
            counter += 1

        connectionStatus.displayConnectionStatus()


# this file is encoded 50 times in base 64

import base64

file = open(r"E:\Downloads\b64.txt", "r")

flag = file.read()

code = ""

for i in range(0,50):
    code = base64.b64decode(flag)
    flag = code

print(flag)


import requests
from bs4 import BeautifulSoup
from selenium.webdriver.common.by import By
from selenium import webdriver
import re
import unicodedata
import time

def get_movie_info(movie_url):
    HEADERS ={"User-Agent":"Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:66.0) Gecko/20100101 Firefox/66.0", "Accept-Encoding":"gzip, deflate", "Accept":"text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8", "DNT":"1","Connection":"close", "Upgrade-Insecure-Requests":"1"}
    response = requests.get(movie_url,headers=HEADERS)
    soup = BeautifulSoup(response.text, 'lxml')
    title = soup.find_all(class_='sc-afe43def-1 fDTGTb')[0].get_text()
    rating = soup.find_all(class_='ipc-inline-list ipc-inline-list--show-dividers sc-afe43def-4 kdXikI baseAlt')[0].find_all('li')[1].get_text()
    length = soup.find_all(class_='ipc-inline-list ipc-inline-list--show-dividers sc-afe43def-4 kdXikI baseAlt')[0].find_all('li')[2].get_text()
    release_date= soup.find_all(class_='ipc-inline-list ipc-inline-list--show-dividers ipc-inline-list--inline ipc-metadata-list-item__list-content base')[3].find_all('li')[0].get_text()
    genres = soup.find_all(class_='ipc-chip-list__scroller')[0].find_all('a')
    list1=[]
    for genre in genres:

        list1.append(genre.get_text())
    list1= ','.join(list1)
    director=[]
    writer=[]
    star=[]
    directors=soup.find_all(class_='ipc-inline-list ipc-inline-list--show-dividers ipc-inline-list--inline ipc-metadata-list-item__list-content baseAlt')[0]
    writers=soup.find_all(class_='ipc-inline-list ipc-inline-list--show-dividers ipc-inline-list--inline ipc-metadata-list-item__list-content baseAlt')[1]
    stars=soup.find_all(class_='ipc-inline-list ipc-inline-list--show-dividers ipc-inline-list--inline ipc-metadata-list-item__list-content baseAlt')[2]
    review_url= 'https://www.imdb.com'+soup.find_all(class_='ipc-link ipc-link--baseAlt ipc-link--touch-target sc-9e83797f-2 jJzPWH isReview')[0]['href']
    budget= soup.find_all(class_='ipc-inline-list ipc-inline-list--show-dividers ipc-inline-list--inline ipc-metadata-list-item__list-content base')[10].find_all('li')[0].get_text()
    #box_office = soup.find_all(class_='ipc-inline-list ipc-inline-list--show-dividers ipc-inline-list--inline ipc-metadata-list-item__list-content base')[13].find_all('li')[0].get_text()
    for item in directors:
        director.append(item.get_text())
    director='/'.join(director)
    for item in writers:
        writer.append(item.get_text())
    writer='/'.join(writer)
    for item in stars:
        star.append(item.get_text())
    star='/'.join(star)
    return (title,rating,list1,length,release_date,director,writer,star,budget,review_url)

def get_movie_id(review_url):
    return review_url.split('/')[4]

def get_review_info(review):
    #driver = webdriver.Chrome()
    try:
        rating = review.find_element(By.CSS_SELECTOR,'span.rating-other-user-rating').text
    except:
        rating = None

    try:
        title = review.find_element(By.CSS_SELECTOR,'a.title').text
    except:
        title = None

    try:
        user_name = review.find_element(By.CSS_SELECTOR,'span.display-name-link').text
    except:
        user_name = None

    try:
        user_url = review.find_element(By.CSS_SELECTOR,'span.display-name-link a').get_attribute('href')
    except:
        user_url = None

    try:
        review_date = review.find_element(By.CSS_SELECTOR,'span.review-date').text
    except:
        review_date = None
    try:
        helpful = review.find_element(By.CSS_SELECTOR,'div.text-muted').text
    except:
        helpful = None

    try:
        expander = review.find_element(By.CSS_SELECTOR,'div.spoiler-warning__control')
        expander.click()
    except:
        expander = None

    try:
        content = review.find_element(By.CSS_SELECTOR,'div.content').text
    except:
        content = None

    return [rating, title, user_name, user_url, review_date, helpful, expander is not None, content]

def load_page(driver):
    i=0
    while i<5:
        try:
            loadMoreButton = driver.find_element(By.ID,"load-more-trigger")
            time.sleep(2)
            loadMoreButton.click()
            time.sleep(2)
        except Exception as e:
            print('e')
            break

# Convert length to min
def parse_time(x):
    if not x:
        return x
    length = 0
    h = re.search('(\d*)h', x)

    if h:
        length += int(h.group(1)) * 60

    m = re.search('(\d*)m', x)

    if m:
        length += int(m.group(1))
    return length

def unicodeToAscii(s):

        return ''.join(
            c for c in unicodedata.normalize('NFD', s)
            if unicodedata.category(c) != 'Mn')


import os
import re
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import glob
from PIL import Image
from pathlib import Path
from torch import distributions as pyd
from torch.distributions.utils import _standard_normal

__REDUCE__ = lambda b: "mean" if b else "none"

def l1(pred, target, reduce=False):
    """Computes the L1-loss between predictions and targets."""
    return F.l1_loss(pred, target, reduction=__REDUCE__(reduce))

def mse(pred, target, reduce=False):
    """Computes the MSE loss between predictions and targets."""
    return F.mse_loss(pred, target, reduction=__REDUCE__(reduce))

def _get_out_shape(in_shape, layers):
    """Utility function. Returns the output shape of a network for a given input shape."""
    x = torch.randn(*in_shape).unsqueeze(0)
    return (
        (nn.Sequential(*layers) if isinstance(layers, list) else layers)(x)
        .squeeze(0)
        .shape
    )

def orthogonal_init(m):
    """Orthogonal layer initialization."""
    if isinstance(m, nn.Linear):
        nn.init.orthogonal_(m.weight.data)
        if m.bias is not None:
            nn.init.zeros_(m.bias)
    elif isinstance(m, nn.Conv2d):
        gain = nn.init.calculate_gain("relu")
        nn.init.orthogonal_(m.weight.data, gain)
        if m.bias is not None:
            nn.init.zeros_(m.bias)

def soft_update_params(m, m_target, tau):
    """Update slow-moving average of online network (target network) at rate tau."""
    with torch.no_grad():
        for p, p_target in zip(m.parameters(), m_target.parameters()):
            p_target.data.lerp_(p.data, tau)

def set_requires_grad(net, value):
    """Enable/disable gradients for a given (sub)network."""
    for param in net.parameters():
        param.requires_grad_(value)

class TruncatedNormal(pyd.Normal):
    """Utility class implementing the truncated normal distribution."""

    def __init__(self, loc, scale, low=-1.0, high=1.0, eps=1e-6):
        super().__init__(loc, scale, validate_args=False)
        self.low = low
        self.high = high
        self.eps = eps

    def _clamp(self, x):
        clamped_x = torch.clamp(x, self.low + self.eps, self.high - self.eps)
        x = x - x.detach() + clamped_x.detach()
        return x

    def sample(self, clip=None, sample_shape=torch.Size()):
        shape = self._extended_shape(sample_shape)
        eps = _standard_normal(shape, dtype=self.loc.dtype, device=self.loc.device)
        eps *= self.scale
        if clip is not None:
            eps = torch.clamp(eps, -clip, clip)
        x = self.loc + eps
        return self._clamp(x)

class NormalizeImg(nn.Module):
    """Module that divides (pixel) observations by 255."""

    def __init__(self):
        super().__init__()

    def forward(self, x):
        return x.div(255.0)

class Flatten(nn.Module):
    """Module that flattens its input to a (batched) vector."""

    def __init__(self):
        super().__init__()

    def forward(self, x):
        return x.view(x.size(0), -1)

def enc(cfg):
    """Returns our MoDem encoder that takes a stack of 224x224 frames as input."""
    C = int(3 * cfg.frame_stack)
    layers = [
        NormalizeImg(),
        nn.Conv2d(C, cfg.num_channels, 7, stride=2),
        nn.ReLU(),
        nn.Conv2d(cfg.num_channels, cfg.num_channels, 5, stride=2),
        nn.ReLU(),
        nn.Conv2d(cfg.num_channels, cfg.num_channels, 3, stride=2),
        nn.ReLU(),
        nn.Conv2d(cfg.num_channels, cfg.num_channels, 3, stride=2),
        nn.ReLU(),
        nn.Conv2d(cfg.num_channels, cfg.num_channels, 3, stride=2),
        nn.ReLU(),
        nn.Conv2d(cfg.num_channels, cfg.num_channels, 3, stride=2),
        nn.ReLU(),
    ]
    out_shape = _get_out_shape((C, cfg.img_size, cfg.img_size), layers)
    layers.extend([Flatten(), nn.Linear(np.prod(out_shape), cfg.latent_dim)])
    return nn.Sequential(*layers)

def state_enc(cfg):
    """Returns a proprioceptive state encoder + modality fuse."""
    return (
        nn.Sequential(
            nn.Linear(cfg.state_dim, cfg.enc_dim),
            nn.ELU(),
            nn.Linear(cfg.enc_dim, cfg.latent_dim),
        ),
        nn.Sequential(
            nn.Linear(cfg.latent_dim, cfg.enc_dim),
            nn.ELU(),
            nn.Linear(cfg.enc_dim, cfg.latent_dim),
        ),
    )

def mlp(in_dim, mlp_dim, out_dim, act_fn=nn.ELU()):
    """Returns an MLP."""
    if isinstance(mlp_dim, int):
        mlp_dim = [mlp_dim, mlp_dim]
    return nn.Sequential(
        nn.Linear(in_dim, mlp_dim[0]),
        act_fn,
        nn.Linear(mlp_dim[0], mlp_dim[1]),
        act_fn,
        nn.Linear(mlp_dim[1], out_dim),
    )

def q(cfg, act_fn=nn.ELU()):
    """Returns a Q-function that uses Layer Normalization."""
    return nn.Sequential(
        nn.Linear(cfg.latent_dim + cfg.action_dim, cfg.mlp_dim),
        nn.LayerNorm(cfg.mlp_dim),
        nn.Tanh(),
        nn.Linear(cfg.mlp_dim, cfg.mlp_dim),
        act_fn,
        nn.Linear(cfg.mlp_dim, 1),
    )

class RandomShiftsAug(nn.Module):
    """
    Random shift image augmentation.
    Adapted from https://github.com/facebookresearch/drqv2
    """

    def __init__(self, cfg):
        super().__init__()
        self.pad = int(
            cfg.img_size / 21
        )  # maintain same padding ratio as in original implementation

    def forward(self, x):
        n, _, h, w = x.size()
        assert h == w
        padding = tuple([self.pad] * 4)
        x = F.pad(x, padding, "replicate")
        eps = 1.0 / (h + 2 * self.pad)
        arange = torch.linspace(
            -1.0 + eps, 1.0 - eps, h + 2 * self.pad, device=x.device, dtype=x.dtype
        )[:h]
        arange = arange.unsqueeze(0).repeat(h, 1).unsqueeze(2)
        base_grid = torch.cat([arange, arange.transpose(1, 0)], dim=2)
        base_grid = base_grid.unsqueeze(0).repeat(n, 1, 1, 1)
        shift = torch.randint(
            0, 2 * self.pad + 1, size=(n, 1, 1, 2), device=x.device, dtype=x.dtype
        )
        shift *= 2.0 / (h + 2 * self.pad)
        grid = base_grid + shift
        return F.grid_sample(x, grid, padding_mode="zeros", align_corners=False)

class Episode(object):
    """Storage object for a single episode."""

    def __init__(self, cfg, init_obs, init_state=None):
        self.cfg = cfg
        self.device = torch.device("cpu")
        self.obs = torch.empty(
            (cfg.episode_length + 1, *init_obs.shape),
            dtype=torch.uint8,
            device=self.device,
        )
        self.obs[0] = torch.tensor(init_obs, dtype=torch.uint8, device=self.device)
        self.state = torch.empty(
            (cfg.episode_length + 1, *init_state.shape),
            dtype=torch.float32,
            device=self.device,
        )
        self.state[0] = torch.tensor(
            init_state, dtype=torch.float32, device=self.device
        )
        self.action = torch.empty(
            (cfg.episode_length, cfg.action_dim),
            dtype=torch.float32,
            device=self.device,
        )
        self.reward = torch.empty(
            (cfg.episode_length,), dtype=torch.float32, device=self.device
        )
        self.cumulative_reward = 0
        self.done = False
        self._idx = 0

    @classmethod
    def from_trajectory(cls, cfg, obs, states, action, reward, done=None):
        """Constructs an episode from a trajectory."""
        episode = cls(cfg, obs[0], states[0])
        episode.obs[1:] = torch.tensor(
            obs[1:], dtype=episode.obs.dtype, device=episode.device
        )
        episode.state[1:] = torch.tensor(
            states[1:], dtype=episode.state.dtype, device=episode.device
        )
        episode.action = torch.tensor(
            action, dtype=episode.action.dtype, device=episode.device
        )
        episode.reward = torch.tensor(
            reward, dtype=episode.reward.dtype, device=episode.device
        )
        episode.cumulative_reward = torch.sum(episode.reward)
        episode.done = True
        episode._idx = cfg.episode_length
        return episode

    def __len__(self):
        return self._idx

    @property
    def first(self):
        return len(self) == 0

    def __add__(self, transition):
        self.add(*transition)
        return self

    def add(self, obs, state, action, reward, done):
        self.obs[self._idx + 1] = torch.tensor(
            obs, dtype=self.obs.dtype, device=self.obs.device
        )
        self.state[self._idx + 1] = torch.tensor(
            state, dtype=self.state.dtype, device=self.state.device
        )
        self.action[self._idx] = action
        self.reward[self._idx] = reward
        self.cumulative_reward += reward
        self.done = done
        self._idx += 1

def get_demos(cfg):
    fps = glob.glob(str(Path(cfg.demo_dir) / "demonstrations" / f"{cfg.task}/*.pt"))
    episodes = []
    for fp in fps:
        data = torch.load(fp)
        frames_dir = Path(os.path.dirname(fp)) / "frames"
        assert frames_dir.exists(), "No frames directory found for {}".format(fp)
        frame_fps = [frames_dir / fn for fn in data["frames"]]
        obs = np.stack([np.array(Image.open(fp)) for fp in frame_fps]).transpose(
            0, 3, 1, 2
        )
        state = torch.tensor(data["states"], dtype=torch.float32)
        if cfg.task.startswith("mw-"):
            state = torch.cat((state[:, :4], state[:, 18 : 18 + 4]), dim=-1)
        elif cfg.task.startswith("adroit-"):
            if cfg.task == "adroit-door":
                state = np.concatenate([state[:, :27], state[:, 29:32]], axis=1)
            elif cfg.task == "adroit-hammer":
                state = state[:, :36]
            elif cfg.task == "adroit-pen":
                state = np.concatenate([state[:, :24], state[:, -9:-6]], axis=1)
            else:
                raise NotImplementedError()
        actions = np.array(data["actions"], dtype=np.float32).clip(-1, 1)
        if cfg.task.startswith("mw-") or cfg.task.startswith("adroit-"):
            rewards = (
                np.array(
                    [
                        _data[
                            "success" if "success" in _data.keys() else "goal_achieved"
                        ]
                        for _data in data["infos"]
                    ],
                    dtype=np.float32,
                )
                - 1.0
            )
        else:  # use dense rewards for DMControl
            rewards = np.array(data["rewards"])
        episode = Episode.from_trajectory(cfg, obs, state, actions, rewards)
        episodes.append(episode)
    return episodes

class ReplayBuffer(object):
    """
    Storage and sampling functionality for training MoDem.
    Uses prioritized experience replay by default.
    """

    def __init__(self, cfg):
        self.cfg = cfg
        self.device = torch.device("cpu")
        self.capacity = 2 * cfg.train_steps + 1
        obs_shape = (3, *cfg.obs_shape[-2:])
        self._state_dim = cfg.state_dim
        self._obs = torch.empty(
            (self.capacity + 1, *obs_shape), dtype=torch.uint8, device=self.device
        )
        self._last_obs = torch.empty(
            (self.capacity // cfg.episode_length, *cfg.obs_shape),
            dtype=torch.uint8,
            device=self.device,
        )
        self._action = torch.empty(
            (self.capacity, cfg.action_dim), dtype=torch.float32, device=self.device
        )
        self._reward = torch.empty(
            (self.capacity,), dtype=torch.float32, device=self.device
        )
        self._state = torch.empty(
            (self.capacity, self._state_dim), dtype=torch.float32, device=self.device
        )
        self._last_state = torch.empty(
            (self.capacity // cfg.episode_length, self._state_dim),
            dtype=torch.float32,
            device=self.device,
        )
        self._priorities = torch.ones(
            (self.capacity,), dtype=torch.float32, device=self.device
        )
        self._eps = 1e-6
        self.idx = 0

    def __len__(self):
        return self.idx

    def __add__(self, episode: Episode):
        self.add(episode)
        return self

    def add(self, episode: Episode):
        obs = episode.obs[:-1, -3:]
        if episode.obs.shape[1] == 3:
            last_obs = episode.obs[-self.cfg.frame_stack :].view(
                self.cfg.frame_stack * 3, *self.cfg.obs_shape[-2:]
            )
        else:
            last_obs = episode.obs[-1]
        self._obs[self.idx : self.idx + self.cfg.episode_length] = obs
        self._last_obs[self.idx // self.cfg.episode_length] = last_obs
        self._action[self.idx : self.idx + self.cfg.episode_length] = episode.action
        self._reward[self.idx : self.idx + self.cfg.episode_length] = episode.reward
        states = torch.tensor(episode.state, dtype=torch.float32)
        self._state[
            self.idx : self.idx + self.cfg.episode_length, : self._state_dim
        ] = states[:-1]
        self._last_state[
            self.idx // self.cfg.episode_length, : self._state_dim
        ] = states[-1]
        max_priority = (
            1.0
            if self.idx == 0
            else self._priorities[: self.idx].max().to(self.device).item()
        )
        mask = (
            torch.arange(self.cfg.episode_length)
            >= self.cfg.episode_length - self.cfg.horizon
        )
        new_priorities = torch.full(
            (self.cfg.episode_length,), max_priority, device=self.device
        )
        new_priorities[mask] = 0
        self._priorities[self.idx : self.idx + self.cfg.episode_length] = new_priorities
        self.idx = (self.idx + self.cfg.episode_length) % self.capacity

    def update_priorities(self, idxs, priorities):
        self._priorities[idxs] = priorities.squeeze(1).to(self.device) + self._eps

    def _get_obs(self, arr, idxs):
        obs = torch.empty(
            (self.cfg.batch_size, self.cfg.frame_stack * 3, *arr.shape[-2:]),
            dtype=arr.dtype,
            device=torch.device("cuda"),
        )
        obs[:, -3:] = arr[idxs].cuda(non_blocking=True)
        _idxs = idxs.clone()
        mask = torch.ones_like(_idxs, dtype=torch.bool)
        for i in range(1, self.cfg.frame_stack):
            mask[_idxs % self.cfg.episode_length == 0] = False
            _idxs[mask] -= 1
            obs[:, -(i + 1) * 3 : -i * 3] = arr[_idxs].cuda(non_blocking=True)
        return obs.float()

    def sample(self):
        probs = self._priorities[: self.idx] ** self.cfg.per_alpha
        probs /= probs.sum()
        total = len(probs)
        idxs = torch.from_numpy(
            np.random.choice(
                total, self.cfg.batch_size, p=probs.cpu().numpy(), replace=True
            )
        ).to(self.device)
        weights = (total * probs[idxs]) ** (-self.cfg.per_beta)
        weights /= weights.max()
        obs = (
            self._get_obs(self._obs, idxs)
            if self.cfg.frame_stack > 1
            else self._obs[idxs].cuda(non_blocking=True)
        )
        next_obs_shape = (3 * self.cfg.frame_stack, *self._last_obs.shape[-2:])
        next_obs = torch.empty(
            (self.cfg.horizon + 1, self.cfg.batch_size, *next_obs_shape),
            dtype=obs.dtype,
            device=obs.device,
        )
        action = torch.empty(
            (self.cfg.horizon + 1, self.cfg.batch_size, *self._action.shape[1:]),
            dtype=torch.float32,
            device=self.device,
        )
        reward = torch.empty(
            (self.cfg.horizon + 1, self.cfg.batch_size),
            dtype=torch.float32,
            device=self.device,
        )
        state = self._state[idxs, : self._state_dim]
        next_state = torch.empty(
            (self.cfg.horizon + 1, self.cfg.batch_size, *state.shape[1:]),
            dtype=state.dtype,
            device=state.device,
        )
        for t in range(self.cfg.horizon + 1):
            _idxs = idxs + t
            next_obs[t] = (
                self._get_obs(self._obs, _idxs + 1)
                if self.cfg.frame_stack > 1
                else self._obs[_idxs + 1].cuda(non_blocking=True)
            )
            action[t] = self._action[_idxs]
            reward[t] = self._reward[_idxs]
            next_state[t] = self._state[_idxs + 1, : self._state_dim]

        mask = (_idxs + 1) % self.cfg.episode_length == 0
        next_obs[-1, mask] = (
            self._last_obs[_idxs[mask] // self.cfg.episode_length]
            .to(next_obs.device, non_blocking=True)
            .float()
        )
        state = state.cuda(non_blocking=True)
        next_state[-1, mask] = (
            self._last_state[_idxs[mask] // self.cfg.episode_length, : self._state_dim]
            .to(next_state.device)
            .float()
        )
        next_state = next_state.cuda(non_blocking=True)
        next_obs = next_obs.cuda(non_blocking=True)
        action = action.cuda(non_blocking=True)
        reward = reward.unsqueeze(2).cuda(non_blocking=True)
        idxs = idxs.cuda(non_blocking=True)
        weights = weights.cuda(non_blocking=True)

        return obs, next_obs, action, reward, state, next_state, idxs, weights

def linear_schedule(schdl, step):
    """Outputs values following a linear decay schedule"""
    try:
        return float(schdl)
    except ValueError:
        match = re.match(r"linear\((.+),(.+),(.+)\)", schdl)
        if match:
            init, final, duration = [float(g) for g in match.groups()]
            mix = np.clip(step / duration, 0.0, 1.0)
            return (1.0 - mix) * init + mix * final
    raise NotImplementedError(schdl)