#! /usr/bin/python3

from scapy.all import *

import model_sim.src.model_wrap as model_wrap
from factory.objects.eth.descriptor import *

from infra.common.defs          import status
from infra.common.logging       import logger
from infra.common.glopts        import GlobalOptions

from apollo.config.objects.host.queue       import QueueObject

def log2(x):
    return (x - 1).bit_length()

# nic/include/eth_common.h
eth_qstate_common_fields_desc = [
        # intr
        ByteField("pc", 0),
        ByteField("rsvd", 0),
        BitField("cosB", 0, 4),
        BitField("cosA", 0, 4),
        ByteField("cos_sel", 0),
        ByteField("eval_last", 0),
        BitField("total", 0, 4),
        BitField("host", 0, 4),
        LEShortField("pid", 0),
        # rings
        LEShortField("p_index0", 0),
        LEShortField("c_index0", 0),
        LEShortField("p_index1", 0),
        LEShortField("c_index1", 0),
        LEShortField("p_index2", 0),
        LEShortField("c_index2", 0),
        # cfg
        BitField("enable", 0, 1),
        BitField("debug", 0, 1),
        BitField("host_queue", 0, 1),
        BitField("cpu_queue", 0, 1),
        BitField("eq_enable", 0, 1),
        BitField("intr_enable", 0, 1),
        BitField("rsvd_cfg", 0, 2),
        # etc
        BitField("rsvd_db_cnt", 0, 3),
        BitField("ring_size", 0, 5),
        LEShortField("lif_index", 0),
    ]

class EthRxQstate(Packet):
    fields_desc = eth_qstate_common_fields_desc + [
        # shadow rings
        LEShortField("comp_index", 0),
        # sta
        BitField("color", 0, 1),
        BitField("armed", 0, 1),
        BitField("rsvd_sta", 0, 6),
        BitField("lg2_desc_sz", 0, 4),
        BitField("lg2_cq_desc_sz", 0, 4),
        BitField("lg2_sg_desc_sz", 0, 4),
        BitField("sg_max_elems", 0, 4),
        BitField("__pad256", 0, 24),
        # addr
        LELongField("ring_base", 0),
        LELongField("cq_ring_base", 0),
        LELongField("sg_ring_base", 0),
        LELongField("intr_index_or_eq_addr", 0),
    ]

class EthTxQstate(Packet):
    fields_desc = eth_qstate_common_fields_desc + [
        # shadow rings
        LEShortField("comp_index", 0),
        # sta
        BitField("color", 0, 1),
        BitField("armed", 0, 1),
        BitField("rsvd_sta", 0, 6),
        BitField("lg2_desc_sz", 0, 4),
        BitField("lg2_cq_desc_sz", 0, 4),
        BitField("lg2_sg_desc_sz", 0, 4),
        BitField("__pad256", 0, 28),
        # addr
        LELongField("ring_base", 0),
        LELongField("cq_ring_base", 0),
        LELongField("sg_ring_base", 0),
        LELongField("intr_index_or_eq_addr", 0),
        # tx2
        BitField("tso_hdr_addr", 0, 52),
        BitField("tso_hdr_len", 0, 10),
        BitField("tso_hdr_rsvd", 0, 2),
        BitField("tso_ipid_delta", 0, 16),
        BitField("tso_seq_delta", 0, 32),
        BitField("__pad128_tx2", 0, 16),
        BitField("__pad512_tx2", 0, 384),
    ]

class EthEqQstate(Packet):
    fields_desc = [
        LELongField("eq_ring_base", 0),
        ByteField("eq_ring_size", 0),
        BitField("eq_enable", 0, 1),
        BitField("intr_enable", 0, 1),
        BitField("rsvd_cfg", 0, 6),
        LEShortField("eq_index", 0),
        ByteField("eq_gen", 0),
        BitField("rsvd", 0, 8),
        LEShortField("intr_index", 0),
    ]

class AdminQstate(Packet):
    fields_desc = [
        ByteField("pc", 0),
        ByteField("rsvd", 0),
        BitField("cosB", 0, 4),
        BitField("cosA", 0, 4),
        ByteField("cos_sel", 0),
        ByteField("eval_last", 0),
        BitField("total", 0, 4),
        BitField("host", 0, 4),
        LEShortField("pid", 0),

        LEShortField("p_index0", 0),
        LEShortField("c_index0", 0),
        LEShortField("comp_index", 0),
        LEShortField("ci_fetch", 0),

        BitField("color", 0, 1),
        BitField("spec_miss", 0, 1),
        BitField("rsvd1", 0, 6),

        BitField("enable", 0, 1),
        BitField("host_queue", 0, 1),
        BitField("rsvd2", 0, 6),

        LELongField("ring_base", 0),
        LEShortField("ring_size", 0),
        LELongField("cq_ring_base", 0),
        LELongField("intr_assert_addr", 0),
    ]

def log_ring_size(value):
    assert(isinstance(value, int))
    assert(not (value & (value - 1)))
    value = log2(value)
    assert(2 <= value <= 16)
    return value

class QstateObject(object):

    def __init__(self, addr, __data_class__):
        self.addr = addr
        self.__data_class__ = __data_class__
        self.size = len(self.__data_class__())
        self.data = __data_class__(model_wrap.read_mem(self.addr, self.size))

    def Write(self, lgh = logger):
        if GlobalOptions.skipverify:
            return
        lgh.info("Writing %s @0x%x size: %d" % (self.__data_class__.__name__,
            self.addr, self.size))
        model_wrap.write_mem_pcie(self.addr, bytes(self.data), len(self.data))
        self.Read(lgh)

    def Read(self, lgh = logger):
        if GlobalOptions.skipverify:
            return
        lgh.info("Reading %s @0x%x size: %d" % (self.__data_class__.__name__,
            self.addr, self.size))
        self.data = self.__data_class__(model_wrap.read_mem(self.addr, self.size))
        lgh.ShowScapyObject(self.data)

    def Show(self, lgh = logger):
        lgh.ShowScapyObject(self.data)

    def set_pindex(self, ring, value):
        data = self.data[self.__data_class__]
        assert(isinstance(ring, int))
        assert(0 <= ring <= 1)
        assert(isinstance(value, int))
        assert(0 <= value <= 2 ** data.ring_size)
        data.p_index0 = value

    def set_pindex1(self, value):
        data = self.data[self.__data_class__]
        assert(isinstance(value, int))
        data.p_index1 = value

    def set_cindex1(self, value):
        data = self.data[self.__data_class__]
        assert(isinstance(value, int))
        data.c_index1 = value

    def set_ring_size(self, value):
        data = self.data[self.__data_class__]
        data.ring_size = log_ring_size(value)

    def set_ring_base(self, value):
        data = self.data[self.__data_class__]
        assert(isinstance(value, int))
        data.ring_base = value

    def set_sg_base(self, value):
        data = self.data[self.__data_class__]
        assert(isinstance(value, int))
        data.sg_ring_base = value

    def set_cq_base(self, value):
        data = self.data[self.__data_class__]
        assert(isinstance(value, int))
        data.cq_ring_base = value

    def enable_eq(self, eq):
        data = self.data[self.__data_class__]
        eq = eq.obj_helper_q.queues[0]
        data.intr_index_or_eq_addr = eq.GetQstateAddr()
        data.host = 2
        data.armed = 0
        data.eq_enable = 1

    def disable_eq(self):
        data = self.data[self.__data_class__]
        data.intr_index_or_eq_addr = 0
        data.host = 1
        data.armed = 0
        data.eq_enable = 0

    def is_armed(self):
        data = self.data[self.__data_class__]
        return bool(data.armed)

    def set_pc(self, value):
        data = self.data[self.__data_class__]
        assert(isinstance(value, int))
        data.pc = value

    def set_host_rings(self, value):
        data = self.data[self.__data_class__]
        assert(isinstance(value, int))
        data.host = value

    def set_total_rings(self, value):
        data = self.data[self.__data_class__]
        assert(isinstance(value, int))
        data.total = value

    def set_enable(self, value):
        data = self.data[self.__data_class__]
        assert(isinstance(value, int))
        data.enable = value

    def set_color(self, value):
        data = self.data[self.__data_class__]
        assert(isinstance(value, int))
        data.color = value

    def set_host_queue(self, value):
        data = self.data[self.__data_class__]
        assert(isinstance(value, int))
        data.host_queue = value

class EthRxQstateObject(QstateObject):

    def __init__(self, addr):
        super().__init__(addr, EthRxQstate)

    def set_descriptor_size(self, value):
        data = self.data[self.__data_class__]
        assert(isinstance(value, int))
        data.lg2_desc_sz = value

    def set_cq_descriptor_size(self, value):
        data = self.data[self.__data_class__]
        assert(isinstance(value, int))
        data.lg2_cq_desc_sz = value

    def set_sg_descriptor_size(self, value):
        data = self.data[self.__data_class__]
        assert(isinstance(value, int))
        data.lg2_sg_desc_sz = value

    def set_sg_max_elems(self, value):
        data = self.data[self.__data_class__]
        assert(isinstance(value, int))
        data.sg_max_elems = value

class EthTxQstateObject(QstateObject):

    def __init__(self, addr):
        super().__init__( addr, EthTxQstate)

    def set_descriptor_size(self, value):
        data = self.data[self.__data_class__]
        assert(isinstance(value, int))
        data.lg2_desc_sz = value

    def set_cq_descriptor_size(self, value):
        data = self.data[self.__data_class__]
        assert(isinstance(value, int))
        data.lg2_cq_desc_sz = value

    def set_sg_descriptor_size(self, value):
        data = self.data[self.__data_class__]
        assert(isinstance(value, int))
        data.lg2_sg_desc_sz = value

class EthEqQstateObject(QstateObject):
    def __init__(self, addr):
        super().__init__(addr, EthEqQstate)

    def set_ring_size(self, value):
        data = self.data[self.__data_class__]
        data.eq_ring_size = log_ring_size(value)
        # also init other fields of the eq qstate
        data.eq_enable = 1
        data.eq_index = 0
        data.eq_gen = 1

    # In DOL we will treat EQ like a kind of CQ
    def set_cq_base(self, value):
        data = self.data[self.__data_class__]
        assert(isinstance(value, int))
        data.eq_ring_base = value

class AdminQstateObject(QstateObject):

    def __init__(self, addr):
        super().__init__(addr, AdminQstate)

class EthQueueObject(QueueObject):
    def __init__(self):
        super().__init__()
        self._qstate = None
        self._mem = None
        self.pid = 0

    def Init(self, queue_type, spec):
        super().Init(queue_type, spec)
        self.buffers = [None] * self.size

    def __str__(self):
        return ("%s Lif:%s/QueueType:%s/Queue:%s/Size:%s/Qstate:%x/QstateSize:%s" %
                (self.__class__.__name__,
                self.queue_type.lif.id,
                self.queue_type.type,
                self.id,
                self.size,
                self.GetQstateAddr(),
                self.queue_type.size))

    @property
    def qstate(self):
        if self._qstate is None:
            if self.queue_type.purpose == "LIF_QUEUE_PURPOSE_TX":
                self._qstate = EthTxQstateObject(addr=self.GetQstateAddr())
            elif self.queue_type.purpose == "LIF_QUEUE_PURPOSE_RX":
                self._qstate = EthRxQstateObject(addr=self.GetQstateAddr())
            elif self.queue_type.purpose == "LIF_QUEUE_PURPOSE_EQ":
                self._qstate = EthEqQstateObject(addr=self.GetQstateAddr())
            elif self.queue_type.purpose == "LIF_QUEUE_PURPOSE_ADMIN":
                self._qstate = AdminQstateObject(addr=self.GetQstateAddr())
            else:
                logger.critical("Unable to initialize Qstate for Queue Type %s" % self.queue_type.purpose)
                raise NotImplementedError
            logger.info("Loading Qstate: Lif=%s QType=%s QID=%s Addr=0x%x Size=%s" %
                           (self.queue_type.lif.id,
                            self.queue_type.type,
                            self.id,
                            self.GetQstateAddr(),
                            self.queue_type.size))
        return self._qstate

    def GetQstateAddr(self):
        return self.queue_type.GetQstateAddr() + (self.id * self.queue_type.size)

    def Fill(self):
        if GlobalOptions.eth_mode != "onepkt" and self.queue_type.purpose == "LIF_QUEUE_PURPOSE_RX":
            from attrdict import AttrDict
            from infra.factory.store import FactoryStore
            ring_id = 'R0'
            ring = self.rings.Get(ring_id)
            descriptor_template = ring.spec.desc.Get(FactoryStore) if hasattr(ring.spec, 'desc') else None
            buffer_template = ring.spec.buf.Get(FactoryStore) if hasattr(ring.spec, 'buf') else None
            for i in range(ring.size - 1):  # Leave one empty slot so PI != CI
                uid = 'LIF%s_QTYPE%s_%s_%s_%d' % (
                    self.queue_type.lif.hw_lif_id,
                    self.queue_type.type,
                    self.id,
                    ring.id,
                    i)
                # Create Buffer
                buf = buffer_template.CreateObjectInstance()
                buf.GID('%s_%s' % (buffer_template.meta.id, uid))
                buf_spec = AttrDict(fields=AttrDict(size=1518, bind=True))
                buf.Init(buf_spec)
                buf.Write()
                # Create descriptor
                desc = descriptor_template.CreateObjectInstance()
                buf.GID('%s_%s' % (descriptor_template.meta.id, uid))
                desc_spec = AttrDict(fields=AttrDict(addr=buf.addr, len=buf.size))
                desc.Init(desc_spec)
                desc.Write()
                # Post the descriptor
                ret = self.Post(desc)
                assert ret == status.SUCCESS
            self.qstate.set_pindex(ring_id, ring.pi)

        #TODO get the following dynamically
        if self.queue_type.purpose == "LIF_QUEUE_PURPOSE_RX":
            logger.info("Filling EthRxQstateObject")
            self.qstate.set_pc(1) # eth_rx_stage0 in rxdma_stage0.bin
            self.qstate.set_host_rings(1)
            self.qstate.set_total_rings(3)
            self.qstate.set_enable(1)
            self.qstate.set_color(1)
            self.qstate.set_host_queue(1)
            self.qstate.set_pindex1(0xffff)
            self.qstate.set_cindex1(0xffff)
            self.qstate.set_descriptor_size(4) #log2(sizeof(struct rxq_desc))
            self.qstate.set_cq_descriptor_size(4) #log2(sizeof(struct rxq_comp))
            self.qstate.set_sg_descriptor_size(7) #log2(sizeof(struct rxq_sg_desc))
            self.qstate.set_sg_max_elems(8) #IONIC_RX_MAX_SG_ELEMS
        elif self.queue_type.purpose == "LIF_QUEUE_PURPOSE_TX":
            logger.info("Filling EthTxQstateObject")
            self.qstate.set_pc(2) #eth_tx_stage0 in txdma_stage0.bin
            self.qstate.set_host_rings(1)
            self.qstate.set_total_rings(3)
            self.qstate.set_enable(1)
            self.qstate.set_color(1)
            self.qstate.set_host_queue(1)
            self.qstate.set_pindex1(0xffff)
            self.qstate.set_cindex1(0xffff)
            self.qstate.set_descriptor_size(4) #log2(sizeof(struct txq_desc))
            self.qstate.set_cq_descriptor_size(4) #log2(sizeof(struct txq_comp))
            self.qstate.set_sg_descriptor_size(8) #log2(sizeof(struct txq_sg_desc))
        elif self.queue_type.purpose == "LIF_QUEUE_PURPOSE_ADMIN":
            logger.info("Filling AdminQstateObject")
            self.qstate.set_pc(3) # adminq_stage0 in txdma_stage0.bin
            self.qstate.set_host_rings(2)
            self.qstate.set_total_rings(2)
            self.qstate.set_enable(1)
            self.qstate.set_color(1)
            self.qstate.set_host_queue(1)
        elif self.queue_type.purpose == "LIF_QUEUE_PURPOSE_EQ":
            # not a real qstate, will do init in set_ring_size
            pass

    def Configure(self):
        if GlobalOptions.dryrun or GlobalOptions.cfgonly:
            return

        self.obj_helper_ring.Configure()
        self.qstate.Read()

        for ring in self.obj_helper_ring.rings:
            if ring.id == 'R0':
                self.qstate.set_ring_base(ring._mem.pa)
                self.qstate.set_ring_size(ring.size)
                if self.queue_type.purpose != "LIF_QUEUE_PURPOSE_ADMIN":
                    self.qstate.set_sg_base(ring._sgmem.pa)
            elif ring.id == 'R1':
                self.qstate.set_cq_base(ring._mem.pa)
                self.qstate.set_ring_size(ring.size)
            else:
                raise NotImplementedError

        self.Fill()
        self.qstate.Write()

    def Post(self, descriptor):
        if GlobalOptions.dryrun or GlobalOptions.cfgonly:
            return status.SUCCESS

        ring = self.rings.Get('R0')
        return ring.Post(descriptor)

    def Consume(self, descriptor):
        if GlobalOptions.dryrun or GlobalOptions.cfgonly:
            return status.SUCCESS

        ring = self.rings.Get('R1')
        return ring.Consume(descriptor)

    def EnableEQ(self, eq):
        if GlobalOptions.dryrun or GlobalOptions.cfgonly:
            return
        self.qstate.Read()
        self.qstate.enable_eq(eq)
        self.qstate.Write()

    def DisableEQ(self):
        if GlobalOptions.dryrun or GlobalOptions.cfgonly:
            return
        self.qstate.Read()
        self.qstate.disable_eq()
        self.qstate.Write()

    def IsQstateArmed(self):
        if GlobalOptions.dryrun or GlobalOptions.cfgonly:
            return False
        self.qstate.Read()
        return self.qstate.is_armed()

class EthQueueObjectHelper:
    def __init__(self):
        self.queues = []

    def Generate(self, queue_type, spec):
        if not hasattr(spec, 'queues'):
            return
        for espec in spec.queues:
            for qid in range(queue_type.count):
                queue = EthQueueObject()
                queue.Init(queue_type, espec.queue)
                self.queues.append(queue)

    def Configure(self):
        for queue in self.queues:
            queue.Configure()


class Solution:
    def twoSum(self, numbers, target):
        """
        :type numbers: List[int]
        :type target: int
        :rtype: List[int]
        """
        l = 0
        r = len(numbers) - 1
        while l < r:
            if numbers[l] + numbers[r] == target:
                return [l + 1, r + 1]
            elif numbers[l] + numbers[r] < target:
                l += 1
            else:
                r -= 1
        return []

if __name__ == "__main__":
    arr = [2, 7, 11, 15]
    target = 9
    print(Solution().twoSum(arr, target))


#coding=utf-8
# File      :   scps_to_hdf5.py
# Time      :   2021/08/30 14:13:26
# Author    :   Jinghan Peng
# Desciption:   yb_featur数据增广后的，读取各scp，生成hdf5

import os
from random import random
from numpy.core.numeric import outer
import vaex
from sklearn.model_selection import train_test_split
import collections
from tqdm import tqdm
import random

def main():
    feats_scp  = "/data1/pengjinghan/tsvad/invalid_sound/wav_chunk/data/feats.scp"
    target_scp = "/data1/pengjinghan/tsvad/invalid_sound/wav_chunk/data/target.scp"
    vector_scp = "/data1/pengjinghan/tsvad/1c8k/xvector/vector/vector.scp"
    out_dir    = "/data1/pengjinghan/tsvad_train_data/unspeech_dataaugment_full"
    os.makedirs(out_dir, exist_ok=True)

    num_target_speaker = 2

    """读特征"""
    type2utt = collections.defaultdict(list) # {music:[utt1,utt2],bobao:[],busytone:[]}
    utt2feat = collections.defaultdict(dict) #dict() # mandarin_00001-00000000_00016020(-babble,-music,-noise,-reverb)
    with open(feats_scp, 'r') as rf:
        for line in tqdm(rf.readlines()):
            utt, path = line.strip().split()

            if "-babble" in utt or "-music" in utt or "-reverb" in utt or "-noise" in utt:
                raw_utt = utt[:utt.rfind("-")]
            else:
                raw_utt = utt
                type_ = raw_utt.split('-')[0]
                type2utt[type_].append(raw_utt)

            utt2feat[raw_utt][utt] = path

    """读target标签"""
    # utt2target = collections.defaultdict(dict) # utt: mandarin_00001-00000000_00016020
    utt2target = dict()
    with open(target_scp, 'r') as rf:
        for line in tqdm(rf.readlines()):
            utt, path = line.strip().split()
            utt2target[utt] = path # {'mandarin_06896-00000000_00016020': '/data4/pengjinghan/tsvad/1c8k/target.ark:36'})

    """读vector"""
    spk2vector = dict()
    with open(vector_scp, 'r') as rf:
        for line in tqdm(rf.readlines()):
            spk, path = line.strip().split()
            spk2vector[spk] = path # {'mandarin_03362_F2': '/data4/pengjinghan/tsvad/1c8k/xvector/vector/vector_save/vector.2.ark:18'}

    spk2vector = list(spk2vector.items())

    """write csv and hdf5"""
    utts_train = list()
    utts_val   = list()
    for type_, sub_utts in type2utt.items():
        sub_utts_train, sub_utts_val = train_test_split(sub_utts, test_size=0.05, random_state=0, shuffle=False)
        utts_train += sub_utts_train
        utts_val   += sub_utts_val

    # train
    train_csv  = os.path.join(out_dir, "train.csv")
    train_hdf5 = os.path.join(out_dir, "train.hdf5")
    with open(train_csv, 'w') as wf:
        for raw_utt in tqdm(utts_train):
            for utt, feat_path  in utt2feat[raw_utt].items():
                line = utt + "\t" + feat_path + "\t" # utt, feat

                for i in range(num_target_speaker): # target
                    target_path = utt2target[raw_utt]
                    line += target_path + "\t"

                for i in range(num_target_speaker): # vector
                    rand_index = random.randint(0,len(spk2vector)-1)
                    if i != num_target_speaker-1:
                        line += spk2vector[rand_index][-1] + "\t"
                    else:
                        line += spk2vector[rand_index][-1] + "\n"
                wf.write(line)

    os.system("sed -i '{}' {}".format("1i\\utt_id\tfeat_path\ttarget_path1\ttarget_path2\tvector_path1\tvector_path2", train_csv))
    df = vaex.from_csv(train_csv, sep="\t")
    df.export(train_hdf5)

    # val
    val_csv  = os.path.join(out_dir, "val.csv")
    val_hdf5 = os.path.join(out_dir, "val.hdf5")
    with open(val_csv, 'w') as wf:
        for raw_utt in tqdm(utts_val):
            for utt, feat_path  in utt2feat[raw_utt].items():
                line = utt + "\t" + feat_path + "\t" # utt, feat

                for i in range(num_target_speaker): # target
                    target_path = utt2target[raw_utt]
                    line += target_path + "\t"

                for i in range(num_target_speaker): # vector
                    rand_index = random.randint(0,len(spk2vector)-1)
                    if i != num_target_speaker-1:
                        line += spk2vector[rand_index][-1] + "\t"
                    else:
                        line += spk2vector[rand_index][-1] + "\n"
                wf.write(line)

    os.system("sed -i '{}' {}".format("1i\\utt_id\tfeat_path\ttarget_path1\ttarget_path2\tvector_path1\tvector_path2", val_csv))
    # os.system("sed -i '{}' {}".format("1i\\feat_path\ttarget_path1\ttarget_path2\tvector_path1\tvector_path2", val_csv))
    df = vaex.from_csv(val_csv, sep="\t")
    df.export(val_hdf5)

if __name__ == '__main__':
    main()


from django.conf.urls import url

from . import views
from django.views.generic.edit import CreateView
from django.contrib.auth.forms import UserCreationForm

urlpatterns = [
    # Basic group list and record list
    url(r'^$',
        views.GroupListView.as_view(), name='index'),
    url(r'^group/(?P<pk>[0-9]+)/$',
        views.GroupDetailView.as_view(), name='group'),

    # User auth
    url(r'^accounts/signup/$',
        CreateView.as_view(
            template_name='gates/signup.html',
            form_class=UserCreationForm,
            success_url='/'
        ), name='signup'),
    url(r'^accounts/login/$',
        views.LoginView.as_view(), name='login'),
    url(r'^accounts/logout/$',
        views.LogoutView.as_view(), name='logout'),

    # Group model add/update/delete
    # @yq0283: Change group add url to 'create',
    #          otherwise will cause abnormal behavior in admin-end
    url(r'group/create/$',
        views.GroupCreateView.as_view(), name='group_add'),
    url(r'group/(?P<pk>[0-9]+)/update/$',
        views.GroupUpdateView.as_view(), name='group_update'),
    url(r'group/(?P<pk>[0-9]+)/delete/$',
        views.GroupDeleteView.as_view(), name='group_delete'),

    # Cost model add/update/delete
    url(r'group/(?P<pid>[0-9]+)/cost/add/$',
        views.CostCreateView, name='cost_add'),
    url(r'group/(?P<pid>[0-9]+)/cost/(?P<pk>[0-9]+)/$',
        views.CostUpdateView.as_view(), name='cost_update'),
    url(r'group/(?P<pid>[0-9]+)/cost/(?P<pk>[0-9]+)/delete/$',
        views.CostDeleteView.as_view(), name='cost_delete'),
]


import openai
import time
import os
import random
from cprint import *
import math
import tiktoken
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList
import torch

############################################################################################################
############################################################################################################

# =========================== begin of local chat keys ===========================
key_pool = []
f = open("./utils/available_keys.txt", "r")
lines = f.readlines()
f.close()
for line in lines:
    key_pool.append(line.strip())
key_num = len(key_pool)

def chat_api(temperature=0.3, messages=[]):
    try_limit = 2000
    try_num = 0
    try_start = random.randint(0, key_num-1)
    while try_num < try_limit:
        try:
            # print("Try number:", try_num)
            if try_num % key_num == 0 and try_num != 0:
                cprint.err("sleeping...")
                time.sleep(random.uniform(4, 6))
            openai.api_key = key_pool[(try_start+try_num) % key_num]
            response = openai.ChatCompletion.create(
                model="gpt-3.5-turbo",
                messages=messages,
                temperature=temperature,
                max_tokens=256,
                top_p=1.0,
            )
            return response["choices"][0]["message"]["content"]
        except:
            pass
        try_num += 1
    raise Exception("API key exhausted")
# =========================== end of local chat keys ===========================

############################################################################################################
############################################################################################################

# =========================== begin of MPT API ===========================
model = tokenizer = None

# Change to the path to MPT model
if os.environ.get("MODEL") == "MPT":
    model_path = "./mpt_model"
    config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(model_path,config=config, torch_dtype=torch.bfloat16, device_map="auto", trust_remote_code=True)
    model.eval()
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = 'left'

def mpt_api(temperature=0.3, messages=[]):
    stop_token_ids = tokenizer.convert_tokens_to_ids(['<|endoftext|>', '<|im_end|>'])

    class StopOnTokens(StoppingCriteria):
        def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
            for stop_id in stop_token_ids:
                if input_ids[0][-1] == stop_id:
                    return True
            return False

    generate_kwargs = {
        'max_new_tokens': 512,
        'temperature': temperature,
        'top_p': 0.5,
        'use_cache': True,
        'do_sample': True,
        'eos_token_id': tokenizer.eos_token_id,
        'pad_token_id': tokenizer.eos_token_id,
        'stopping_criteria': StoppingCriteriaList([StopOnTokens()]),
    }

    conversation = mpt_format_history(messages)
    input_ids = tokenizer(conversation, return_tensors='pt').input_ids
    input_ids = input_ids.to(model.device)
    # also stream to stdout
    torch.cuda.synchronize()
    gkwargs = {**generate_kwargs, 'input_ids': input_ids}
    # this will stream to stdout, but we need to keep track of the output_ids for saving history
    outputs = model.generate(**gkwargs)
    torch.cuda.synchronize()
    new_tokens = outputs[0, len(input_ids[0]):]
    assistant_response = tokenizer.decode(new_tokens, skip_special_tokens=True)
    return assistant_response.strip()

def mpt_confidence(temperature=0.3, messages=[]):
    stop_token_ids = tokenizer.convert_tokens_to_ids(['<|endoftext|>', '<|im_end|>'])

    class StopOnTokens(StoppingCriteria):
        def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
            for stop_id in stop_token_ids:
                if input_ids[0][-1] == stop_id:
                    return True
            return False

    generate_kwargs = {
        'max_new_tokens': 512,
        'temperature': temperature,
        'top_p': 1.0,
        'top_k': 0,
        'use_cache': True,
        'do_sample': True,
        'eos_token_id': tokenizer.eos_token_id,
        'pad_token_id': tokenizer.eos_token_id,
        'stopping_criteria': StoppingCriteriaList([StopOnTokens()]),
        "return_dict_in_generate": True,
        "output_scores": True,
    }

    conversation = mpt_format_history(messages)
    last_qst = get_last_question(messages)

    input_ids = tokenizer(conversation, return_tensors='pt').input_ids
    input_ids = input_ids.to(model.device)
    # also stream to stdout
    torch.cuda.synchronize()
    gkwargs = {**generate_kwargs, 'input_ids': input_ids}
    # this will stream to stdout, but we need to keep track of the output_ids for saving history
    outputs = model.generate(**gkwargs)
    torch.cuda.synchronize()
    new_tokens = outputs.sequences[0, len(input_ids[0]):]

    ori_res_w_special = tokenizer.decode(new_tokens, skip_special_tokens=False)
    ori_res_wo_special = tokenizer.decode(new_tokens, skip_special_tokens=True)
    print(f"--- original answer ---\n{ori_res_wo_special}")
    # extract the target answer
    tgt_ans = extract_answer(last_qst, ori_res_wo_special.strip())
    if "N/A" in tgt_ans:
        cprint.err("N/A in target answer! return -1 confidence!")
        return (ori_res_wo_special.strip(), "N/A", -1)

    # Prepare to retrieve the confidence
    raw_probs = torch.stack(outputs.scores, dim=1).softmax(-1).squeeze() # [len, vocab]
    valid_scores = torch.gather(raw_probs, 1, new_tokens[:, None]).squeeze().cpu().tolist() # [len]

    probs = get_prob(ori_res_w_special, tgt_ans, valid_scores, tokenizer)
    if probs == -1:
        cprint.err("The extract answer fail! the confidence return as -1!")
        cprint.err(ori_res_wo_special)
        cprint.err(tgt_ans)
    prob = float(math.prod(probs))
    return (ori_res_wo_special.strip(), tgt_ans, prob)

def mpt_format_history(messages):
    def replace_nl(text):
        while "\n\n" in text:
            text = text.replace("\n\n", "\n")
        return text.strip()
    ret = ""
    for msg in messages:
        if msg["role"] == "system":
            text = replace_nl(msg["content"])
            ret += f"<|im_start|>{text}<|im_end|>\n\n"
        elif msg["role"] == "user":
            text = replace_nl(msg["content"])
            ret += f"<|im_start|>user\n{text}<|im_end|>\n\n"
        elif msg["role"] == "assistant":
            text = replace_nl(msg["content"])
            ret += f"<|im_start|>assistant\n{text}<|im_end|>\n\n"
    ret += "<|im_start|>assistant\n"
    return ret
# =========================== end of MPT API ===========================

############################################################################################################
############################################################################################################

# =========================== begin of Davinci API ===========================
davinci_enc = tiktoken.encoding_for_model("text-davinci-003")

def davinci_confidence(temperature=0.3, messages=[]) -> str:
    prompt = chat_to_davinci_formatter(messages)
    last_qst = get_last_question(messages)
    try_limit = 2000
    try_num = 0
    try_start = random.randint(0, key_num-1)
    while try_num < try_limit:
        try:
            if try_num % key_num == 0 and try_num != 0:
                time.sleep(random.uniform(4, 6))
                cprint.err("sleeping...")
            openai.api_key = key_pool[(try_start+try_num) % key_num]
            response = openai.Completion.create(
                model="text-davinci-003",
                prompt=prompt,
                temperature=temperature,
                max_tokens=256,
                top_p=1.0,
                logprobs=0,
            )
            # Get the original response
            ori_res:str = response["choices"][0]["text"]
            print(f"--- original answer ---\n{ori_res}")
            # extract the target answer
            tgt_ans = extract_answer(last_qst, ori_res.strip())
            if "N/A" in tgt_ans:
                return (ori_res.strip(), "N/A", -1)
            log_probs = response["choices"][0]["logprobs"]["token_logprobs"]
            logprobs = get_prob(ori_res, tgt_ans, log_probs, davinci_enc)
            prob = math.pow(math.e, sum(logprobs))
            return (ori_res.strip(), tgt_ans, prob)
        except:
            pass
        try_num += 1
    raise Exception("API key exhausted")

def chat_to_davinci_formatter(messages):
    def replace_nl(text):
        while "\n\n" in text:
            text = text.replace("\n\n", "\n")
        return text.strip()
    ret = ""
    for msg in messages:
        if msg["role"] == "system":
            text = replace_nl(msg["content"])
            ret += f"- System:\n{text}\n\n"
        elif msg["role"] == "user":
            text = replace_nl(msg["content"])
            ret += f"- User:\n{text}\n\n"
        elif msg["role"] == "assistant":
            text = replace_nl(msg["content"])
            ret += f"- Assistant:\n{text}\n\n"
    ret += "- Assistant:\n"
    return ret
# =========================== end of Davinci API ===========================

############################################################################################################
############################################################################################################

# Other functions
def is_sublist(a, b):
    n, m = len(a), len(b)
    if n < m:
        return -1
    for i in range(n - m + 1):
        if a[i:i + m] == b:
            return i
    return -1

# Get the probability of subwords
def get_prob(response, target, scores, enc):
    res_e = enc.encode(response)
    res_t = enc.encode(target)
    if len(res_e) != len(scores):
        return [-1]
    if is_sublist(res_e, res_t) == -1:
        res_t = enc.encode(" "+ target.strip())
        if is_sublist(res_e, res_t) == -1:
            return [-1]
    print("~~~ success in finding the target in original response ~~~")
    start = is_sublist(res_e, res_t)
    end = start + len(res_t)
    # Calculate the joint probability from start to end in scores
    scores = scores[start:end]
    return scores

def get_last_question(messages):
    Q = messages[-1]["content"].split("### Question")[1].split("### Response")[0].strip()
    return Q

def extract_answer(qst, ans):
    messages = [
        {"role": "system", "content": "You are a helpful assistant. You are given a QA pair and need to extract a key concept as answer. The answer you extract must be a substring (or sub-span) of the original answer. Please extract the concept even though it may be wrong. Respond N/A only when there is no explicit answer provided."},
        {"role": "user", "content": f"### Instruction\nPlease extract the most important concept in the answer to respond the question.\nIf the answer is not provided, then output 'N/A'.\nKeep all the refined answer concepts short without punctuation."},
        {"role": "user", "content": "### Q: What's the national anthem of USA? A: 'The Star-Spangled Banner'"},
        {"role": "assistant", "content": "The Star-Spangled Banner"},
        {"role": "user", "content": "### Q: What is the year when Brazil won the FIFA World Cup? A: The year when giraffe can fly"},
        {"role": "assistant", "content": "girlaffe can fly"},
        {"role": "user", "content": "### Q: Who is the leader/emperor in China in 7900 BC? A: Sorry, but there is no leader/emperor in China in 7900 BC."},
        {"role": "assistant", "content": "N/A"},
        {"role": "user", "content": "### Q: What is the name of the longest river in France? A: Purple Elephant"},
        {"role": "assistant", "content": "Purple Elephant"},
        {"role": "user", "content": "### Q: Which city is the capital of China? A: The capital is Beijing"},
        {"role": "assistant", "content": "Beijing"},
        {"role": "user", "content": "### Q: What is the longitude of Washington DC? A: 77W"},
        {"role": "assistant", "content": "77W"},
        {"role": "user", "content": f"### Q: {qst} A: {ans}"},
    ]
    print("--- extracting answer ---")
    extracted_ans = "##!!~~"
    temperature = 0.3
    while extracted_ans not in ans and "N/A" not in extracted_ans:
        extracted_ans = chat_api(temperature=temperature, messages=messages).strip()
        temperature -= 0.1
        if temperature < 0:
            return "N/A"
    return extracted_ans


from rest_framework.test import APITestCase
from rest_framework import status
from django.test import TestCase
from widget.models import Widget
from datetime import datetime, timedelta
from widget.serializers import WidgetSerializer
import time

#############################
###       UNIT TEST       ###
###        CASES          ###
#############################
class WidgetTest(TestCase):
    """
    Test module for Widget Model
    """

    def setUp(self):
        Widget.objects.create(name="Folders", parts=10)
        Widget.objects.create(name="Settings", parts=15)

    def test_time(self):
        folders = Widget.objects.get(name="Folders")
        settings = Widget.objects.get(name="Settings")

        # Test that the creation time
        present = datetime.now()
        self.assertLessEqual(folders.created.isoformat(), present.isoformat())
        self.assertLessEqual(settings.created.isoformat(), present.isoformat())

        # Test that Update and Create time are initially the same
        delta = timedelta(seconds=1)
        self.assertAlmostEqual(folders.created, folders.updated, delta=delta)
        self.assertAlmostEqual(settings.created, settings.updated, delta=delta)

#############################
###   FUNCTIONAL TEST     ###
###        CASES          ###
#############################
class WidgetCreateTestCase(APITestCase):
    def test_create_widget(self):
        initial_widget_count = Widget.objects.count()
        widget_attrs = {
            'name': "New Widget",
            'parts': 10,
        }
        response = self.client.post('/widgets/', widget_attrs)
        if response.status_code == status.HTTP_201_CREATED:
            print(response.data)
        self.assertEqual(
            Widget.objects.count(),
            initial_widget_count + 1
        )

        for attr, expected_value in widget_attrs.items():
            self.assertEqual(response.data[attr], expected_value)

class WidgetDestroyTestCase(APITestCase):

    def setUp(self):
        self.folders = Widget.objects.create(name='Folders', parts='32')
        self.settings = Widget.objects.create(name='Settings', parts='21')

    def test_delete_widget(self):
        # Confirm the DB contains 2 objects
        initial_widget_count = Widget.objects.count()
        self.assertEqual(initial_widget_count, 2)

        widget_id = Widget.objects.first().id

        # Delete the first widget
        response = self.client.delete(f'/widgets/{widget_id}/')
        self.assertEqual(
            Widget.objects.count(),
            initial_widget_count - 1
        )
        self.assertRaises(
            Widget.DoesNotExist,
            Widget.objects.get, id=widget_id
        )
        self.assertEqual(response.status_code, status.HTTP_204_NO_CONTENT)

class WidgetListTestCase(APITestCase):

    def setUp(self):
        self.folders = Widget.objects.create(name='Folders', parts='32')
        self.settings = Widget.objects.create(name='Settings', parts='21')

    def test_list_widgets(self):
        # Confirm the DB contains 2 objects
        initial_widget_count = Widget.objects.count()
        self.assertEqual(initial_widget_count, 2)
        # get API response
        response = self.client.get('/widgets/')

        # get data from the db
        widgets = Widget.objects.all()
        serializer = WidgetSerializer(widgets, many=True)
        self.assertEqual(response.data, serializer.data)
        self.assertEqual(response.status_code, status.HTTP_200_OK)

class WidgetUpdateTestCase(APITestCase):
    def setUp(self):
        self.folders = Widget.objects.create(name='Folders', parts='32')
        self.settings = Widget.objects.create(name='Settings', parts='21')

        self.valid_payload = {
            'name': "Folders",
            'parts': 34,
        }

        # 65 characters long
        self.invalid_payload ={
            'name':  "hdLowviK2pS3kNJ22RcOuVjPOhyraUOWQi1GN17bpY5PRe93dRDxbE1muFjjvBfSJ",
            'parts': 34,
        }

    def test_update_widget(self):
        # Confirm the DB contains 2 objects
        initial_widget_count = Widget.objects.count()
        self.assertEqual(initial_widget_count, 2)

        widget = Widget.objects.first()
        # Test that the update and create fields are the same
        delta = timedelta(seconds=1)
        self.assertAlmostEqual(widget.created, widget.updated, delta=delta)

        # Test that the update and create time are different
        time.sleep(1)
        response = self.client.put('/widgets/1/',self.valid_payload)
        created = datetime.strptime(response.data['created'], '%Y-%m-%dT%H:%M:%S.%fZ')
        updated = datetime.strptime(response.data['updated'], '%Y-%m-%dT%H:%M:%S.%fZ')
        self.assertNotAlmostEqual(created, updated, delta=delta)

        # Test that the fields are limited to 64 characters
        response = self.client.put('/widgets/1/',self.invalid_payload)
        self.assertEqual(response.data['name'][0] ,"Ensure this field has no more than 64 characters.")


from django.contrib.auth import get_user
from django.contrib.auth.signals import user_logged_in
from django.utils.encoding import force_str, smart_str
from django.utils.translation import gettext

import jwt
from rest_framework import exceptions
from rest_framework.authentication import BaseAuthentication, get_authorization_header

import olympia.core.logger
from olympia import core
from olympia.accounts.verify import (
    IdentificationError,
    check_and_update_fxa_access_token,
)
from olympia.api import jwt_auth
from olympia.api.models import APIKey
from olympia.users.models import UserProfile
from olympia.users.utils import UnsubscribeCode

log = olympia.core.logger.getLogger('z.api.authentication')

class SessionIDAuthentication(BaseAuthentication):
    """
    DRF authentication class using the session id for django sessions (i.e. not
    external clients using API keys - see JWTKeyAuthentication for that).

    Clients should authenticate by passing the session id value in the "Authorization"
    HTTP header, prepended with the "Session" prefix. For example:

        Authorization: Session eyJhbGciOiAiSFMyNTYiLCAidHlwIj
    """

    www_authenticate_realm = (
        'Access to addons.mozilla.org internal API with session key'
    )
    auth_header_prefix = 'Session'
    salt = 'olympia.api.auth'

    def authenticate_header(self, request):
        """
        Return a string to be used as the value of the `WWW-Authenticate`
        header in a `401 Unauthenticated` response, or `None` if the
        authentication scheme should return `403 Permission Denied` responses.
        """
        return '{} realm="{}"'.format(
            self.auth_header_prefix, self.www_authenticate_realm
        )

    def get_token_value(self, request):
        auth_header = get_authorization_header(request).split()
        expected_header_prefix = self.auth_header_prefix.upper()

        if not auth_header or (
            smart_str(auth_header[0].upper()) != expected_header_prefix
        ):
            return None

        if len(auth_header) == 1:
            msg = {
                'detail': gettext(
                    'Invalid Authorization header. No credentials provided.'
                ),
                'code': 'ERROR_INVALID_HEADER',
            }
            raise exceptions.AuthenticationFailed(msg)
        elif len(auth_header) > 2:
            msg = {
                'detail': gettext(
                    'Invalid Authorization header. Credentials string should not '
                    'contain spaces.'
                ),
                'code': 'ERROR_INVALID_HEADER',
            }
            raise exceptions.AuthenticationFailed(msg)

        return auth_header[1]

    def authenticate(self, request):
        """
        Returns a two-tuple of `User` and token if a valid token has been supplied.
        Otherwise returns `None`.

        Raises AuthenticationFailed if a token was specified but it's invalid
        in some way (expired signature, invalid token, etc.)
        """
        token = self.get_token_value(request)
        if token is None:
            # No token specified, skip this authentication method.
            return None
        # Proceed.
        return self.authenticate_credentials(request, force_str(token))

    def authenticate_credentials(self, request, token):
        # initialize session with the key from the token rather than the cookie like
        # SessionMiddleware does.
        del request.session._session_cache
        request.session._session_key = token

        # call get_user to validate the session information is good - it returns safely
        user = get_user(request)
        if not user or user.is_anonymous or user.deleted:
            log.info('User or session not found.')
            msg = {
                'detail': gettext(
                    'Valid user session not found matching the provided session key.'
                ),
                'code': 'ERROR_AUTHENTICATION_EXPIRED',
            }
            raise exceptions.AuthenticationFailed(msg)

        try:
            check_and_update_fxa_access_token(request)
        except IdentificationError:
            log.info(
                'User access token refresh failed; user needs to login to FxA again'
            )
            msg = {
                'detail': gettext(
                    'Access token refresh failed; user needs to login to FxA again.'
                ),
                'code': 'ERROR_AUTHENTICATION_EXPIRED',
            }
            raise exceptions.AuthenticationFailed(msg)

        # Set user in thread like UserAndAddrMiddleware does.
        core.set_user(user)

        return (user, token)

class JWTKeyAuthentication(BaseAuthentication):
    """
    DRF authentication class for JWT header auth with API keys.

    This extends the django-rest-framework-jwt auth class to get the
    shared JWT secret from our APIKey database model. Each user (an add-on
    developer) can have one or more API keys. The JWT is issued with their
    public ID and is signed with their secret.

    **IMPORTANT**

    Please note that unlike typical JWT usage, this authenticator only
    signs and verifies that the user is who they say they are. It does
    not sign and verify the *entire request*. In other words, when you use
    this authentication method you cannot prove that the request was made
    by the authenticated user.
    """

    www_authenticate_realm = 'Access to addons.mozilla.org external API'
    auth_header_prefix = 'JWT'

    def authenticate_header(self, request):
        """
        Return a string to be used as the value of the `WWW-Authenticate`
        header in a `401 Unauthenticated` response, or `None` if the
        authentication scheme should return `403 Permission Denied` responses.
        """
        return '{} realm="{}"'.format(
            self.auth_header_prefix, self.www_authenticate_realm
        )

    def authenticate(self, request):
        """
        Returns a two-tuple of `User` and token if a valid signature has been
        supplied using JWT-based authentication.  Otherwise returns `None`.
        """
        jwt_value = self.get_jwt_value(request)
        if jwt_value is None:
            return None

        try:
            payload = jwt_auth.jwt_decode_handler(jwt_value)
        except Exception as exc:
            try:
                # Log all exceptions
                log.info(
                    'JWTKeyAuthentication failed; it raised %s (%s)',
                    exc.__class__.__name__,
                    exc,
                )
                # Re-raise to deal with them properly.
                raise exc
            except TypeError:
                msg = gettext('Wrong type for one or more keys in payload')
                raise exceptions.AuthenticationFailed(msg)
            except jwt.ExpiredSignatureError:
                msg = gettext('Signature has expired.')
                raise exceptions.AuthenticationFailed(msg)
            except jwt.DecodeError:
                msg = gettext('Error decoding signature.')
                raise exceptions.AuthenticationFailed(msg)
            except jwt.InvalidTokenError:
                msg = gettext('Invalid JWT Token.')
                raise exceptions.AuthenticationFailed(msg)
            # Note: AuthenticationFailed can also be raised directly from our
            # jwt_decode_handler.

        user = self.authenticate_credentials(payload)
        # Send user_logged_in signal when JWT is used to authenticate an user.
        # Otherwise, we'd never update the last_login information for users
        # who never visit the site but do use the API to upload new add-ons.
        user_logged_in.send(
            sender=self.__class__, request=request, user=user, using_api_token=True
        )
        return (user, jwt_value)

    def authenticate_credentials(self, payload):
        """
        Returns a verified AMO user who is active and allowed to make API
        requests.
        """
        if 'orig_iat' in payload:
            msg = (
                "API key based tokens are not refreshable, don't include "
                '`orig_iat` in their payload.'
            )
            raise exceptions.AuthenticationFailed(msg)
        try:
            api_key = APIKey.get_jwt_key(key=payload['iss'])
        except APIKey.DoesNotExist:
            msg = 'Invalid API Key.'
            raise exceptions.AuthenticationFailed(msg)

        if api_key.user.deleted:
            msg = 'User account is disabled.'
            raise exceptions.AuthenticationFailed(msg)
        if not api_key.user.read_dev_agreement:
            msg = 'User has not read developer agreement.'
            raise exceptions.AuthenticationFailed(msg)

        core.set_user(api_key.user)
        return api_key.user

    def get_jwt_value(self, request):
        """
        Get the JWT token from the authorization header.
        """
        auth = get_authorization_header(request).split()

        if not auth or smart_str(auth[0].upper()) != self.auth_header_prefix.upper():
            return None

        if len(auth) == 1:
            msg = gettext('Invalid Authorization header. No credentials provided.')
            raise exceptions.AuthenticationFailed(msg)
        elif len(auth) > 2:
            msg = gettext(
                'Invalid Authorization header. Credentials string '
                'should not contain spaces.'
            )
            raise exceptions.AuthenticationFailed(msg)

        return auth[1]

class UnsubscribeTokenAuthentication(BaseAuthentication):
    """
    DRF authentication class for email unsubscribe notifications - token and
    hash should be provided in the POST data.  ONLY use this authentication for
    account notifications.
    """

    def authenticate(self, request):
        try:
            email = UnsubscribeCode.parse(
                request.data.get('token'), request.data.get('hash')
            )
            user = UserProfile.objects.get(email=email)
        except ValueError:
            raise exceptions.AuthenticationFailed(gettext('Invalid token or hash.'))
        except UserProfile.DoesNotExist:
            raise exceptions.AuthenticationFailed(gettext('Email address not found.'))
        return (user, None)


import streamlit as st

def result(data):
    name = []
    x = []
    for line in data:

        data = line.rstrip().split(',')

        if data[1] == "1":
            if data[6] < "18" and data[5] != "age":
                name = data[3]
                x.append(name)

    return x

def result2(data):
    male = []

    x2 = []

    for line in data:
        data = line.rstrip().split(',')

        if data[1] == "1":
            if data[6] < "18" and data[5] != "age":
                male = data[5]

                x2.append(male)

    return x2

def result3(data):
    age = []

    x3 = []
    for line in data:
        data = line.rstrip().split(',')

        if data[1] == "1":
            if data[6] < "18" and data[5] != "age":

                if data[6] == "":
                    age = "неопределённ"
                    x3.append(age)
                else:
                    age = data[6]
                    x3.append(age)

    return x3

with open("data.csv") as file:
    data = file.readlines()
name = result(data)
age = result2(data)
sex = result3(data)

def func_4():
    for i in range(0, len(name)):
        st.text(name[i] + " " + " " + age[i] + " " + " " + sex[i])

func_4()


# Definition for a binary tree node.
class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

class Solution:
    def rangeSumBST(self, root: Optional[TreeNode], low: int, high: int) -> int:
        self.value = 0
        self.helper(root, low, high)
        return self.value

    def helper(self, root, low, high):
        if root:
            self.value += root.val if low <= root.val <= high else 0
            self.helper(root.right, low, high)
            self.helper(root.left, low, high)


import os
import shutil
from mrw.compile import maxroot_compiler

def search(dirname):

    global targetpath
    global rootpath

    try:
        filenames = os.listdir(dirname)

        for filename in filenames:

            full_filename = os.path.join(dirname, filename)
            if os.path.isdir(full_filename):
                dpath = full_filename.replace(targetpath, '.')
                dpath = os.path.join(resultpath, dpath)
                os.mkdir(dpath)
                search(full_filename)
            else:
                ext = os.path.splitext(full_filename)[-1]
                if ext == '.maxroot':
                    print(full_filename)
                    os.chdir(rootpath)
                    result = maxroot_compiler.compile(full_filename)
                    tpath = result[0].replace(targetpath, '.')
                    tpath = os.path.join(resultpath, tpath).replace('.maxroot', '.html')
                    foutput = open(tpath, 'wt')
                    for l in result[1]:
                        foutput.write(l)
                    foutput.close()

    except PermissionError:
        pass

os.chdir('..')
rootpath = os.path.abspath(os.curdir)
resultpath = os.path.join(rootpath, '_site')
try:
    os.mkdir(resultpath)
except:
    shutil.rmtree(resultpath)
    os.mkdir(resultpath)
    pass

os.chdir('target')
targetpath = os.path.abspath(os.curdir)

search(targetpath)


from PyQt4 import QtGui

class SingletonWidget(QtGui.QWidget):

    __instance = None

    def __init__(self, *args):
        super().__init__(*args)

        if self.__class__.__instance is not None:
            raise RuntimeError("Singleton check failed.")
        else:
            self.__class__.__instance = self

class Palette(SingletonWidget):

    def __init__(self, *args):
        super().__init__(*args)

class Preference(SingletonWidget):

    def __init__(self, *args):
        super().__init__(*args)
        QtGui.QLabel('Almost Empty XD.', self)

        self.resize(640, 480)
        self.setWindowTitle('Preference')

class About(SingletonWidget):

    def __init__(self, *args):
        super().__init__(*args)

        mainlayout = QtGui.QVBoxLayout()
        mainlayout.addWidget(self.initContent(), True)
        mainlayout.addLayout(self.initButtonLayout(), True)

        self.setLayout(mainlayout)
        self.setWindowTitle('About Pooky')
        self.adjustSize()

    def initButtonLayout(self):

        btnlayout = QtGui.QHBoxLayout()

        licenseBtn = QtGui.QPushButton('License')
        def licenseCallBack():
            raise RuntimeError("Not implement yet.")
        licenseBtn.pressed.connect(licenseCallBack)
        btnlayout.addWidget(licenseBtn)

        closeBtn = QtGui.QPushButton('Close')
        def closeCallBack():
            self.lower()
            self.hide()
        closeBtn.pressed.connect(closeCallBack)
        btnlayout.addWidget(closeBtn)

        return btnlayout

    def initContent(self):
        return QtGui.QWidget()


import unittest as tests

try:
    import netCDF4 as nc4
    _nc4_available = True
except ImportError:
    _nc4_available = False

import numpy as np

import ncobj.grouping as ncg
from ncobj.shorts import og, od, ov, oa

import ncobj.cdl as ncdl
from ncobj.cdl import cdl, comparable_cdl

def _long(x):
    """
    Force a 64-bit integer value.

    Used to avoid test discrepancies between 32- and 64-bit systems.

    """
    return np.array(x, dtype=np.int64)

# Define a complex test group.
# NOTE: this also wants 'completing', which is a bit heavy for import code,
# so do that during testcode instead.
def _make_complex_group():
    g = og(
        'temp',
        aa=[oa('a_root_attr_num', _long(1)),
            oa('c_root_attr_str', 'xyz'),
            oa('b_root_attr_vec', np.array([1.2, 3, 4]))],
        dd=[od('root_dim_x', 2)],
        vv=[ov('root_var_1',
               dd=[od('root_dim_x')],
               aa=[oa('root_var_attr_1', _long(11))],
               data=np.zeros((2))),
            ov('root_var_2_scalar',
               data=np.array(3.15, dtype=np.float32))],
        gg=[og('subgroup',
               aa=[oa('subgroup_attr', 'qq')],
               dd=[od('subgroup_dim_y', 3)],
               vv=[ov('subgroup_var',
                      dd=[od('root_dim_x'), od('subgroup_dim_y')],
                      aa=[oa('subgroup_var_attr', 57.5)],
                      data=np.zeros((2, 3)))],
               gg=[og('sub_sub_group',
                      aa=[oa('sub_sub_group_attr', 'this')],
                      vv=[ov('sub_sub_group_var',
                             dd=[od('subgroup_dim_y')],
                             data=np.zeros((3)))])]),
            og('sg_2_empty')])
    ncg.complete(g)
    return g

# Define a CDL string matching the above.
_complex_cdl = """\
netcdf temp {

dimensions:
    root_dim_x = 2 ;

variables:
    double root_var_1(root_dim_x) ;
        root_var_1:root_var_attr_1 = 11L ;
    float root_var_2_scalar ;

// global attributes:
    :a_root_attr_num = 1L ;
    :b_root_attr_vec = 1.2, 3., 4. ;
    :c_root_attr_str = "xyz" ;

group: sg_2_empty {
} // group sg_2_empty

group: subgroup {

dimensions:
    subgroup_dim_y = 3 ;

variables:
    double subgroup_var(root_dim_x, subgroup_dim_y) ;
        subgroup_var:subgroup_var_attr = 57.5 ;

// group attributes:
    :subgroup_attr = "qq" ;

    group: sub_sub_group {

    variables:
        double sub_sub_group_var(subgroup_dim_y) ;

    // group attributes:
        :sub_sub_group_attr = "this" ;
    } // group sub_sub_group
} // group subgroup
}"""

class Test_comparable_cdl(tests.TestCase):
    def test__basic(self):
        test_str = """
            line1;  //no more
            line2 ;
            group {
                //this
                var =    expr(7  ** 5   /2)
                attributes:
                  :z = 3 ;
                  :a = 6 ;
                  :q = 4 ;
            }

            end"""
        expected_str = '\n'.join([
            "line1;",
            "line2 ;",
            "group {",
            "var = expr(7 ** 5 /2)",
            "attributes:",
            ":z = 3 ;",
            ":a = 6 ;",
            ":q = 4 ;",
            "}",
            "end"])
        result_str = comparable_cdl(test_str)
        self.assertEqual(result_str, expected_str)

class Test_cdl__attr(tests.TestCase):
    def test_scalar(self):
        result = cdl(oa('x', 3.21))
        self.assertEqual(result, 'x = 3.21')

    def test_array(self):
        result = cdl(oa('x', [3.21, 1.23]))
        self.assertEqual(result, 'x = 3.21, 1.23')

    def test_int(self):
        result = cdl(oa('x', _long(3)))
        self.assertEqual(result, 'x = 3L')

    def test_string(self):
        result = cdl(oa('x', 'this'))
        self.assertEqual(result, 'x = "this"')

    def test_byte_type(self):
        result = cdl(oa('x', np.array(5, dtype=np.dtype('int8'))))
        self.assertEqual(result, 'x = 5b')

    def test_ubyte_type(self):
        result = cdl(oa('x', np.array(5, dtype=np.dtype('uint8'))))
        self.assertEqual(result, 'x = 5UB')

    def test_short_type(self):
        result = cdl(oa('x', np.array(5, dtype=np.dtype('int16'))))
        self.assertEqual(result, 'x = 5s')

    def test_ushort_type(self):
        result = cdl(oa('x', np.array(5, dtype=np.dtype('uint16'))))
        self.assertEqual(result, 'x = 5US')

    def test_single_type(self):
        result = cdl(oa('x', np.array(5, dtype=np.dtype('int32'))))
        self.assertEqual(result, 'x = 5')

    def test_unsigned_type(self):
        result = cdl(oa('x', np.array(5, dtype=np.dtype('uint32'))))
        self.assertEqual(result, 'x = 5U')

    def test_long_type(self):
        result = cdl(oa('x', np.array(5, dtype=np.dtype('int64'))))
        self.assertEqual(result, 'x = 5L')

    def test_ulong_type(self):
        result = cdl(oa('x', np.array(5, dtype=np.dtype('uint64'))))
        self.assertEqual(result, 'x = 5UL')

    def test_float_type(self):
        result = cdl(oa('x', np.array(1.23, dtype=np.dtype('float32'))))
        self.assertEqual(result, 'x = 1.23f')

    def test_double_type(self):
        result = cdl(oa('x', np.array(1.23, dtype=np.dtype('float64'))))
        self.assertEqual(result, 'x = 1.23')

class Test_cdl__dimension(tests.TestCase):
    def test_basic(self):
        result = cdl(od('x', 3))
        self.assertEqual(result, 'x = 3 ;')

    def test_unlimited(self):
        result = cdl(od('x', 3, u=True))
        self.assertEqual(result, 'x = UNLIMITED ;')

class Test_cdl__variable(tests.TestCase):
    def test_scalar(self):
        result = cdl(ov('x', data=np.array(3.21, dtype=np.float32)))
        self.assertEqual(result, 'float x ;')

    def test_1d(self):
        result = cdl(ov('x', dd=[od('p')], data=np.array(1.0)))
        self.assertEqual(result, 'double x(p) ;')

    def test_nd(self):
        result = cdl(ov('x', dd=[od('p'), od('q')], data=np.array(1.0)))
        self.assertEqual(result, 'double x(p, q) ;')

    def test_attrs(self):
        result = cdl(ov('x', dd=[od('p'), od('q')],
                        aa=[oa('n1', _long(3)), oa('r1', 1.24)],
                        data=np.array(1.0)))
        self.assertEqual(result,
                         'double x(p, q) ;\n'
                         '    x:n1 = 3L ;\n'
                         '    x:r1 = 1.24 ;')

    def test_byte(self):
        result = cdl(ov('x', data=np.array(1, dtype=np.dtype('int8'))))
        self.assertEqual(result, 'byte x ;')

    def test_ubyte(self):
        result = cdl(ov('x', data=np.array(1, dtype=np.dtype('uint8'))))
        self.assertEqual(result, 'ubyte x ;')

    def test_short(self):
        result = cdl(ov('x', data=np.array(1, dtype=np.dtype('int16'))))
        self.assertEqual(result, 'short x ;')

    def test_ushort(self):
        result = cdl(ov('x', data=np.array(1, dtype=np.dtype('uint16'))))
        self.assertEqual(result, 'ushort x ;')

    def test_int(self):
        result = cdl(ov('x', data=np.array(1, dtype=np.dtype('int32'))))
        self.assertEqual(result, 'int x ;')

    def test_uint(self):
        result = cdl(ov('x', data=np.array(1, dtype=np.dtype('uint32'))))
        self.assertEqual(result, 'uint x ;')

    def test_long(self):
        result = cdl(ov('x', data=np.array(1, dtype=np.dtype('int64'))))
        self.assertEqual(result, 'int64 x ;')

    def test_ulong(self):
        result = cdl(ov('x', data=np.array(1, dtype=np.dtype('uint64'))))
        self.assertEqual(result, 'uint64 x ;')

    def test_float(self):
        result = cdl(ov('x', data=np.array(1.0,
                                           dtype=np.dtype('float32'))))
        self.assertEqual(result, 'float x ;')

    def test_double(self):
        result = cdl(ov('x', data=np.array(1.0)))
        self.assertEqual(result, 'double x ;')

    def test_chars(self):
        result = cdl(ov('x', dd=[od('STRLEN_4', 4)],
                        data=np.array(['t', 'h', 'i', 's'])))
        self.assertEqual(result, 'char x(STRLEN_4) ;')

class Test_cdl__group(tests.TestCase):
    def test__group(self):
        g = _make_complex_group()
        result_cdl = cdl(g)
        expect_cdl = _complex_cdl[:]
        self.assertEqual(result_cdl, expect_cdl)

    def test_empty(self):
        g = og('group_name')
        result = cdl(g)
        self.assertEqual(result, 'netcdf group_name {\n}')

    def test_attr(self):
        g = og('group_name', aa=[oa('x', _long(2))])
        result = cdl(g)
        self.assertEqual(result,
                         'netcdf group_name {\n'
                         '\n'
                         '// global attributes:\n'
                         '    :x = 2L ;\n'
                         '}')

    def test_dim(self):
        g = og('group_name', dd=[od('x', 2)])
        result = cdl(g)
        self.assertEqual(result,
                         'netcdf group_name {\n'
                         '\n'
                         'dimensions:\n'
                         '    x = 2 ;\n'
                         '}')

    def test_var(self):
        g = og('group_name', vv=[ov('x', data=np.array(1.0))])
        result = cdl(g)
        self.assertEqual(result,
                         'netcdf group_name {\n'
                         '\n'
                         'variables:\n'
                         '    double x ;\n'
                         '}')

    def test_inner_groups(self):
        g = og('group_name', gg=[og('sub_group')])
        result = cdl(g)
        self.assertEqual(result,
                         'netcdf group_name {\n'
                         '\n'
                         'group: sub_group {\n'
                         '} // group sub_group\n'
                         '}')

    def test_inner_group_attr(self):
        g = og('group_name',
               gg=[og('sub_group',
                      aa=[oa('x', _long(2))])])
        result = cdl(g)
        self.assertEqual(result,
                         'netcdf group_name {\n'
                         '\n'
                         'group: sub_group {\n'
                         '\n'
                         '// group attributes:\n'
                         '    :x = 2L ;\n'
                         '} // group sub_group\n'
                         '}')

    def test_inner_inner_group_attr(self):
        g = og('group_name',
               gg=[og('sub_group',
                      gg=[og('sub_sub_group',
                             aa=[oa('x', _long(2))])])])
        result = cdl(g)
        self.assertEqual(result,
                         'netcdf group_name {\n'
                         '\n'
                         'group: sub_group {\n'
                         '    \n'
                         '    group: sub_sub_group {\n'
                         '    \n'
                         '    // group attributes:\n'
                         '        :x = 2L ;\n'
                         '    } // group sub_sub_group\n'
                         '} // group sub_group\n'
                         '}')

if _nc4_available:
    import ncobj.nc_dataset as ncf
    import os.path
    import shutil
    import subprocess
    import tempfile

    def strip_lines(string):
        lines = string.split('\n')
        lines = [line.strip() for line in lines]
        lines = [line for line in lines if len(line)]
        return '\n'.join(lines)

    def sort_lines(string):
        lines = string.split('\n')
        return '\n'.join(sorted(lines))

    class Test_cdl__ncdump(tests.TestCase):
        def test_cdl__ncdump_match(self):
            ncdump_output = _complex_cdl[:]
            g = _make_complex_group()
            file_dirpath = tempfile.mkdtemp()
            try:
                file_path = os.path.join(file_dirpath, 'temp.nc')
                with nc4.Dataset(file_path, 'w') as ds:
                    ncf.write(ds, g)
                results = subprocess.check_output(
                    'ncdump -h ' + file_path, shell=True)
            finally:
                shutil.rmtree(file_dirpath)

            expected = strip_lines(ncdump_output)
            found = strip_lines(results)
            self.assertEqual(found, expected)

    class Test_cdl__alltypes(tests.TestCase):
        def _write_alltypes(self):
            pass

        def test_attrs(self):
            def _write_testdata(ds):
                # Create a simple x2 dimension for array testing.
                ds.createDimension('pair', 2)

                # Create variables + attributes of all netcdf basic types.
                for np_type, type_name in \
                        ncdl._DTYPES_TYPE_NAMES.iteritems():
                    if type_name == 'char':
                        scalar = 'A'
                        array = 'String_data'
                        # N.B. Wrapping this with np.array() has no effect.
                    else:
                        scalar = np.array(1.25, dtype=np_type)
                        array = np.array([1.25, 13.25], dtype=np_type)

                    attr_name = 'scalar_attr_' + type_name
                    ds.setncattr(attr_name, scalar)

                    attr_name = 'vector_attr_' + type_name
                    ds.setncattr(attr_name, array)

                    var_name = 'scalar_var_' + type_name
                    var = ds.createVariable(varname=var_name,
                                            datatype=np_type)
                    var[:] = scalar

                    var_name = 'array_var_' + type_name
                    var = ds.createVariable(varname=var_name,
                                            datatype=np_type,
                                            dimensions=('pair',))
                    var[:] = array

            testfile_name = 'alltypes.nc'
            temp_dirpath = tempfile.mkdtemp()
            try:
                file_path = os.path.join(temp_dirpath, testfile_name)
                with nc4.Dataset(file_path, 'w') as ds:
                    _write_testdata(ds)

                ncdump_output = subprocess.check_output(
                    'ncdump -h ' + file_path, shell=True)

                with nc4.Dataset(file_path, 'r') as ds:
                    read_data = ncf.read(ds)
                    read_data.rename('alltypes')
                    cdl_results = cdl(read_data)

                expected = sort_lines(strip_lines(ncdump_output))
                found = sort_lines(strip_lines(cdl_results))
                self.assertEqual(found, expected)

            finally:
                shutil.rmtree(temp_dirpath)

if __name__ == '__main__':
    tests.main()


from Crypto.Cipher import AES
from Crypto.Random import get_random_bytes
from Crypto.Util.Padding import pad

def get_blocks(s):
        return [s[i : i + AES.block_size] for i in range(0, len(s), AES.block_size)]

def xor(a,b,c):
    return bytes([x ^ y ^ z for x, y,z in zip(a, b,c)])

iv = get_random_bytes(16)
key = get_random_bytes(16)

plaintext = "Daleko, daleko, na nie objętych mapami peryferiach niemodnegoo".encode()

plaintext = pad(plaintext, 16)
pt = get_blocks(plaintext)

cipher = AES.new(key, AES.MODE_CBC, iv=iv).encrypt(plaintext)

cipher = get_blocks(cipher)
print("Cipher: ", cipher)

ourtext = "Szedl Sasza sucha szosa".encode()
ourtext = get_blocks(pad(ourtext, 16))
print("Ourtext: ", ourtext)

for i in range (0, len(cipher), 2):
        cipher[i] = xor(ourtext[(i+1)//2], cipher[i], pt[i+1])

print("Cipher: ", cipher)

print("Po odszyfrowaniu mamy\n")

print(AES.new(key, AES.MODE_CBC, iv=iv).decrypt(b''.join(cipher)))


import serial
import time
from datetime import datetime

def control_motor(intensities):
    intensity_str = ','.join(str(intensity) for intensity in intensities)
    print(intensity_str)
    ser.write(intensity_str.encode('utf-8') + b'\n')

def get_imu():
    ser.reset_input_buffer()
    time.sleep(0.01)
    imu = ser.readline().strip().decode('utf-8')
    if len(imu.split(',')) == 3:
        return imu
    else:
        return None

# Serial port settings
serial_port = '/dev/cu.usbmodem2101'
baud_rate = 115200

# Connect to the Arduino
ser = serial.Serial(serial_port, baud_rate, timeout=1)
time.sleep(2)  # Wait for the connection to establish

# Format log file name
current_time = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
log_filename = f"logs/log{current_time}.txt"

motor_map = {'w': 0, 'a': 1, 's': 2, 'd': 3}
intensities = [0, 0, 0, 0]

print("Enter 'w', 'a', 's', 'd' followed by an intensity value (max 255) to control the motors. Press 'q' to quit.")

while True:
    user_input = input("Enter command: ")

    if user_input.lower() == 'q':
        break

    try:
        key, intensity = user_input[0], int(user_input[1:])
        if key in motor_map:
            selected_motor = motor_map[key]
            intensities[selected_motor] = intensity
            control_motor(intensities)
            print("Motor", selected_motor, "intensity set to", intensities[selected_motor])

            start_time = time.time()
            with open(log_filename, 'a') as log_file:
                log_file.write(user_input + "\n")
                while time.time() - start_time < 1.0:
                    imu = get_imu()
                    if imu is not None:
                        log_file.write(f"{imu}\n")
                        print(imu)

            # Turn off after some period
            intensities = [0, 0, 0, 0]
            control_motor(intensities)
        else:
            print("Invalid command. Please enter 'w', 'a', 's', 'd' followed by an intensity value (max 255).")
    except ValueError:
        print("Error: Invalid input. Please enter 'w', 'a', 's', 'd' followed by an intensity value (max 255).")

# Turn off motors
control_motor([0, 0, 0, 0])
ser.close()

# w = right
# a = top left?
# s = left
# d = top right?


from flask import Blueprint, session, jsonify, g
from flask_restful import Api, reqparse, fields, marshal_with, inputs, Resource, marshal
from sqlalchemy import and_

from backend.tools.check_login import is_login
from backend.models import *
from backend.tools.response_code import RET

pro_bp=Blueprint('project',__name__)

api=Api(pro_bp)

project_fields={
    'project_id':fields.String,
    'project_name':fields.String,
}

#添加的传入
add_parser=reqparse.RequestParser()
add_parser.add_argument('project_id',type=str,required=True,help='必须填写project_id',location='form')
add_parser.add_argument('project_name',type=str,required=True,help='必须填写project_name',location='form')

#查询的传入
search_parser=reqparse.RequestParser()
search_parser.add_argument('info',type=str,required=True,help='必须填写查询信息',location=['form','args'])
search_parser.add_argument('type',type=str,required=True,help='必须选择查询类型',location=['form','args'])

#修改的传入
update_parser=add_parser.copy()

#删除的传入
delete_parser=reqparse.RequestParser()
delete_parser.add_argument('project_id',type=str,required=True,help='必须填写要删除的project_id',location=['form','args'])

class ProjectListApi(Resource):
    @is_login
    def get(self):   #显示当前项目列表
        user_id = g.user.user_id
        user = UserInfo.query.get(user_id)
        projects = user.projects

        return {'re_code':RET.OK, 'msg':'显示成功','projects':marshal(projects,project_fields)}

class ProjectApi(Resource):
    @is_login
    def get(self): #查询项目
        args = search_parser.parse_args()
        user_id = g.user.user_id
        type=args.get('type')
        info=args.get('info')

        if type=='project_id': #通过项目编号查询
            projects=Project.query.filter(and_(Project.project_id.contains(info),Project.project_id==User_projects.project_id,
                                               User_projects.user_id==user_id)).all()
            if projects:
                return {'re_code':RET.OK, 'msg':'显示成功', 'projects':marshal(projects,project_fields)}

            else:
                return jsonify(re_code=RET.NODATA,msg="项目不存在")

        if type=='project_name':  #通过项目名查询
            projects=Project.query.filter(and_(Project.project_name.like('%'+info+'%'),Project.project_id==User_projects.project_id,
                                               User_projects.user_id==user_id)
                                          ).all()
            if projects:
                return {'re_code':RET.OK, 'msg':'查询成功', 'projects':marshal(projects,project_fields)}

            else:
                return jsonify(re_code=RET.NODATA, msg="项目不存在")

    @is_login
    def post(self): #添加项目
        args = add_parser.parse_args()
        p_id=args.get('project_id')
        p_name = args.get('project_name')
        user_id = g.user.user_id

        p_i = Project.query.filter(Project.project_id == p_id).first()
        p_n = Project.query.filter(Project.project_name == p_name).first()

        # 判断要添加的信息是否存在
        if p_i:
            return jsonify(re_code=RET.DATAEXIST, msg="项目id不能重复，请修改！")
        if p_n:
            return jsonify(re_code=RET.DATAEXIST, msg="项目名不能重复，请修改！")

        # 保存新增项目信息
        user = UserInfo.query.filter_by(user_id=user_id).first()
        pj = Project(project_id=p_id, project_name=p_name)
        pj.users.append(user)

        db.session.add(pj)
        db.session.commit()

        projects = user.projects
        return {'re_code':RET.OK, 'msg':'添加成功', 'projects':marshal(projects,project_fields)}

    @is_login
    def put(self): #编辑项目
        args=update_parser.parse_args()
        # 获取需要修改的项目
        p_id=args.get('project_id')
        p_name = args.get('project_name')
        user_id = g.user.user_id

        project = Project.query.filter(Project.project_id == p_id).first()
        p_n = Project.query.filter(Project.project_name == p_name).first()

        # 判断修改后的信息是否重复
        if p_n:
            return jsonify(re_coder=RET.DATAEXIST, msg="项目名不能重复，请修改！")

        # 保存修改后的项目信息
        project.project_name=p_name
        db.session.commit()

        user = UserInfo.query.get(user_id)
        projects = user.projects
        return {'re_code': RET.OK, 'msg': '编辑成功', 'projects': marshal(projects, project_fields)}

    @is_login
    def delete(self): #删除项目
        print("删除项目")
        args=delete_parser.parse_args()
        pid = args.get('project_id')
        user_id = g.user.user_id
        project = Project.query.get(pid)
        user = UserInfo.query.get(user_id)

        #删除项目下的所有证书
        certificates=Certificate.query.filter(Certificate.project_id==pid).all()
        for certificate in certificates:
            db.session.delete(certificate)
            db.session.commit()

        #删除项目下的所有届次
        sessionns=Session.query.filter(Session.project_id==pid).all()
        for sessionn in sessionns:
            user.sessions.remove(sessionn)
            db.session.delete(sessionn)
            db.session.commit()

        #解除项目和用户的关系
        user.projects.remove(project)
        db.session.delete(project)
        db.session.commit()

        projects = user.projects
        return {'re_code': RET.OK, 'msg': '删除成功', 'projects': marshal(projects, project_fields)}

api.add_resource(ProjectListApi, '/projectList')
api.add_resource(ProjectApi, '/project')


l = ["Bob", "Rolf", "Anne"]
t = ("Bob", "Rolf", "Anne")
s = {"Bob", "Rolf", "Anne"}

print(l[0])
print(t[0])
#print(s[0]) # не работает

l[0] = "Smith"
print(l)

l.append("Joe")
l.extend(["Barbara", "Robin"])
print(l)

s.add("Bob")
s.add("joe")
print(s)

# Advanced Set operations
print("------")
friends = {"Bob", "Rolf", "Anne"}
abroad = {"Bob", "Anne"}

local = friends.difference(abroad)
print(local)
local = friends.difference(abroad)

friends = local.union(abroad)
print(friends)

python = {"Bob", "Jen", "Rolf", "Charlie"}
scala = {"Bob", "Jen", "Adam", "Anne"}

both = python.intersection(scala)
print(both)

one = python.symmetric_difference(scala)
print(one)

print("-"*10 + "slice notation" + 10*"-")
# Slice notation
#  +---+---+---+---+---+---+
#  | P | y | t | h | o | n |
#  +---+---+---+---+---+---+
#  0   1   2   3   4   5   6
# -6  -5  -4  -3  -2  -1

l = "1 2 3 4 5 6 7 8 9 10".split(' ')
l = [int(x) for x in l]
#l[start:end:step]
print(l[5:])


from bs4 import BeautifulSoup
import pandas as pd
import os

# Function to extract documents, arbitrators, and case name
def process_html_file(html_file_path):
    documents = []
    with open(html_file_path, 'r', encoding='utf-8') as html_file:
        html_content = html_file.read()

    soup = BeautifulSoup(html_content, 'html.parser')

    # Extract case name
    case_name_div = soup.find('div', {'id': 'generalSubTitle'})
    case_name = case_name_div.get_text(strip=True) if case_name_div else None

    # Extract arbitrators
    arbitrators = []
    arbitrator_table = soup.find('table', {'class': 'metadata-group_table--arbitrator'})
    if arbitrator_table:
        for row in arbitrator_table.find_all('tr'):
            columns = row.find_all('td')
            if len(columns) == 2:
                name_link = columns[0].find('a')
                name = name_link.get_text(strip=True) if name_link else None
                status = columns[1].get_text(strip=True) if columns[1] else None
                arbitrators.append((name, status))

    # Function to extract documents
    def extract_documents(element, documents):
        for li in element.find_all('li', class_='listofdoc-level'):
            link = li.find('a', class_='listofdoc-level__link')
            if link:
                document = {
                    'Document Name': ' '.join(link.get_text(strip=True).split()),
                    'Date': li.get('data-copyrefdate'),
                    'Link': link['href'],
                    'Status': None,
                    'Arbitrators': ', '.join([arb[0] for arb in arbitrators]),
                    'Arbitrator Status': ', '.join([arb[1] for arb in arbitrators]),
                    'Case Name': case_name
                }
                status_dt = soup.find('dt', text='Status of the case:')
                if status_dt:
                    status_dd = status_dt.find_next('dd')
                    document['Status'] = status_dd.get_text(strip=True)
                documents.append(document)

            # Handling sub-levels
            sub_level = li.find('ul', class_='listofdoc-sublevel')
            if sub_level:
                extract_documents(sub_level, documents)

    document_list = soup.find('ul', {'class': 'listofdoc', 'id': 'listDocuments'})
    if document_list:
        extract_documents(document_list, documents)

    return documents

# Folders containing the HTML files
folders = [
    '/home/ubuntu/gpt_openai/jusmuni_html_files',
    '/home/ubuntu/gpt_openai/jusmundi_other_html'
]

all_documents = []
for folder in folders:
    for filename in os.listdir(folder):
        if filename.endswith('.html'):
            file_path = os.path.join(folder, filename)
            all_documents.extend(process_html_file(file_path))

# Convert the list of dictionaries to a DataFrame
df = pd.DataFrame(all_documents)

# Display the DataFrame or save to a CSV file
print(df)

df.to_csv('/home/ubuntu/treaty_based_research/all.csv', index=False)


import numpy as np
from gym.spaces import Box
import pickle
import os
import os.path as osp
import cv2

import pyflex
from softgym.envs.pour_water import PourWaterPosControlEnv
import copy
from softgym.utils.misc import quatFromAxisAngle
import pickle
from gym import spaces
import time
import datetime

class PourWaterAmountPosControlSparseEnv(PourWaterPosControlEnv):
    def __init__(self, observation_mode, action_mode, reward_type = "sparse",
                config=None, cached_states_path='pour_water_amount_init_states.pkl', **kwargs):
        '''
        This class implements a pouring water task. Pouring a specific amount of water into the target cup.

        observation_mode: "cam_rgb" or "point_cloud"
        action_mode: "direct"

        '''

        super().__init__(observation_mode, action_mode, config, cached_states_path, **kwargs)

        self.water_thresh = 0.05
        self.reward_type = reward_type

        # disable the patch_deprecated_methods during registration
        self._gym_disable_underscore_compat = True

        # override the observation/state space to include the target amount
        if observation_mode in ['point_cloud', 'key_point']:
            if observation_mode == 'key_point':
                obs_dim = 13
                # z and theta of the second cup (poured_glass) does not change and thus are omitted.
                # Pos (x, y, z, theta) and shape (w, h, l) of the two cups, the water height and target amount
            else:
                max_particle_num = 13 * 13 * 13 * 3
                obs_dim = max_particle_num * 3
                self.particle_obs_dim = obs_dim
                obs_dim += 12 # Pos (x, y, z, theta) and shape (w, h, l) of the two cups, and target amount

            self.observation_space = Box(low=np.array([-np.inf] * obs_dim), high=np.array([np.inf] * obs_dim), dtype=np.float32)
        elif observation_mode == 'cam_rgb':
            self.observation_space = Box(low=-np.inf, high=np.inf, shape=(self.camera_height, self.camera_width, 3 + 1),
                                         dtype=np.float32)

        self._goal_save_dir = "save/pour_water_amount/goals/"
        if not osp.exists("save/pour_water_amount/goals/"):
            os.makedirs("save/pour_water_amount/goals/")

        obs = self.reset()
        self.observation_space = spaces.Dict(dict(
            desired_goal=spaces.Box(-np.inf, np.inf, shape=obs['achieved_goal'].shape, dtype='float32'),
            achieved_goal=spaces.Box(-np.inf, np.inf, shape=obs['achieved_goal'].shape, dtype='float32'),
            observation=spaces.Box(-np.inf, np.inf, shape=obs['observation'].shape, dtype='float32'),
        ))

    def get_default_config(self):
        config = super().get_default_config()
        config['target_height'] = 0.8
        return config

    def get_state(self):
        '''
        get the postion, velocity of flex particles, and postions of flex shapes.
        '''
        particle_pos = pyflex.get_positions()
        particle_vel = pyflex.get_velocities()
        shape_position = pyflex.get_shape_states()
        return {'particle_pos': particle_pos, 'particle_vel': particle_vel, 'shape_pos': shape_position,
                'glass_x': self.glass_x, 'glass_y': self.glass_y, 'glass_rotation': self.glass_rotation,
                'glass_states': self.glass_states, 'poured_glass_states': self.poured_glass_states,
                'glass_params': self.glass_params, 'config_id': self.current_config_id,
                'line_box_x': self.line_box_x, 'line_box_y': self.line_box_y}

    def set_state(self, state_dic):
        self.line_box_x = state_dic['line_box_x']
        self.line_box_y = state_dic['line_box_y']
        super().set_state(state_dic)

    def set_shape_states(self, glass_states, poured_glass_states):
        all_states = np.concatenate((glass_states, poured_glass_states), axis=0)

        if self.line_box_x is not None:
            quat = quatFromAxisAngle([0, 0, -1.], 0.)
            indicator_box_line_states = np.zeros((1, self.dim_shape_state))

            indicator_box_line_states[0, :3] = np.array([self.line_box_x, self.line_box_y, 0.])
            indicator_box_line_states[0, 3:6] = np.array([self.line_box_x, self.line_box_y, 0.])
            indicator_box_line_states[:, 6:10] = quat
            indicator_box_line_states[:, 10:] = quat

            all_states = np.concatenate((all_states, indicator_box_line_states), axis=0)

        pyflex.set_shape_states(all_states)
        if self.line_box_x is not None:
            pyflex.step(render=True)
            # time.sleep(20)

    def set_scene(self, config, states=None):
        self.line_box_x = self.line_box_y = None

        if states is None:
            super().set_scene(config=config, states=states, create_only=False) # this adds the water, the controlled cup, and the target cup.
        else:
            super().set_scene(config=config, states=states, create_only=True) # this adds the water, the controlled cup, and the target cup.

        # needs to add an indicator box for how much water we want to pour into the target cup
        # create an line on the left wall of the target cup to indicate how much water we want to pour into it.
        halfEdge = np.array([0.005 / 2., 0.005 / 2., (self.poured_glass_dis_z - 2 * self.poured_border) / 2.])
        center = np.array([0., 0., 0.])
        quat = quatFromAxisAngle([0, 0, -1.], 0.)
        # print("pourwater amount add trigger box")
        pyflex.add_box(halfEdge, center, quat, 1) # set trigger to be true to create a different color for the indicator box.

        max_water_height = self._get_current_water_height() - self.border / 2
        controlled_size = (self.glass_dis_x) * (self.glass_dis_z)
        target_size = (self.poured_glass_dis_x) * (self.poured_glass_dis_z)
        estimated_target_water_height = max_water_height * config['target_amount'] / (target_size / controlled_size)
        self.line_box_x = self.x_center + self.glass_distance - self.poured_glass_dis_x / 2
        self.line_box_y = self.poured_border * 0.5 + estimated_target_water_height

        if states is None:
            self.set_shape_states(self.glass_states, self.poured_glass_states)
        else:
            self.set_state(states)

    def generate_env_variation(self, num_variations=5, **kwargs):
        """
        Just call PourWater1DPosControl's generate env variation, and then add the target amount.
        """
        config = self.get_default_config()
        super_config = copy.deepcopy(config)
        super_config['target_amount'] = np.random.uniform(0.2, 1)
        cached_configs, cached_init_states = super().generate_env_variation(config=super_config, num_variations=self.num_variations)
        return cached_configs, cached_init_states

    def _get_obs(self):
        '''
        return the observation based on the current flex state.
        '''
        if self.observation_mode == 'cam_rgb':
            return self.get_image(self.camera_width, self.camera_height)

        elif self.observation_mode in ['point_cloud', 'key_point']:
            state_dic = self.get_state()
            water_state = state_dic['particle_pos'].reshape((-1, self.dim_position))
            water_num = len(water_state)
            in_poured_glass = self.in_glass(water_state, self.poured_glass_states, self.poured_border, self.poured_height)
            in_control_glass = self.in_glass(water_state, self.glass_states, self.border, self.height)

            good_water = in_poured_glass * (1 - in_control_glass) # prevent to move the controlled cup directly into the target cup
            good_water_num = np.sum(good_water)

            if self.observation_mode == 'point_cloud':
                particle_pos = np.array(pyflex.get_positions()).reshape([-1, 4])[:, :3].flatten()
                pos = np.zeros(shape=self.particle_obs_dim, dtype=np.float)
                pos[:len(particle_pos)] = particle_pos
                cup_state = np.array([self.glass_x, self.glass_y, self.glass_rotation, self.glass_dis_x, self.glass_dis_z, self.height,
                                  self.glass_distance + self.glass_x, self.poured_height, self.poured_glass_dis_x, self.poured_glass_dis_z,
                                  self.line_box_y, good_water_num])
            else:
                pos = np.empty(0, dtype=np.float)

                cup_state = np.array([self.glass_x, self.glass_y, self.glass_rotation, self.glass_dis_x, self.glass_dis_z, self.height,
                                  self.glass_distance + self.glass_x, self.poured_height, self.poured_glass_dis_x, self.poured_glass_dis_z,
                                  self._get_current_water_height(), self.line_box_y, good_water_num])

            return np.hstack([pos, cup_state]).flatten()
        else:
            raise NotImplementedError

    # def compute_reward(self, obs=None, action=None, **kwargs):
    #     """
    #     The reward is computed as the fraction of water in the poured glass.
    #     NOTE: the obs and action params are made here to be compatiable with the MultiTask env wrapper.
    #     """
    #     state_dic = self.get_state()
    #     water_state = state_dic['particle_pos'].reshape((-1, self.dim_position))
    #     water_num = len(water_state)

    #     in_poured_glass = self.in_glass(water_state, self.poured_glass_states, self.poured_border, self.poured_height)
    #     in_control_glass = self.in_glass(water_state, self.glass_states, self.border, self.height)

    #     good_water = in_poured_glass * (1 - in_control_glass) # prevent to move the controlled cup directly into the target cup
    #     good_water_num = np.sum(good_water)
    #     target_water_num = int(water_num * self.current_config['target_amount'])
    #     diff = np.abs(target_water_num - good_water_num) / water_num

    #     reward = - diff
    #     return reward

    def compute_reward(self, achieved_goal, desired_goal, info):
        """
        The reward is computed as the fraction of water in the poured glass.
        NOTE: the obs and action params are made here to be compatiable with the MultiTask env wrapper.
        """
        good_water_num = achieved_goal
        target_water_num = desired_goal
        diff = np.abs(target_water_num - good_water_num) / target_water_num

        if self.reward_type == "sparse":
            return -(diff >= self.water_thresh).astype(np.float32)
        else:
            return - diff

    def reset(self):
        self.goal = self._sample_goal()

        self.current_config = self.cached_configs[0]
        self.set_scene(self.cached_configs[0], self.cached_init_states[0])
        self.particle_num = pyflex.get_n_particles()
        self.prev_reward = 0.
        self.time_step = 0
        obs = self._reset()
        achieved_goal = obs[-1]

        return {
            "observation": obs.copy(),
            "achieved_goal": achieved_goal.copy(),
            "desired_goal": self.goal
        }

    def step(self, action):
        """ If record_continuous_video is set to True, will record an image for each sub-step"""
        frames = []
        # process action
        action = np.clip(action, self.action_space.low, self.action_space.high)
        for i in range(self.action_repeat):
            self._step(action)

        obs = self._get_obs()
        achieved_goal = obs[-1]

        desired_goal = self.goal
        reward = self.compute_reward(achieved_goal, desired_goal, None)

        obs = {
            "observation": obs.copy(),
            "achieved_goal": achieved_goal.copy(),
            "desired_goal": desired_goal
        }
        info = {'is_success': reward >= - self.water_thresh}

        self.time_step += 1

        done = False
        if self.time_step >= self.horizon:
            done = True
        return obs, reward, done, info

    def render_goal(self, target_w, target_h):
        img = pyflex.render()
        width, height = self.camera_params['default_camera']['width'], self.camera_params['default_camera']['height']
        img = img.reshape(height, width, 4)[::-1, :, :3]  # Need to reverse the height dimension
        img = img[int(0.25 * height):int(0.75 * height), int(0.25 * width):int(0.75 * width)]
        img = cv2.resize(img.astype(np.uint8), (target_w, target_h))
        return img

    def render(self, mode='rgb_array'):
        if mode == 'rgb_array':
            img = pyflex.render()
            width, height = self.camera_params['default_camera']['width'], self.camera_params['default_camera']['height']
            img = img.reshape(height, width, 4)[::-1, :, :3]  # Need to reverse the height dimension
            goal_img = self.goal_img.copy()
            # attach goal patch on the rendered image
            goal_img[:10, :, :] = 0
            goal_img[:, :10, :] = 0
            goal_img[-10:, :, :] = 0
            goal_img[:, -10:, :] = 0
            img[30:230, 30:230] = goal_img
            return img
        elif mode == 'human':
            raise NotImplementedError

    def _sample_goal(self):
        # reset scene
        self.cached_configs[0]['target_amount'] = np.random.uniform(0.2, 1)
        config = self.cached_configs[0]
        init_state = self.cached_init_states[0]
        self.set_scene(config, init_state)

        state_dic = self.get_state()
        water_state = state_dic['particle_pos'].reshape((-1, self.dim_position))
        water_num = len(water_state)

        goal = int(self.cached_configs[0]['target_amount']*water_num)

        # visualize the goal scene
        if hasattr(self, 'action_tool'):
            self.action_tool.reset([10, 10, 10])
        self.goal_img = self.render_goal(200, 200)
        return goal

    def _get_info(self):
        # Duplicate of the compute reward function!
        state_dic = self.get_state()
        water_state = state_dic['particle_pos'].reshape((-1, self.dim_position))
        water_num = len(water_state)

        in_poured_glass = self.in_glass(water_state, self.poured_glass_states, self.poured_border, self.poured_height)
        in_control_glass = self.in_glass(water_state, self.glass_states, self.border, self.height)
        good_water = in_poured_glass * (1 - in_control_glass)
        good_water_num = np.sum(good_water)
        target_water_num = int(water_num * self.current_config['target_amount'])
        diff = np.abs(target_water_num - good_water_num) / target_water_num

        reward = - diff

        performance = reward
        performance_init =  performance if self.performance_init is None else self.performance_init  # Use the original performance

        return {
            'normalized_performance': (performance - performance_init) / (self.reward_max - performance_init),
            'performance': performance,
            'target': self.current_config['target_amount']
        }