from fileinput import filename
import socket
from tqdm import tqdm
import os

def main():
    s = socket.socket()
    s.bind((socket.gethostname(), 3339))
    s.listen(10)

    conn, address = s.accept()
    data = conn.recv(1024).decode("utf-8")
    print(data)

    item = data.split("_")
    filename = item[0]
    filesize = int(item[1])

    print("filename and filesize recieved")
    conn.send("filename and filesize recieved".encode("utf-8"))

    bar = tqdm(range(filesize), f"receiving {filename}", unit="B", unit_scale=True, unit_divisor=1024)

    with open(f"recv_{filename}", 'w') as f:
        while True:
            data = conn.recv(1024).decode("utf-8")
            if not data:
                break
            f.write(data)
            bar.update(len(data))
    conn.close()
    s.close()

if __name__ == "__main__":
    main()


# -*- coding: UTF-8 -*-
'''=================================================
@Author ：zhenyu.yang
@Date   ：2020/11/5 11:37 AM
=================================================='''

import sys
sys.path.append('./')
sys.path.insert(0,'/data/zhenyu.yang/modules')

import cv2
import json
import numpy as np
import random
import copy
from multiprocessing import Process
import os

try:
    from .tools import crop_img,eye_crop,get_box
except:
    from tools import crop_img,eye_crop,get_box

def cal_Flow(img1,img2):
    scale = 0.5
    levels = 1
    winsize = 5
    iterations = 10
    poly_n = 5
    poly_sigma = 1.1
    flow = cv2.calcOpticalFlowFarneback(img1, img2,
                                        flow=None, pyr_scale=scale,
                                        levels=levels, iterations=iterations,
                                        winsize=winsize, poly_n=poly_n,
                                        poly_sigma=poly_sigma, flags=0)
    return flow

def scale_height(bbox,scale=1.5):
    w,h = bbox[2]-bbox[0],bbox[3]-bbox[1]
    x,y = (bbox[2]+bbox[0])/2,(bbox[3]+bbox[1])/2
    h = h*scale
    rect = [x-w/2,y-h/2,x+w/2,y+h/2]
    rect = list(map(int,rect))
    return rect

def cellMotion(flow_y):
    h,w = flow_y.shape[:2]
    sep_h = h/3
    sep_w = w/3

    cell_motion = np.zeros((3,3))

    for i in range(3):
        for j in range(3):
            yl = int(i*sep_h)
            yr = int(i * sep_h+sep_h)

            xl = int(j*sep_w)
            xr = int(j * sep_w+sep_w)

            cell_motion[i,j] = np.mean(flow_y[yl:yr,xl:xr])

    return cell_motion

def cal_distance(rect_1,rect_2):
    x_1,y_1 = (rect_1[0]+rect_1[2])/2,(rect_1[1]+rect_1[3])/2
    x_2, y_2 = (rect_2[0] + rect_2[2]) / 2, (rect_2[1] + rect_2[3]) / 2
    return ((x_1-x_2)**2+(y_1-y_2)**2)**0.5

def get_treshold(distance):
    return 0.02*distance/25

def get_single_status(flow,rect,thres):
    eye_rect = scale_height(rect,2)
    eye_flow = eye_crop(flow,eye_rect,1)

    cell_motion = cellMotion(eye_flow)
    cell_motion = cell_motion[:2].reshape(-1)
    average = np.mean(cell_motion)
    var = np.var(cell_motion)

    if var < thres:
        return 0

    if cell_motion[4] > 0 and cell_motion[4] > cell_motion[1]:
        return -1 # close

    if cell_motion[4] < 0 and cell_motion[4] < cell_motion[1]:
        return 1 # open

    return 0

def get_doube_eye_status(flow,info):
    left_eye = get_single_status(flow,info[0],info[-1])
    right_eye = get_single_status(flow,info[1],info[-1])
    return left_eye,right_eye

def getFiles(path, suffix,prefix):
    return [os.path.join(root, file) for root, dirs, files in os.walk(path)
            for file in files if file.endswith(suffix) and file.startswith(prefix)]

def blinklist2num(blink_list):
    num_list = []
    for blink in blink_list:
        num_list.extend(list(range(blink[0], blink[1] + 1)))
    return num_list

def is_good_jsons(img_info):
    ldmks = []
    if 'left_eye_ldmk' in img_info and img_info['left_eye_ldmk'] is not None and  len(img_info['left_eye_ldmk']) > 4:
        ldmks.append([img_info['left_eye_ldmk'],img_info['left_status']])
    if 'right_eye_ldmk' in img_info and img_info['right_eye_ldmk'] is not None and len(img_info['right_eye_ldmk']) > 4:
        ldmks.append([img_info['right_eye_ldmk'],img_info['right_status']])
    if len(ldmks) != 2:
        return False
    return True

def get_batch_data(video_list,suffix,batch_size,start_id,dst_dir,random_ratio = 0.2):
    random.shuffle(video_list)
    while True:
        video_path = video_list.pop()

        json_path = video_path.replace('.mp4', suffix)
        with open(json_path, 'r') as f:
            big_json = f.readlines()

        jsons_dict = {}
        for json_info in big_json:
            try:
                json_info = json.loads(json_info.strip())
            except:
                break
            jsons_dict[int(json_info['frame'])] = json_info

        blink_path = video_path.replace(os.path.basename(video_path), suffix)

        with open(blink_path, 'r') as f:
            jsons = json.load(f)

        close_list = blinklist2num(jsons['blink_close'])
        open_list = blinklist2num(jsons['blink_open'])

        cap = cv2.VideoCapture(video_path)
        frame_id = -1
        data_id = -1

        frames = []
        jsons = []

        while True:
            frame_id += 1
            ret, frame = cap.read()

            if not ret:
                break

            frames.append(frame)
            jsons.append(jsons_dict[frame_id])
            if len(frames) > 2:
                _ = frames.pop(0)
                _ = jsons.pop(0)

            if random.random() > random_ratio:
                continue

            if frame_id not in close_list and frame_id not in open_list:
                if random.random() < 0.6:
                    continue

            if not is_good_jsons(jsons[0]):
                continue

            left_eye = get_box(jsons[0]['left_eye_ldmk'])
            right_eye = get_box(jsons[0]['right_eye_ldmk'])

            thres = get_treshold(cal_distance(left_eye, right_eye))

            flow = cal_Flow(frames[0], frames[1])
            flow_y = flow[:, :, 0]

            left_eye_label, right_eye_label = get_doube_eye_status(flow_y, [left_eye,right_eye,thres])

            label = 0
            if frame_id in close_list:
                label = -1
            if frame_id in open_list:
                label = 1

            temp_info = {'thres': thres,
                         'left_eye': left_eye,
                         'right_eye': right_eye,
                         'left_eye_label':left_eye_label,
                         'right_eye_label':right_eye_label,
                         'label':label}

            data_id += 1

            temp_data_name = '{}.json'.format(start_id +data_id )
            with open(os.path.join(dst_dir,temp_data_name),'w') as f:
                json.dump(temp_info,f)

            flow_name = '{}.npy'.format(start_id +data_id )
            np.save(os.path.join(dst_dir,flow_name),flow_y)

            if data_id >= batch_size:
                return True

        return False

def split(input,num=60):
    random.shuffle(input)

    ans = []
    sep = len(input) //num
    for i in range(num-1):
        ans.append(input[i*sep:(i+1)*sep])

    ans.append(input[(num-1)*sep:])

    return ans

if __name__ == '__main__':

    version = 'v0.1'
    suffix = '_{}.json'.format(version)

    src_dir_dict = {'train':'/data/weiyu.li/DMSData/FatigueView/raw_video',
                    'test':'/data/weiyu.li/DMSData/FatigueView/test_video'

    }
    camera_list = ['ir_down','ir_front','ir_left','ir_left_up','ir_up','rgb_down','rgb_front','rgb_left','rgb_left_up','rgb_up']

    data_type = 'test'
    for camera_id in range(len(camera_list)):

        src_dir = src_dir_dict[data_type]

        camera_type = camera_list[camera_id]

        dst_dir = './data/{}/{}'.format(data_type,camera_type)
        if not os.path.exists(dst_dir):
            os.makedirs(dst_dir)

        video_list = getFiles(src_dir, '.mp4', camera_type)

        all_num = 10000
        running_num = 32
        batch_size = all_num//running_num
        split_videos = split(video_list, running_num)

        process_list = []
        for i in range(running_num):
            temp_p = Process(target=get_batch_data,args=(split_videos[i],suffix,batch_size,batch_size*i,dst_dir,))
            process_list.append(temp_p)

        for temp_p in process_list:
            temp_p.start()

        for temp_p in process_list:
            temp_p.join()

        print('END')


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

from django.db import models, migrations

class Migration(migrations.Migration):

    dependencies = [
        ('pages', '0001_initial'),
    ]

    operations = [
        migrations.AddField(
            model_name='page',
            name='lock_state',
            field=models.CharField(choices=[('UU', 'Unlocked'), ('LL', 'Locked')], max_length=2, default='UU'),
            preserve_default=True,
        ),
    ]


import socket
import time
import base_pb2
import setting
import main_ui
from chessman import *
from PyQt5.Qt import QProgressDialog

class PlayGame(QWidget, main_ui.Ui_Form):
    backSignal = pyqtSignal()  # 返回信号

    def __init__(self, sock, username, parent=None):
        super().__init__(parent=parent)
        self.setupUi(self)
        self.sock = sock
        self.isBeginGame = False
        self.isCanXia = False
        self.label.setCursor(QCursor(Qt.ForbiddenCursor))
        self.username = username
        self.setFixedSize(self.width(), self.height())
        self.sock.readyRead.connect(self.readData)
        self.switch = {
            7: self.cmd7,
            9: self.cmd9,
            10: self.cmd10,
            12: self.cmd12,
            13: self.cmd13,
            14: self.cmd14,
            15: self.cmd15,
            16: self.cmd16,
            18: self.cmd18,
        }
        self.sound_piece = QSound("source/luozisheng.wav")
        self.sound_win = QSound("source/win.wav")
        self.sound_defeated = QSound("source/defeated.wav")
        self.sound_piece_t = self.sound_piece
        self.sound_win_t = self.sound_win
        self.sound_defeated_t = self.sound_defeated
        self.comboBox.currentIndexChanged.connect(self.messageSend)
        self.listView.setEditTriggers(QAbstractItemView.NoEditTriggers)
        self.chessboard = [[None for i in range(19)] for i in range(19)]
        self.history = []
        self.history2 = []
        self.is_over = False
        self.ai = False
        self.lbl = None
        self.isMatch = False
        self.init()
        self.initBeginGame()
        self.setting = setting.Setting(self)
        self.isCloseChat = False
        self.lcdNumber.setMode(QLCDNumber.Dec)
        self.lcdNumber_3.setMode(QLCDNumber.Dec)
        self.lcdNumber_3.setSegmentStyle(QLCDNumber.Flat)
        # self.lcdNumber.setStyleSheet("border: 2px solid black; color: red; ")
        self.time = QTimer(self)
        self.time.setInterval(1000)
        self.time.timeout.connect(self.refresh)

    def init(self):
        self.setWindowTitle("网络五子棋 v1.1")
        # 设置窗口图标
        self.setWindowIcon(QIcon("source/icon.ico"))
        # 设置窗口大小
        # self.setFixedSize(QImage("source/游戏界面.png").size())
        self.listView.clicked.connect(self.listViewClicked)
        self.pushButton.clicked.connect(self.messageSend)
        self.pushButton_2.clicked.connect(self.withdraw)
        self.pushButton_4.clicked.connect(self.beginGame)
        self.pushButton_5.clicked.connect(self.wantToLose)
        self.pushButton_3.clicked.connect(self.draw)
        self.pushButton_6.clicked.connect(self.setUp)
        self.fontText = QFont("Microsoft YaHei", 10, 75)
        self.label_3.setFont(self.fontText)
        self.label_4.setFont(self.fontText)
        self.pushButton_7.clicked.connect(self.aiBegin)
        self.focusPoint = QLabel(self)
        self.focusPoint.setPixmap(QPixmap("source/标识.png"))
        self.pushButton.setShortcut(QKeySequence(Qt.Key_Enter))
        self.pushButton.setShortcut(QKeySequence(Qt.Key_Return))
        self.label_9.setText('在线列表（点击发起对战）：')
        palette = QPalette()
        pix = QPixmap('source/background.png')
        pix = pix.scaled(self.width(), self.height())
        palette.setBrush(QtGui.QPalette.Background, QtGui.QBrush(pix))
        self.setPalette(palette)

    def initBeginGame(self):
        if self.isBeginGame:
            self.pushButton_5.setEnabled(True)
            self.pushButton_4.setEnabled(False)
            self.pushButton_7.setEnabled(False)
            if self.ai == True:
                self.pushButton_3.setEnabled(False)
                self.pushButton_3.setEnabled(False)
            else:
                self.pushButton_3.setEnabled(True)
                self.pushButton_2.setEnabled(True)

            self.time.start()
            self.reStart()
            if self.myColor == 'black':
                self.label_2.setPixmap(QPixmap(r'source\black.png'))
                self.label_5.setPixmap(QPixmap(r'source\white.png'))
                if self.isCanXia:
                    self.turnChessColor = 'black'
                else:
                    self.turnChessColor = 'white'
            elif self.myColor == 'white':
                self.label_2.setPixmap(QPixmap(r'source\white.png'))
                self.label_5.setPixmap(QPixmap(r'source\black.png'))
                if self.isCanXia:
                    self.turnChessColor = 'white'
                else:
                    self.turnChessColor = 'black'
        else:
            self.pushButton_2.setEnabled(False)
            self.pushButton_3.setEnabled(False)
            self.pushButton_4.setEnabled(True)
            self.pushButton_5.setEnabled(False)
            self.pushButton_7.setEnabled(True)
        if self.isCanXia:
            self.label.setCursor(QCursor(Qt.PointingHandCursor))
        else:
            self.label.setCursor(QCursor(Qt.ForbiddenCursor))

        self.endDate = QDateTime.currentMSecsSinceEpoch() + 1000 * 60 * 10
        self.endDateOther = self.endDate

    def refresh(self):
        nowData = QDateTime.currentMSecsSinceEpoch()
        interval = self.endDate - nowData
        interval2 = self.endDateOther - nowData
        if self.isBeginGame == False:
            return
        if interval > 0:
            if self.isCanXia == False:
                self.endDate = self.endDate + 1000
            else:
                self.endDateOther = self.endDateOther + 1000
            self.lcdNumber.display(self.conversionTime(interval))
            self.lcdNumber_2.display(self.conversionTime(interval2))
        else:
            self.lose()

    def conversionTime(self, interval):
        days = interval // (24 * 60 * 60 * 1000)
        hour = (interval - days * 24 * 60 * 60 * 1000) // (60 * 60 * 1000)
        min = (interval - days * 24 * 60 * 60 * 1000 - hour * 60 * 60 * 1000) // (60 * 1000)
        sec = (interval - days * 24 * 60 * 60 * 1000 - hour * 60 * 60 * 1000 - min * 60 * 1000) // 1000
        s = str(min) + ':' + str(sec)
        return s

    def closeEvent(self, a0: QtGui.QCloseEvent):
        self.backSignal.emit()

    def mouseReleaseEvent(self, a0: QtGui.QMouseEvent):
        if self.isCanXia == False:
            self.label.setCursor(QCursor(Qt.ForbiddenCursor))
            return
        else:
            self.label.setCursor(QCursor(Qt.PointingHandCursor))

        if self.isBeginGame == False:
            return None

        pos, chess_index = self.reversePos(a0)
        if pos is None:
            return

        if self.chessboard[chess_index[1]][chess_index[0]] != None:
            return

        self.chess = Chessman(color=self.turnChessColor, parent=self)
        self.chess.setIndex(chess_index[0], chess_index[1])
        self.chess.move()

        s_g_p = base_pb2.server_gobang_position()
        s_g_p.cmd = 6
        s_g_p.x = chess_index[0]
        s_g_p.y = chess_index[1]

        self.writeData(s_g_p.SerializeToString())
        self.isCanXia = False
        self.label.setCursor(QCursor(Qt.ForbiddenCursor))
        if self.ai == True:
            self.autoDown()

    # 坐标转换
    def reversePos(self, a0: QtCore.QPoint):
        if a0.x() <= self.label.x() + 4 or a0.x() >= self.label.x() + 517 + 18 or a0.y() <= self.label.y() + 4 or a0.y() >= self.label.y() + 517 + 18:
            return None, None

        self.x = (a0.x() - self.label.x() - 4) // 36
        self.y = (a0.y() - self.label.y() - 4) // 36

        return QPoint(22 + self.x * 36, 22 + self.y * 36), (self.x, self.y)

    def isWin(self, chessman):
        # print("in iswin,lastChessman:", chessman.color, chessman.x, chessman.y)
        # 水平方向y相同，chessboard[chessman.y][i]
        count = 1
        # 左边
        i = chessman.x - 1
        while i >= 0:
            if self.chessboard[chessman.y][i] == None or self.chessboard[chessman.y][i].color != chessman.color:
                break
            count += 1
            i -= 1
        # 右边
        i = chessman.x + 1
        while i <= 18:
            if self.chessboard[chessman.y][i] == None or self.chessboard[chessman.y][i].color != chessman.color:
                break
            count += 1
            i += 1

        if count >= 5:
            return chessman.color

        count = 1
        j = chessman.y - 1
        while j >= 0:
            if self.chessboard[j][chessman.x] == None or self.chessboard[j][chessman.x].color != chessman.color:
                break
            count += 1
            j -= 1

        j = chessman.y + 1
        while j <= 18:
            if self.chessboard[j][chessman.x] == None or self.chessboard[j][chessman.x].color != chessman.color:
                break
            count += 1
            j += 1

        if count >= 5:
            return chessman.color

        count = 1
        j, i = chessman.y - 1, chessman.x + 1
        while j >= 0 and i <= 18:
            if self.chessboard[j][i] == None or self.chessboard[j][i].color != chessman.color:
                break
            count += 1
            j -= 1
            i += 1

        j, i = chessman.y + 1, chessman.x - 1
        while i >= 0 and j <= 18:
            if self.chessboard[j][i] == None or self.chessboard[j][i].color != chessman.color:
                break
            count += 1
            i -= 1
            j += 1
        if count >= 5:
            return chessman.color

        count = 1
        j, i = chessman.y - 1, chessman.x - 1
        while j >= 0 and i >= 0:
            if self.chessboard[j][i] == None or self.chessboard[j][i].color != chessman.color:
                break
            count += 1
            j -= 1
            i -= 1

        j, i = chessman.y + 1, chessman.x + 1
        while j <= 18 and i <= 18:
            if self.chessboard[j][i] == None or self.chessboard[j][i].color != chessman.color:
                break
            count += 1
            j += 1
            i += 1

        if count >= 5:
            return chessman.color

        return None

    def showResult(self, isWin=None):
        self.lbl = QLabel(self)
        if isWin == 'white':
            self.lbl.setPixmap(QPixmap("source/白棋胜利.png"))
            self.lbl.move(300, 100)
            self.lbl.show()
        elif isWin == 'black':
            self.lbl.setPixmap(QPixmap("source/黑棋胜利.png"))
            self.lbl.move(300, 100)
            self.lbl.show()
        if isWin == self.myColor:
            self.sound_win.play()
        else:
            self.sound_defeated.play()
        self.isBeginGame = False
        self.isCanXia = False
        self.isMatch = False
        self.ai = False
        self.initBeginGame()
        self.time.stop()

    def reStart(self):
        for i in range(19):
            for j in range(19):
                if self.chessboard[i][j] != None:
                    self.chessboard[i][j].close()
                    self.chessboard[i][j] = None
                    self.focusPoint.close()
                else:
                    pass
        if self.lbl != None:
            self.lbl.close()

    def beginGame(self):
        self.reStart()
        if self.isBeginGame == False and self.isMatch == False:
            requestRes = base_pb2.requestResources()
            requestRes.cmd = 17
            requestRes.code = 1
            self.writeData(requestRes.SerializeToString())
            self.progressDialog = QProgressDialog(self)
            self.progressDialog.setRange(0, 0)
            self.progressDialog.setWindowTitle('正在等待对面回应')
            self.progressDialog.canceled.connect(self.progressDialogpdCanceled)
            self.isMatch = True  # 是否匹配中
            self.progressDialog.exec()

    def progressDialogpdCanceled(self):
        requestRes = base_pb2.requestResources()
        requestRes.cmd = 17
        requestRes.code = 2  # 取消匹配
        self.writeData(requestRes.SerializeToString())
        self.isMatch = False

    def wantToLose(self):
        if self.isBeginGame == False:
            return
        reply = QMessageBox.question(self, '提示', '您确定要认输吗?', QMessageBox.Yes | QMessageBox.No)
        if reply == QMessageBox.No:
            return
        else:
            self.lose()

    def lose(self):
        if self.myColor == 'black':
            self.showResult('white')
        elif self.myColor == 'white':
            self.showResult('black')
        w_w = base_pb2.whoWin()
        w_w.cmd = 15
        w_w.code = 2
        w_w.win = ''
        self.writeData(w_w.SerializeToString())

    def draw(self):
        if self.isBeginGame == False:
            return
        reply = QMessageBox.question(self, '提示', '您确定要和棋吗?', QMessageBox.Yes | QMessageBox.No)
        if reply == QMessageBox.No:
            return
        w_w = base_pb2.whoWin()
        w_w.cmd = 15
        w_w.code = 3
        w_w.win = ''
        self.writeData(w_w.SerializeToString())

    def setUp(self):
        self.setting.setWindowModality(Qt.WindowModal)
        self.setting.show()

    def withdraw(self):
        w_d = base_pb2.withDraw()
        w_d.cmd = 16
        if self.isCanXia:  # 自己回合
            w_d.nums = 2
        else:
            w_d.nums = 1
        self.writeData(w_d.SerializeToString())

    def huiBack(self):
        if len(self.history) == 0:
            return
        m = self.history.pop()
        a = self.history2.pop()

        self.chessboard[m.y][m.x] = None
        m.close()
        a.close()
        if self.turnChessColor == "black":
            self.turnChessColor = "white"
        else:
            self.turnChessColor = "black"

    def readData(self):
        headerData = self.sock.read(4)
        dataLen = int.from_bytes(headerData, byteorder='little', signed=False)
        readData = self.sock.read(dataLen)
        if readData == b'':
            return
        cmd = base_pb2.cmd()
        cmd.ParseFromString(readData)
        if cmd.c != 9:
            print("cmd: %s" % cmd.c)
        self.switch[cmd.c](readData)

    def writeData(self, data):
        print("writeData")
        headerData = QByteArray()
        wData = QDataStream(headerData, QIODevice.WriteOnly)
        wData.writeUInt32(socket.htonl(len(data)))
        self.sock.write(headerData + data)
        self.sock.flush()

    def cmd7(self, data):
        # 获取对方的棋子位置，并显示
        s_g_p = base_pb2.server_gobang_position()
        s_g_p.ParseFromString(data)
        self.chessman = Chessman(color=self.turnChessColor, parent=self)
        self.chessman.setIndex(s_g_p.x, s_g_p.y)
        self.chessman.move()
        self.isCanXia = True
        self.label.setCursor(QCursor(Qt.PointingHandCursor))

    def cmd9(self, data):
        # 显示在线用户列表
        s_o_i = base_pb2.server_online_infor()
        s_o_i.ParseFromString(data)
        # print(s_o_i.people)
        self.peopleList = s_o_i.people
        self.peopleList.remove(self.username)
        slm = QStringListModel()
        slm.setStringList(self.peopleList)
        self.listView.setModel(slm)
        self.lcdNumber_3.display(len(self.peopleList) + 1)

    def cmd10(self, data):
        s_g_i = base_pb2.server_game_invite()
        s_g_i.ParseFromString(data)
        reply = QMessageBox.question(self, '提示', '"%s"请求对战' % s_g_i.people, QMessageBox.Yes | QMessageBox.No)
        c_g_i = base_pb2.client_game_invite()
        c_g_i.cmd = 11
        if reply == QMessageBox.Yes:
            c_g_i.code = 1
            self.isBeginGame = True
            self.myColor = "white"
            self.initBeginGame()
        else:
            c_g_i.code = 0
            pass
        self.writeData(c_g_i.SerializeToString())

    def cmd12(self, data):
        s_g_i = base_pb2.server_game_isInvite()
        s_g_i.ParseFromString(data)
        reply = ""
        self.fightBox.close()
        if s_g_i.code == 1:
            reply = '对方接受了请求'
            self.isBeginGame = True
            self.isCanXia = True
            self.myColor = 'black'
            self.initBeginGame()
        elif s_g_i.code == 0:
            reply = "对方拒绝了请求"

        QMessageBox.information(self, '提示', reply, QMessageBox.Ok)

    def cmd13(self, data):
        c_m = base_pb2.chatMessage()
        c_m.ParseFromString(data)
        messagetype = c_m.type
        string = c_m.data
        messageTime = c_m.time
        if self.isCloseChat == True:
            return
        s = '对手 %s
%s' % (messageTime, string)
        if messagetype == 1:  # 文字信息
            self.insertChatMessage(s)

    def cmd14(self, data):
        s_u_i = base_pb2.server_user_infor()
        s_u_i.ParseFromString(data)
        self.userGameInfor = s_u_i
        s = """昵称:%-5s\n积分：%-5s  级别：%-5s\n总盘数：%-4s 游戏币：%-5s\n胜：%-5s 负：%-5s平：%-5s
        """ % (
            s_u_i.name, s_u_i.integral, s_u_i.level, s_u_i.numsGame, s_u_i.gameCurrency, s_u_i.win, s_u_i.lose,
            s_u_i.draw)
        img_src = r'source\avatar\%s.png' % s_u_i.avatar
        img = QPixmap(img_src)
        if s_u_i.code == 1:
            self.label_11.setPixmap(img)
            self.label_11.setFixedSize(img.size())
            self.label_11.setScaledContents(True)
            self.label_3.setText(s)
        elif s_u_i.code == 2:
            self.label_12.setPixmap(img)
            self.label_12.setFixedSize(img.size())
            self.label_12.setScaledContents(True)
            self.label_4.setText(s)

    def cmd15(self, data):  # 和棋
        w_w = base_pb2.whoWin()
        w_w.ParseFromString(data)
        if w_w.code == 2:
            self.showResult(self.myColor)
        elif w_w.code == 3:
            reply = QMessageBox.question(self, '提示', '对方请求和棋', QMessageBox.Yes | QMessageBox.No)
            if reply == QMessageBox.Yes:
                w_w.code = 4
                self.showResult(self.myColor)
            else:
                w_w.code = 5
            self.writeData(w_w.SerializeToString())
        elif w_w.code == 4:
            self.showResult(self.myColor)

        elif w_w.code == 5:
            reply = QMessageBox.question(self, '提示', '对方不同意和棋', QMessageBox.Yes)
            if reply == QMessageBox.Yes:
                return

    def cmd16(self, data):  # 悔棋
        w_d = base_pb2.withDraw()
        w_d.ParseFromString(data)
        requestRes = base_pb2.responseResources()
        requestRes.cmd = 17
        reply = QMessageBox.question(self, '提示', '对方请求悔棋', QMessageBox.Yes | QMessageBox.No)
        if reply == QMessageBox.Yes:
            requestRes.code = 3
            if w_d.nums == 1:
                self.huiBack()
                self.isCanXia = False
                self.label.setCursor(QCursor(Qt.ForbiddenCursor))
            elif w_d.nums == 2:
                self.huiBack()
                self.huiBack()
        else:
            requestRes.code = 4
        self.writeData(requestRes.SerializeToString())

    def cmd18(self, data):
        resResources = base_pb2.responseResources()
        resResources.ParseFromString(data)
        if resResources.code == 1:
            self.progressDialog.close()
            if resResources.code2 == 1:
                self.myColor = 'white'
                self.isCanXia = False
                self.label.setCursor(QCursor(Qt.PointingHandCursor))
            elif resResources.code2 == 2:
                self.myColor = 'black'
                self.isCanXia = True
                self.label.setCursor(QCursor(Qt.ForbiddenCursor))
            self.isBeginGame = True
            self.isMatch = False
            self.initBeginGame()
        elif resResources.code == 2:
            if resResources.code2 == 1:
                if self.isCanXia:  # 自己回合
                    self.huiBack()
                    self.huiBack()
                else:
                    self.huiBack()
                    self.isCanXia = True
                    self.label.setCursor(QCursor(Qt.PointingHandCursor))
            elif resResources.code2 == 2:
                QMessageBox.information(self, '提示', '对方不同意悔棋', QMessageBox.Ok)

    def insertChatMessage(self, s):
        self.textBrowser.append(s)
        self.textBrowser.moveCursor(QTextCursor.End)
        self.comboBox.lineEdit().clear()

    def listViewClicked(self, qModelIndex):
        if self.isBeginGame == True:
            QMessageBox.information(self, '警告', '您正在游戏中', QMessageBox.Ok)
            return
        self.withusername = self.peopleList[qModelIndex.row()]
        print(self.withusername)
        c_c_g = base_pb2.client_create_game()
        c_c_g.cmd = 8
        c_c_g.withUsername = self.withusername
        self.writeData(c_c_g.SerializeToString())

        # 创建一个问答框，注意是Question
        self.fightBox = QMessageBox(QMessageBox.Question, '对局请求中', '正在等待对方响应中...', parent=self)
        # 添加按钮，可用中文
        cancel = self.fightBox.addButton('取消对局', QMessageBox.YesRole)

        # 设置消息框中内容前面的图标
        self.fightBox.setIcon(1)

        # 显示该问答框
        self.fightBox.show()
        if self.fightBox.clickedButton() == cancel:
            self.fightBox.close()

    def sendGameRes(self, code, win):
        res = base_pb2.whoWin()
        res.cmd = 15
        res.code = code
        res.win = win
        self.writeData(res.SerializeToString())

    def messageSend(self):
        if self.comboBox.lineEdit().text() == "":
            return
        if self.isCloseChat == True:
            s = '您已开启消息屏蔽功能，暂时无法发送消息，如有需要请在设置面板中开启'
            self.insertChatMessage(s)
            return
        c_m = base_pb2.chatMessage()
        c_m.cmd = 13
        c_m.type = 1
        c_m.time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
        c_m.data = self.comboBox.lineEdit().text()
        s = '我 %s
%s' % (
            c_m.time, self.comboBox.lineEdit().text())
        self.insertChatMessage(s)
        self.writeData(c_m.SerializeToString())

    def soundSetting(self, state):
        if state == 2:  # 点击，关闭
            self.sound_piece = QSound('')
            self.sound_win = QSound('')
            self.sound_defeated = QSound('')
        else:
            self.sound_piece = self.sound_piece_t
            self.sound_win = self.sound_win_t
            self.sound_defeated = self.sound_defeated_t

    def chatSetting(self, state):
        if state == 2:  # 点击，关闭
            self.isCloseChat = True
            s = '您已开启消息屏蔽功能'
        else:
            self.isCloseChat = False
            s = '您已取消消息屏蔽功能'
        self.insertChatMessage(s)

    def getPoint(self):
        '''
        返回落子位置
        :return:
        '''
        # 简单实现：返回一个空白交点
        # for i in range(19):
        #     for j in range(19):
        #         if self.chessboard[i][j] == None:
        #             return QPoint(j, i)
        #
        #  没有找到合适的点
        white_score = [[0 for i in range(15)] for j in range(15)]
        black_score = [[0 for i in range(15)] for j in range(15)]

        for i in range(15):
            for j in range(15):
                if self.chessboard[i][j] != None:
                    continue
                # 模拟落子
                self.chessboard[i][j] = Chessman(color='white', parent=self)
                white_score[i][j] = self.getPointScore(j, i, 'white')
                self.chessboard[i][j].close()
                self.chessboard[i][j] = None
                self.chessboard[i][j] = Chessman(color='black', parent=self)
                black_score[i][j] = self.getPointScore(j, i, 'black')
                self.chessboard[i][j].close()
                self.chessboard[i][j] = None
        # 将二维坐标转换成以为进行计算
        r_white_score = []
        r_black_score = []
        for i in white_score:
            r_white_score.extend(i)
        for i in black_score:
            r_black_score.extend(i)

        # 找到分数最大值
        score = [max(x, y) for x, y in zip(r_white_score, r_black_score)]

        # 找到分数做大的下标
        chess_index = score.index(max(score))

        # print(score, '\n', max(score))

        y = chess_index // 15
        x = chess_index % 15

        return QPoint(x, y)

    def aiBegin(self):
        self.isBeginGame = True
        self.isCanXia = True
        self.myColor = "black"
        self.ai = True
        self.initBeginGame()
        s_u_i = base_pb2.server_user_infor()
        s_u_i.code = 2
        s_u_i.name, s_u_i.integral, s_u_i.level, s_u_i.numsGame, s_u_i.gameCurrency, s_u_i.win, s_u_i.lose, s_u_i.draw, s_u_i.avatar = (
            '电脑', 9999, 9999, 9999, 9999, 9999, 0, 0, '1')
        self.cmd14(s_u_i.SerializeToString())

    def autoDown(self):
        point = self.getPoint()
        s_g_p = base_pb2.server_gobang_position()
        s_g_p.x = point.x()
        s_g_p.y = point.y()
        self.cmd7(s_g_p.SerializeToString())

    def getPointScore(self, x, y, color):
        '''
        返回每个点的得分
        y:行坐标
        x:列坐标
        color：棋子颜色
        :return:
        '''
        # 分别计算点周围5子以内，空白、和同色的分数
        blank_score = 0
        color_score = 0

        # 记录每个方向的棋子分数
        blank_score_plus = [0, 0, 0, 0]  # 横向 纵向 正斜线 反斜线
        color_score_plus = [0, 0, 0, 0]

        # 横线
        # 右侧
        i = x  # 横坐标
        j = y  # 纵坐标
        while i < 19:
            if self.chessboard[j][i] is None:
                blank_score += 1
                blank_score_plus[0] += 1
                break
            elif self.chessboard[j][i].color == color:
                color_score += 1
                color_score_plus[0] += 1
            else:
                break
            if i >= x + 4:
                break
            i += 1
        # print('123123')
        # 左侧
        i = x  # 横坐标
        j = y  # 纵坐标
        while i >= 0:
            if self.chessboard[j][i] is None:
                blank_score += 1
                blank_score_plus[0] += 1
                break
            elif self.chessboard[j][i].color == color:
                color_score += 1
                color_score_plus[0] += 1
            else:
                break
            if i <= x - 4:
                break
            i -= 1

        # 竖线
        # 上方
        i = x  # 横坐标
        j = y  # 纵坐标
        while j >= 0:
            if self.chessboard[j][i] is None:
                blank_score += 1
                blank_score_plus[1] += 1
                break
            elif self.chessboard[j][i].color == color:
                color_score += 1
                color_score_plus[1] += 1
            else:
                break
            if j <= y - 4:
                break
            j -= 1
        # 竖线
        # 下方
        i = x  # 横坐标
        j = y  # 纵坐标
        while j < 19:
            if self.chessboard[j][i] is None:
                blank_score += 1
                blank_score_plus[1] += 1
                break
            elif self.chessboard[j][i].color == color:
                color_score += 1
                color_score_plus[1] += 1
            else:
                break

            if j >= y + 4:  # 最近五个点
                break
            j += 1
        # 正斜线
        # 右上
        i = x
        j = y
        while i < 19 and j >= 0:
            if self.chessboard[j][i] is None:
                blank_score += 1
                blank_score_plus[2] += 1
                break
            elif self.chessboard[j][i].color == color:
                color_score += 1
                color_score_plus[2] += 1
            else:
                break

            if i >= x + 4:  # 最近五个点
                break
            i += 1
            j -= 1
        # 左下
        i = x
        j = y
        while j < 19 and i >= 0:
            if self.chessboard[j][i] is None:
                blank_score += 1
                blank_score_plus[2] += 1
                break
            elif self.chessboard[j][i].color == color:
                color_score += 1
                color_score_plus[2] += 1
            else:
                break

            if j >= y + 4:  # 最近五个点
                break
            i -= 1
            j += 1
        # 反斜线
        # 左上
        i = x
        j = y
        while i >= 0 and j >= 0:
            if self.chessboard[j][i] is None:
                blank_score += 1
                blank_score_plus[3] += 1
                break
            elif self.chessboard[j][i].color == color:
                color_score += 1
                color_score_plus[3] += 1
            else:
                break
            if i <= x - 4:
                break
            i -= 1
            j -= 1
        # 右上
        i = x
        j = y
        while i < 19 and j < 19:
            if self.chessboard[j][i] is None:
                blank_score += 1
                blank_score_plus[3] += 1
                break
            elif self.chessboard[j][i].color == color:
                color_score += 1
                color_score_plus[3] += 1
            else:
                break
            if i >= x + 4:
                break
            i += 1
            j += 1

        for k in range(4):
            if color_score_plus[k] >= 5:
                return 100

        # color_score *= 5
        return max([x + y for x, y in zip(color_score_plus, blank_score_plus)])


import tensorflow as tf
import keras
import pandas as pd
from keras.models import Sequential
from keras.layers import Dense
import os
from keras import optimizers
from keras import regularizers

#read in data using pandas
cwd = os.getcwd()
train_df = pd.read_csv(cwd + "/data.csv")
#validation_df = pd.read_csv(cwd + "/dataTest.csv")

#view data structure
train_df.head()
#validation_df.head()

#create a dataframe with all training data except the target column
train_X = train_df.drop(columns=['Class'])
#validation_X = validation_df.drop(columns=['Class'])

#check that the target variable has been removed
train_X.head()
#validation_X.head()

from keras.utils import to_categorical

#one-hot encode target column
train_y = to_categorical(train_df.Class)
#validation_y = to_categorical(validation_df.Class)

#vcheck that target column has been converted
train_y[0:5]
#validation_y[0:5]

#create model
model = Sequential()

#get number of columns in training data
n_cols = train_X.shape[1]

#early_stop = keras.callbacks.EarlyStopping(monitor='val_loss', min_delta=0, patience=3, verbose=1, mode='auto')
#model.add(Dense(52, input_dim=52,
#            kernel_regularizer=regularizers.l2(0.01),
#            activity_regularizer=regularizers.l1(0.01)))
#add layers to model
model.add(Dense(52, activation='sigmoid', input_shape=(n_cols,)))
model.add(Dense(20, activation='sigmoid'))
model.add(Dense(20, activation='sigmoid'))
model.add(Dense(4, activation='softmax'))

#opt=optimizers.SGD(lr=0.1, momentum=0.1)
#compile model using accuracy to measure model performance
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

#train model
model.fit(train_X, train_y, epochs=200, validation_split=1)

print(model.evaluate(train_X, train_y, verbose=0))

# serialize model to JSON
model_json = model.to_json()
with open("model.json", "w") as json_file:
    json_file.write(model_json)
# serialize weights to HDF5
model.save_weights("model.h5")
print("Saved model to disk")


"""Validate and parse expressions into SQLAlchemy expressions"""

import functools
from dataclasses import dataclass
from datetime import date, datetime

from lark import GrammarError
from sqlalchemy import func
from sqlalchemy.sql.expression import ClauseElement

from sqlalchemy_recipe.dbinfo import DBInfo, ReflectedTable
from sqlalchemy_recipe.expression.datatype import DataType
from sqlalchemy_recipe.expression.grammar import make_columns_for_table
from sqlalchemy_recipe.expression.transformer import TransformToSQLAlchemyExpression
from sqlalchemy_recipe.expression.validator import SQLALchemyValidator

@dataclass
class BuilderResponse:
    datatype: DataType
    expression: ClauseElement
    is_aggr: bool

class SQLAlchemyBuilder:
    def __init__(self, dbinfo: DBInfo, reflected_table: ReflectedTable):
        """Parse an expression by building a custom grammar that
        uses the columns in a table.

        Args:
            dbinfo (DBInfo): _description_
            table (Table): _description_
            grammar (str): _description_
        """
        self.dbinfo = dbinfo
        self.reflected_table = reflected_table
        self.drivername = self.dbinfo.engine.url.drivername
        self.transformer = TransformToSQLAlchemyExpression(
            reflected_table, self.dbinfo.engine.url.drivername
        )

    @functools.lru_cache(maxsize=None)
    def parse(
        self,
        text: str,
        forbid_aggregation: bool = False,
        enforce_aggregation: bool = False,
        convert_dates_with: str = None,
        convert_datetimes_with: str = None,
        default_aggregation=func.sum,
        debug: bool = False,
    ) -> BuilderResponse:
        """Return a parse tree for text

        Args:
            text (str): A field expression
            forbid_aggregation (bool, optional):
              The expression may not contain aggregations. Defaults to False.
            enforce_aggregation (bool, optional):
              Wrap the expression in an aggregation if one is not provided. Defaults to False.
            debug (bool, optional): Show some debug info. Defaults to False.
            convert_dates_with (str, optional): A converter to use for date fields
            convert_datetimes_with (str, optional): A converter to use for datetime fields

        Raises:
            GrammarError: A description of any errors and where they occur

        Returns:
            BuilderResponse
        """
        tree = self.reflected_table.parser.parse(text, start="col")
        validator = SQLALchemyValidator(
            text, forbid_aggregation, self.dbinfo.engine.url.drivername
        )
        # Populate errors and check aggregation
        validator.visit(tree)
        datatype = validator.last_datatype

        if validator.errors:
            if debug:
                print("".join(map(str, validator.errors)))
                print("Tree:\n" + tree.pretty())
            # Validator.errors is a list of VisitError
            ge = GrammarError("".join(map(str, validator.errors)))
            ge.validator_errors = validator.errors
            raise ge

        if debug:
            print("Tree:\n" + tree.pretty())

        self.transformer.text = text
        self.transformer.convert_dates_with = convert_dates_with
        self.transformer.convert_datetimes_with = convert_datetimes_with

        # Build an expression from the bottom of the tree up.
        expr = self.transformer.transform(tree)

        # Expressions that return literal values can't be labeled
        # Possibly we could wrap them in text() but this may be unsafe
        # instead we will disallow them.
        if isinstance(expr, (str, float, int, date, datetime)):
            raise GrammarError("Must return an expression, not a constant value")

        if (
            enforce_aggregation
            and not validator.found_aggregation
            and datatype == "num"
        ):
            return BuilderResponse(
                expression=default_aggregation(expr),
                datatype=datatype,
                is_aggr=validator.found_aggregation,
            )
        else:
            return BuilderResponse(
                expression=expr, datatype=datatype, is_aggr=validator.found_aggregation
            )


import logging
from urllib.parse import urljoin

from cdpcli.exceptions import WorkloadServiceDiscoveryError

LOG = logging.getLogger('cdpcli.extensions.workload')

def register(operation_callers, operation_model):
    """
    Register an extension to run before or after the CLI command.
    To replace the original CLI caller:
    * operation_callers.insert(0, ReplacementCaller())
    * return False by the ReplacementCaller.invoke(...)
    """
    operation_callers.insert(0, WorkloadServiceDiscovery())

def set_workload_access_token(iam_client,
                              parsed_globals,
                              workload_name,
                              environment_crn):
    """
    Set global Access Token and Endpoint URL for Workload and Environment CRN
    """
    req_params = {
        'workloadName': workload_name,
        'environmentCrn': environment_crn
    }
    response = iam_client.generate_workload_auth_token(**req_params)

    workload_url = response.get('endpointUrl', None)
    workload_access_token = response.get('token', None)
    LOG.debug('Workload service-discovery succeeded. '
              'endpoint_url=%s, access_token=%s...',
              workload_url,
              workload_access_token[0:16] if workload_access_token else None)

    # remove path from the workload URL
    if workload_url is not None:
        workload_url = urljoin(workload_url, '/')
    # add 'Bearer' prefix to access token
    if workload_access_token is not None and \
       not workload_access_token.startswith('Bearer '):
        workload_access_token = 'Bearer ' + workload_access_token

    # Set the workload URL and access-token in parsed_globals, so the
    # next CLI operation caller will use them to make the connection to workload.
    parsed_globals.endpoint_url = workload_url
    parsed_globals.access_token = workload_access_token

class WorkloadServiceDiscovery(object):
    def invoke(self,
               client_creator,
               operation_model,
               parameters,
               parsed_args,
               parsed_globals):
        if parsed_globals.access_token:
            # access-token was explicitly set by user,
            # skip the workload service-discovery
            LOG.debug('Skip workload service-discovery')
            return True

        LOG.debug('Run workload service-discovery')
        service_name = operation_model.service_model.service_name
        if not service_name:
            raise WorkloadServiceDiscoveryError(err_msg='Missing service name')

        if service_name == 'dfworkload':
            workload_name = 'DF'
            environment_crn = parameters.get('environmentCrn', None)
            if not environment_crn:
                raise WorkloadServiceDiscoveryError(err_msg='Missing environment CRN')
        else:
            raise WorkloadServiceDiscoveryError(
                err_msg='Unknown service name \'%s\'' % service_name)

        iam_client = client_creator('iam')
        set_workload_access_token(iam_client, parsed_globals, workload_name,
                                  environment_crn)

        # The command processing is not finished, continue with other operation callers.
        return True


#  By submitting this assignment, I agree to the following:
#  "Aggies do not lie, cheat, or steal, or tolerate those who do"
#  "I have not given or received any unauthorized aid on this assignment"
#
# Name:        Chase Albright
# Section:     480
# Assignment:  Lab Assignment 10b
# Date:        27 October 2020

import csv

# program to open and read the csv file
file = input('Enter the name of the loan csv file you wish to visualize: ')

monthList = []
balanceList = []

#open the file
with open(file,'r') as csvfile:
    #read each line seperated by a comma
    readCSV = csv.reader(csvfile, delimiter=',')
    #loop to run through each and while loop to limit to 30 values
    next(readCSV) # skip the label field
    for row in readCSV:
        monthList.append(row[0])
        balanceList.append(row[1])
#print(monthList)
#print(balanceList)

# convert list vals into floats
monthNum = []
for numM in monthList:
    x = float(numM)
    monthNum.append(x)

balNum = []
for numB in balanceList:
    y = float(numB)
    balNum.append(y)

#print(monthNum)
#print(balNum)

#limits length to 30 points if longer than 30
if len(monthNum) >= 30:
    del monthNum[29:]
    del balNum[29:]

# here we begin the graphing
import matplotlib.pyplot as plt

plt.scatter(monthNum,balNum,color='k',label='Balance')
plt.plot(monthNum,balNum,'--k')
plt.xlabel('Months')
plt.ylabel('Principal Balance')
plt.title('Balance vs. Time in Months')
plt.legend()
plt.show()


""" Unit tests for surface simulations

"""

import logging
import unittest

import astropy.units as u
import numpy
from astropy.coordinates import SkyCoord

from rascil.data_models.memory_data_models import Skycomponent
from rascil.data_models.polarisation import PolarisationFrame
from rascil.processing_components.imaging.primary_beams import (
    create_vp_generic_numeric,
    create_vp,
)
from rascil.processing_components.simulation import create_named_configuration
from rascil.processing_components.simulation.surface import (
    simulate_gaintable_from_zernikes,
)
from rascil.processing_components.visibility.base import create_blockvisibility
from rascil.processing_components import create_image

log = logging.getLogger("rascil-logger")

log.setLevel(logging.WARNING)

class TestSurface(unittest.TestCase):
    def setUp(self):

        from numpy.random import default_rng

        self.rng = default_rng(1805550721)

        from rascil.data_models.parameters import rascil_path, rascil_data_path

        self.doplot = True

        self.midcore = create_named_configuration("MID", rmax=100.0)
        self.nants = len(self.midcore.names)
        self.results_dir = rascil_path("test_results")
        self.ntimes = 100
        interval = 10.0
        self.times = numpy.arange(0.0, float(self.ntimes)) * interval
        self.times *= numpy.pi / 43200.0

        self.frequency = numpy.array([1.4e9])
        self.channel_bandwidth = numpy.array([1e7])
        self.phasecentre = SkyCoord(
            ra=+15.0 * u.deg, dec=-45.0 * u.deg, frame="icrs", equinox="J2000"
        )
        self.vis = create_blockvisibility(
            self.midcore,
            self.times,
            self.frequency,
            channel_bandwidth=self.channel_bandwidth,
            phasecentre=self.phasecentre,
            weight=1.0,
            polarisation_frame=PolarisationFrame("stokesI"),
        )
        self.vis["vis"].data *= 0.0

        # Create model
        self.model = create_image(
            npixel=512,
            cellsize=0.001,
            polarisation_frame=PolarisationFrame("stokesI"),
            frequency=self.frequency,
            channel_bandwidth=self.channel_bandwidth,
            phasecentre=self.phasecentre,
        )

    def test_simulate_gaintable_from_voltage_patterns(self):
        numpy.random.seed(18051955)
        offset_phasecentre = SkyCoord(
            ra=+15.0 * u.deg, dec=-44.58 * u.deg, frame="icrs", equinox="J2000"
        )
        component = [
            Skycomponent(
                frequency=self.frequency,
                direction=offset_phasecentre,
                polarisation_frame=PolarisationFrame("stokesI"),
                flux=[[1.0]],
            )
        ]

        key_nolls = [3, 5, 6, 7]
        vp_list = list()
        vp_list.append(
            create_vp_generic_numeric(
                self.model,
                pointingcentre=None,
                diameter=15.0,
                blockage=0.0,
                taper="gaussian",
                edge=0.03162278,
                padding=2,
                use_local=True,
            )
        )
        vp_coeffs = numpy.ones([self.nants, len(key_nolls) + 1])
        for inoll, noll in enumerate(key_nolls):
            zernike = {"coeff": 1.0, "noll": noll}
            vp_coeffs[:, inoll + 1] = self.rng.normal(0.0, 0.03, self.nants)
            vp_list.append(
                create_vp_generic_numeric(
                    self.model,
                    pointingcentre=None,
                    diameter=15.0,
                    blockage=0.0,
                    taper="gaussian",
                    edge=0.03162278,
                    zernikes=[zernike],
                    padding=2,
                    use_local=True,
                )
            )

        for vp in vp_list:
            assert (
                vp.image_acc.wcs.wcs.ctype[0] == "AZELGEO long"
            ), vp.image_acc.wcs.wcs.ctype[0]
        gt = simulate_gaintable_from_zernikes(self.vis, component, vp_list, vp_coeffs)

        import matplotlib.pyplot as plt

        assert gt[0].gain.shape == (self.ntimes, self.nants, 1, 1, 1), gt[0].gain.shape

        plt.clf()
        for ant in range(self.nants):
            plt.plot(gt[0].time, 1.0 / numpy.real(gt[0].gain[:, ant, 0, 0, 0]), ".")
        plt.xlabel("Time (s)")
        plt.ylabel("Gain")
        plt.show(block=False)

if __name__ == "__main__":
    unittest.main()


from typing import NamedTuple

class Resources(NamedTuple):
    ore: int = 0
    clay: int = 0
    obsidian: int = 0
    geode: int = 0

class Robots(NamedTuple):
    ore: int = 0
    clay: int = 0
    obsidian: int = 0
    geode: int = 0

class Cost(NamedTuple):
    ore: Resources = Resources(ore=4)
    clay: Resources = Resources(ore=2)
    obsidian: Resources = Resources(ore=3, clay=14)
    geode: Resources = Resources(ore=2, obsidian=7)

class Round(NamedTuple):
    i: int = 0
    resources: Resources = Resources()
    robots: Robots = Robots(ore=1)

def can_buy(resources: Resources, cost: Resources) -> bool:
    for field in Resources._fields:
        if getattr(resources, field) < getattr(cost, field):
            return False
    return True

def minus_cost(resources: Resources, cost: Resources) -> Resources:
    data = {}
    for field in Resources._fields:
        data[field] = getattr(resources, field) - getattr(cost, field)
    nr = Resources(**data)
    return nr

def add_robot(robots: Robots, robot: str) -> Robots:
    rs = robots._asdict()
    rs[robot] += 1
    return Robots(**rs)

def add_resources(left: Resources, right: Robots) -> Resources:
    left = left._asdict()
    right = right._asdict()
    res = {}
    for key in left:
        res[key] = left[key] + right[key]
    return Resources(**res)

def run(cost: Cost) -> int:

    cutoff = 32
    res = 0
    seen = {}
    states = [Round()]
    while states:
        state = states.pop()
        if state.i not in seen:
            seen[state.i] = (state.resources.geode, state.robots.geode)
        else:
            seen_geode_re, seen_geode_ro = seen[state.i]
            if seen_geode_re > state.resources.geode and seen_geode_ro > state.robots.geode:
                continue
            if state.resources.geode > seen_geode_re and state.robots.geode > seen_geode_ro:
                seen[state.i] = (state.resources.geode, state.robots.geode)
        round_no = state.i + 1
        if round_no > cutoff:
            if state.resources.geode > res:
                res = state.resources.geode
            continue

        # spend resources
        resources = state.resources
        robots = state.robots

        # 1. do nothing
        new_resources = [resources]
        new_robots = [robots]

        # 2. try to buy geode
        for field in ("geode", "obsidian", "clay", "ore"):
            if can_buy(resources, getattr(cost, field)):
                new_resources.append(minus_cost(
                    resources, getattr(cost, field)))
                new_robots.append(add_robot(robots, field))
                if field == "geode":  # or field == "obsidian":
                    new_resources.pop(0)
                    new_robots.pop(0)
                    break
                last_robots = new_robots[-1]
                if field == "clay" and all(
                   last_robots.clay > c for c in (cost.geode.clay, cost.obsidian.clay, cost.ore.clay, cost.clay.clay)
                   ):
                    new_robots.pop()
                    new_resources.pop()
                if field == "ore" and all(
                            last_robots.ore > c for c in (cost.geode.ore, cost.obsidian.ore, cost.ore.ore, cost.clay.ore)
                ):
                    new_robots.pop()
                    new_resources.pop()
                if field == "obsidian" and all(
                            last_robots.obsidian > c for c in (cost.geode.obsidian, cost.obsidian.obsidian, cost.ore.obsidian, cost.clay.obsidian)
                ):
                    new_robots.pop()
                    new_resources.pop()

                    # collect resources
        new_resources = [add_resources(r, robots) for r in new_resources]

        for resource, robot in zip(new_resources, new_robots):
            states.append(Round(round_no, resources=resource, robots=robot))
    return res

# Blueprint 2:
#  Each ore robot costs 2 ore.
#  Each clay robot costs 3 ore.
#  Each obsidian robot costs 3 ore and 8 clay.
#   Each geode robot costs 3 ore and 12 obsidian.
# Blueprint 2: Each ore robot costs 3 ore. Each clay robot costs 3 ore. Each obsidian robot costs 2 ore and 20 clay. Each geode robot costs 2 ore and 20 obsidian.
# Blueprint 3: Each ore robot costs 3 ore. Each clay robot costs 3 ore. Each obsidian robot costs 3 ore and 16 clay. Each geode robot costs 3 ore and 9 obsidian.
if __name__ == "__main__":
    res = 1
    costs = [
        Cost(ore=Resources(ore=2), clay=Resources(ore=4), obsidian=Resources(
            ore=4, clay=20), geode=Resources(ore=3, obsidian=14)),
        Cost(ore=Resources(ore=3), clay=Resources(ore=3), obsidian=Resources(
            ore=2, clay=20), geode=Resources(ore=2, obsidian=20)),
        Cost(ore=Resources(ore=3), clay=Resources(ore=3), obsidian=Resources(
            ore=3, clay=16), geode=Resources(ore=3, obsidian=9)),
    ]
    for cost in costs:
        print(cost)
        res *= run(cost)
    print(res)


"""
Translate any word to any language
"""

import tkinter as tk
from googletrans import Translator

win = tk.Tk()
win.title("Translator")
win.geometry("400x100")

# Function to translate

def translation():
    word = entry.get()
    translator = Translator(service_urls =['translate.google.com',])
    translation1 = translator.translate(word, dest="hi")
    label1 = tk.Label(win, text=f"Translated in Hindi: {translation1.text}", bg="yellow")
    label1.grid(row=2, column=0)

label = tk.Label(win, text="Enter Word: ")
label.grid(row=0, column=0, sticky="W")

entry = tk.Entry(win)
entry.grid(row=1, column=0)

button=tk.Button(win, text="Translate", command=translation)
button.grid(row=1, column=2)

win.mainloop()


# %%
import os
import torch

import random
import shutil

import numpy as np
import torch
import matplotlib.pyplot as plt
import torchvision.transforms.functional as F
# %%
def del_folder(path, *args):
    '''
    delete the folder which path/version

    Args:
        path (str): path
        version (str): version
    '''
    if os.path.exists(os.path.join(path, *args)):
        shutil.rmtree(os.path.join(path, *args))

def make_folder(path, *args):
    '''
    make folder which path/version

    Args:
        path (str): path
        version (str): version
    '''
    if not os.path.exists(os.path.join(path, *args)):
        os.makedirs(os.path.join(path, *args))
        print("success make dir! where: %s " % os.path.join(path, *args))
    else:
        print("The target path already exists! where: %s " % os.path.join(path, *args))

def tensor2var(x, grad=False):
    '''
    put tensor to gpu, and set grad to false

    Args:
        x (tensor): input tensor
        grad (bool, optional):  Defaults to False.

    Returns:
        tensor: tensor in gpu and set grad to false
    '''
    if torch.cuda.is_available():
        x = x.cuda()
        x.requires_grad_(grad)
    return x

def var2tensor(x):
    '''
    put date to cpu

    Args:
        x (tensor): input tensor

    Returns:
        tensor: put data to cpu
    '''
    return x.data.cpu()

def var2numpy(x):
    return x.data.cpu().numpy()

def get_ckpt_path(config):
    '''
    get the finall train ckpt, for the pretrain model, or some task.

    Args:
        config (parser): parameters of the training

    Returns:
        string: final ckpt file path
    '''

    ckpt_path = os.path.join('/workspace/Walk_Video_PyTorch', 'logs', config.model, config.version, 'checkpoints')
    file_name = os.listdir(ckpt_path)[0]

    ckpt_file_path = os.path.join(ckpt_path, file_name)

    return ckpt_file_path

def plot(imgs, **imshow_kwargs):
    '''
    receive a tensor, and plot the figure

    Args:
        imgs (tensor): a batch of tensor
    '''

    if not isinstance(imgs[0], list):
        # Make a 2d grid even if there's just 1 row
        imgs = [imgs]

    num_rows = len(imgs)
    num_cols = len(imgs[0])

    _, axs = plt.subplots(nrows=num_rows, ncols=num_cols, squeeze=False)
    for row_idx, row in enumerate(imgs):
        for col_idx, img in enumerate(row):
            ax = axs[row_idx, col_idx]
            img = F.to_pil_image(img.to("cpu"))
            ax.imshow(np.asarray(img), **imshow_kwargs)
            ax.set(xticklabels=[], yticklabels=[], xticks=[], yticks=[])

    plt.tight_layout()


import Model.param as param
import torch.nn as nn

class PositionwiseFeedForward(nn.Module):
    def __init__(self):
        super(PositionwiseFeedForward, self).__init__()
        self.ffn = nn.Sequential(
            nn.Linear(param.embed_dim, param.d_ff),
            nn.ReLU(),
            nn.Dropout(param.dropout),
            nn.Linear(param.d_ff, param.embed_dim)
        )

    def forward(self, inputs):
        """
        :param inputs: (batch_size, q_len, embed_dim)的张量
        :return:       (batch_size, q_len, embed_dim)的张量
        """

        return nn.LayerNorm(param.embed_dim)(self.ffn(inputs) + inputs)
        # return nn.LayerNorm(param.embed_dim).cuda()(self.ffn(inputs) + inputs)


from keras.layers import Activation, Reshape, Lambda, concatenate, dot, add
from keras.layers import Conv1D, Conv2D, Conv3D
from keras.layers import MaxPool1D
from keras import backend as K

def non_local_block(ip, shield_computation=True, mode='embedded'):
    channel_dim = 1 if K.image_data_format() == 'channels_first' else -1
    ip_shape = K.int_shape(ip)

    if mode not in ['gaussian', 'embedded', 'dot', 'concatenate']:
        raise ValueError('`mode` must be one of `gaussian`, `embedded`, `dot` or `concatenate`')

    dim1, dim2, dim3 = None, None, None

    if len(ip_shape) == 3:  # time series data
        rank = 3
        batchsize, dim1, channels = ip_shape

    elif len(ip_shape) == 4:  # image data
        rank = 4

        if channel_dim == 1:
            batchsize, channels, dim1, dim2 = ip_shape
        else:
            batchsize, dim1, dim2, channels = ip_shape

    elif len(ip_shape) == 5:  # Video / Voxel data
        rank = 5

        if channel_dim == 1:
            batchsize, channels, dim1, dim2, dim3 = ip_shape
        else:
            batchsize, dim1, dim2, dim3, channels = ip_shape

    else:
        raise ValueError('Input dimension has to be either 3 (temporal), 4 (spatial) or 5 (spatio-temporal)')

    if mode == 'gaussian':  # Gaussian instantiation
        x1 = Reshape((-1, channels))(ip)  # xi
        x2 = Reshape((-1, channels))(ip)  # xj
        f = dot([x1, x2], axes=2)
        f = Activation('softmax')(f)

    elif mode == 'dot':  # Dot instantiation
        # theta path
        if rank == 3:
            theta = Conv1D(channels // 2, 1, padding='same', use_bias=False)(ip)
        elif rank == 4:
            theta = Conv2D(channels // 2, (1, 1), padding='same', use_bias=False)(ip)
        else:
            theta = Conv3D(channels // 2, (1, 1, 1), padding='same', use_bias=False)(ip)

        theta = Reshape((-1, channels // 2))(theta)

        # phi path
        if rank == 3:
            phi = Conv1D(channels // 2, 1, padding='same', use_bias=False)(ip)
        elif rank == 4:
            phi = Conv2D(channels // 2, (1, 1), padding='same', use_bias=False)(ip)
        else:
            phi = Conv3D(channels // 2, (1, 1, 1), padding='same', use_bias=False)(ip)

        phi = Reshape((-1, channels // 2))(phi)

        f = dot([theta, phi], axes=2)

        # scale the values to make it size invariant
        f = Lambda(lambda z: 1./batchsize * f)(f)

    elif mode == 'concatenate':  # Concatenation instantiation
        raise NotImplemented('Concatenation mode has not been implemented yet')

    else:  # Embedded Gaussian instantiation
        # theta path
        if rank == 3:
            theta = Conv1D(channels // 2, 1, padding='same', use_bias=False)(ip)
        elif rank == 4:
            theta = Conv2D(channels // 2, (1, 1), padding='same', use_bias=False)(ip)
        else:
            theta = Conv3D(channels // 2, (1, 1, 1), padding='same', use_bias=False)(ip)

        theta = Reshape((-1, channels // 2))(theta)

        # phi path
        if rank == 3:
            phi = Conv1D(channels // 2, 1, padding='same', use_bias=False)(ip)
        elif rank == 4:
            phi = Conv2D(channels // 2, (1, 1), padding='same', use_bias=False)(ip)
        else:
            phi = Conv3D(channels // 2, (1, 1, 1), padding='same', use_bias=False)(ip)

        phi = Reshape((-1, channels // 2))(phi)

        if shield_computation:
            # shielded computation
            phi = MaxPool1D()(phi)

        f = dot([theta, phi], axes=2)
        f = Activation('softmax')(f)

    # g path
    if rank == 3:
        g = Conv1D(channels // 2, 1, padding='same', use_bias=False)(ip)
    elif rank == 4:
        g = Conv2D(channels // 2, (1, 1), padding='same', use_bias=False)(ip)
    else:
        g = Conv3D(channels // 2, (1, 1, 1), padding='same', use_bias=False)(ip)

    g = Reshape((-1, channels // 2))(g)

    if shield_computation and mode == 'embedded':
        # shielded computation
        g = MaxPool1D()(g)

    # compute output path
    y = dot([f, g], axes=[2, 1])

    # reshape to input tensor format
    if rank == 3:
        y = Reshape((dim1, channels // 2))(y)
    elif rank == 4:
        if channel_dim == -1:
            y = Reshape((dim1, dim2, channels // 2))(y)
        else:
            y = Reshape((channels // 2, dim1, dim2))(y)
    else:
        if channel_dim == -1:
            y = Reshape((dim1, dim2, dim3, channels // 2))(y)
        else:
            y = Reshape((channels // 2, dim1, dim2, dim3))(y)

    # project filters
    y = Conv2D(channels, (1, 1), padding='same', use_bias=False)(y)

    # residual connection
    residual = add([ip, y])

    return residual


# imports
from os.path import join
from datetime import date, datetime
from argparse import ArgumentParser
import subprocess

# third party imports
import numpy as np
import torch
from PIL import Image

# project imports
from database.Allen import data_loader_subset as DL_Allen_subset
from database import databaseConfig

####################################
######## GLOBAL  PARAMETERS ########
####################################
NIFTY_REG_DIR = '/home/acasamitjana/Software_MI/niftyreg-git/build/'
ALADINcmd = NIFTY_REG_DIR + 'reg-apps' + '/reg_aladin'
REScmd = NIFTY_REG_DIR + 'reg-apps' + '/reg_resample'

MRI_PROCESSING = False
LABELS_PROCESSING = False
###################################
########### DATA LOADER ###########
###################################
if MRI_PROCESSING:
    data_loader = DL_Allen_subset.DataLoader(databaseConfig.ALLEN_subset)
    MASK_MRI = False
    COMPUTE_CROP_COORDINATES = False
    mx = 0
    Mx = 448
    my = 0
    My = 320
    for rid, subject in data_loader.subject_dict.items():
        # 1 Mask MRI
        if MASK_MRI:
            MM = subject.load_mri_mask()
            MI = subject.load_mri()
            MI[MM == 0] = 0
            img = Image.fromarray(MI, mode='L')
            img.save(subject.image_mri)

        # 2 Compute cropping coordinates
        if COMPUTE_CROP_COORDINATES:
            MM = subject.load_mri_mask()
            idx = np.where(MM > 0)
            mx = min(mx, np.min(idx[0]))
            Mx = max(Mx, np.max(idx[0]))
            my = min(my, np.min(idx[1]))
            My = max(My, np.max(idx[1]))

        # 3 Crop
        MI = subject.load_mri()
        MM = Image.open(subject.labels_mri)
        MM = np.array(MM)
        HI = subject.load_histo_affine()
        HM = Image.open(subject.labels_histo_affine)
        HM = np.array(HM)

        MMc = MM[mx:Mx, my:My]
        MIc = MI[mx:Mx, my:My]
        HMc = HM[mx:Mx, my:My]
        HIc = HI[mx:Mx, my:My]
        img = Image.fromarray(MIc, mode='L')
        img.save(subject.image_mri)
        img = Image.fromarray(MMc, mode='RGB')
        img.save(subject.labels_mri)
        img = Image.fromarray(HIc, mode='L')
        img.save(subject.image_histo_affine)
        img = Image.fromarray(HMc, mode='RGB')
        img.save(subject.labels_histo_affine)

    print('Done.')

#################
# WM Cerebellum #
#################
import pdb
if LABELS_PROCESSING:
    data_loader = DL_Allen_subset.DataLoader(databaseConfig.ALLEN_subset_cll)
    data_loader_all = DL_Allen_subset.DataLoader(databaseConfig.ALLEN_subset)

    for rid, subject in data_loader.subject_dict.items():
        subject_full = data_loader_all.subject_dict[rid.split('.')[0]]

        # 1. Histo Affine
        labels_cll = subject.load_histo_labels_affine()
        labels = subject_full.load_histo_labels_affine()

        new_labels = np.zeros(labels_cll.shape + (3,), dtype='uint8')
        new_labels[labels == 1, 0] = 255
        new_labels[labels == 2, 1] = 255
        new_labels[labels == 3, 2] = 255

        new_labels[labels_cll == 1, 0] = 255
        new_labels[labels_cll == 2, :] = 255

        img = Image.fromarray(new_labels, mode='RGB')
        img.save(subject_full.labels_histo_affine[:-4] + '_new.png')

        # 2. MRI
        labels_cll = subject.load_mri_labels()
        labels = subject_full.load_mri_labels()

        new_labels = np.zeros(labels_cll.shape + (3,), dtype='uint8')
        new_labels[labels == 1, 0] = 255
        new_labels[labels == 2, 1] = 255
        new_labels[labels == 3, 2] = 255

        new_labels[labels_cll == 1, 0] = 255
        new_labels[labels_cll == 2, :] = 255

        img = Image.fromarray(new_labels, mode='RGB')
        img.save(subject_full.labels_mri[:-4] + '_new.png')


from lib.args import get_hudi_args
from lib.spark import get_spark
from lib.table import calculate_partitions
from lib.constants import CHANGES_METADATA_FIELD_PREFIX, HUDI_METADATA_FIELD_PREFIX

cmd_args = get_hudi_args()
spark = get_spark()

# Load Hudi table
print(f">>> Load Hudi table from {cmd_args.hudi_path}...")
df = spark \
    .read \
    .format("hudi") \
    .load(f"{cmd_args.hudi_path}/*")

# Filter out Qlik and Hudi metadata fields
drop_columns = []
drop_columns.extend([c for c in df.columns if c.startswith(CHANGES_METADATA_FIELD_PREFIX)])
drop_columns.extend([c for c in df.columns if c.startswith(HUDI_METADATA_FIELD_PREFIX)])
df = df.drop(*drop_columns)

# Export to parquet snapshot
partitions = calculate_partitions(spark=spark, df=df)
print(f">>> Storing snapshot: {cmd_args.snapshot_path}, with {partitions} partition(s)...")

df.repartition(partitions) \
    .write \
    .mode("overwrite") \
    .parquet(cmd_args.snapshot_path)

print(">>> Done!")


#!/usr/bin/env python3
"""
ros2 Image to LaneNet Bridge
"""

import os.path as ops
import sys
import time

from points_vector.msg import PointsVector
from geometry_msgs.msg import Point

import sklearn
import cv2
import numpy as np
import tensorflow as tf

sys.path.append('/home/yvxaiver/lanenet-lane-detection')
from lanenet_model import lanenet
from lanenet_model import lanenet_postprocess
from local_utils.config_utils import parse_config_utils
from local_utils.log_util import init_logger

CFG = parse_config_utils.lanenet_cfg
LOG = init_logger.get_logger(log_file_name_prefix='lanenet_test')

sys.path.append('/home/yvxaiver/LaneNet_to_Trajectory')
from LaneNetToTrajectory import LaneProcessing, DualLanesToTrajectory

class LaneNetImageProcessor():

    def __init__(self, weights_path, image_width, image_height, max_lane_y=420, WARP_RADIUS=20, WP_TO_M_Coeff=[1,1]):
        self.weights_path = weights_path
        self.image_width = image_width
        self.image_height = image_height
        self.lane_processor = LaneProcessing(self.image_width,self.image_height,max_lane_y,WARP_RADIUS,WP_TO_M_Coeff)
        self.calibration = False
        self.full_lane_pts = []
        self.WP_TO_M_Coeff = WP_TO_M_Coeff
        self.following_path = []

    def init_lanenet(self):

        self.input_tensor = tf.placeholder(dtype=tf.float32, shape=[1, 256, 512, 3], name='input_tensor')

        self.net = lanenet.LaneNet(phase='test', cfg=CFG)
        self.binary_seg_ret, self.instance_seg_ret = self.net.inference(input_tensor=self.input_tensor, name='LaneNet')

        self.postprocessor = lanenet_postprocess.LaneNetPostProcessor(cfg=CFG)

        # Set sess configuration
        sess_config = tf.ConfigProto()
        sess_config.gpu_options.per_process_gpu_memory_fraction = CFG.GPU.GPU_MEMORY_FRACTION
        sess_config.gpu_options.allow_growth = CFG.GPU.TF_ALLOW_GROWTH
        sess_config.gpu_options.allocator_type = 'BFC'

        self.sess = tf.Session(config=sess_config)

        # define moving average version of the learned variables for eval
        with tf.variable_scope(name_or_scope='moving_avg'):
            variable_averages = tf.train.ExponentialMovingAverage(
                CFG.SOLVER.MOVING_AVE_DECAY)
            variables_to_restore = variable_averages.variables_to_restore()

        # define saver
        self.saver = tf.train.Saver(variables_to_restore)

        with self.sess.as_default():
            self.saver.restore(sess=self.sess, save_path=self.weights_path)

        return True

    def image_to_trajectory(self, cv_image, lane_fit=True):

        T_start = time.time()

        image_vis = cv_image
        image = cv2.resize(cv_image, (512, 256), interpolation=cv2.INTER_LINEAR)
        image = image / 127.5 - 1.0

        T_resize = time.time()
        LOG.info('Image Resize cost time: {:.5f}s'.format(T_resize-T_start))

        binary_seg_image, instance_seg_image = self.sess.run(
            [self.binary_seg_ret, self.instance_seg_ret],
            feed_dict={self.input_tensor: [image]}
        )

        T_seg_inference = time.time()
        LOG.info('TOTAL Segmentation Inference cost time: {:.5f}s'.format(T_seg_inference-T_resize))

        if not lane_fit:
            out_dict = self.postprocessor.postprocess(
                binary_seg_result=binary_seg_image[0],
                instance_seg_result=instance_seg_image[0],
                source_image=image_vis,
                with_lane_fit=lane_fit,
                data_source='tusimple'
            )
            return out_dict

        full_lane_pts = self.postprocessor.postprocess_lanepts(
            binary_seg_result=binary_seg_image[0],
            instance_seg_result=instance_seg_image[0],
            source_image=image_vis,
            data_source='tusimple'
        )
        self.full_lane_pts = full_lane_pts

        self.lane_processor.process_next_lane(full_lane_pts)
        full_lane_pts = self.lane_processor.get_full_lane_pts()
        physical_fullpts = self.lane_processor.get_physical_fullpts()

        if self.calibration:
            print(self.lane_processor.WARP_RADIUS)
            print(self.lane_processor.get_wp_to_m_coeff())
            self.calibration = False
            raise SystemExit

        phy_centerpts = []
        phy_splines = []
        closest_lane_dist = float('inf')
        closest_lane_idx = 0
        if physical_fullpts:
            #print(physical_fullpts)
            for i in range(len(physical_fullpts)):
                if not i: continue
                traj = DualLanesToTrajectory(physical_fullpts[i-1],physical_fullpts[i],N_centerpts=20)
                phy_centerpts.append(traj.get_centerpoints())
                phy_splines.append(traj.get_spline())
            min_center_y_val = float('inf')
            #for lane in phy_centerpts:
                #if min(lane[1]) < min_center_y_val: min_center_y_val = min(lane[1])
            for i in range(len(phy_splines)):
                new_dist = abs(phy_splines[i](0.2)-(self.image_width*self.WP_TO_M_Coeff[0])/2)
                if new_dist < closest_lane_dist:
                    closest_lane_dist = new_dist
                    closest_lane_idx = i
            if phy_centerpts: self.following_path = phy_centerpts[closest_lane_idx]

        # For display output
        centerpts = []
        if full_lane_pts:
            for i in range(len(full_lane_pts)):
                if not i: continue
                traj = DualLanesToTrajectory(full_lane_pts[i-1],full_lane_pts[i],N_centerpts=20)
                centerpts.append(traj.get_centerpoints())

        T_post_process = time.time()
        LOG.info('Image Post-Process cost time: {:.5f}s'.format(T_post_process-T_seg_inference))

        if centerpts: return full_lane_pts, centerpts, self.following_path
        return None, None, None

    def get_point_vector_path(self):
        #print(self.following_path)
        if self.following_path and self.full_lane_pts:
            vector = []
            for i in range(len(self.following_path[0])):
                pt = Point()
                pt.x = self.following_path[0][i]
                pt.y = self.following_path[1][i]
                pt.z = 0.0
                vector.append(pt)
            ptv = PointsVector()
            ptv.points = vector
            ptv.x_coeff = float(self.lane_processor.get_wp_to_m_coeff()[0])
            return ptv
        return None


import os
import sys
import time
import wandb
import numpy as np
import torch
from torch.utils.data import DataLoader
from dataloader_from_numpy import *
from my_models import *
from torch.utils.tensorboard import SummaryWriter
from UniqueDraw import UniqueDraw
from utils import *
import torchvision
from mobile_net_v3 import *
from SincNet import SincConv_fast
import torchaudio

def hook_fn(module, input, output):
    """ Store the output of the hook """
    global hooked_output
    hooked_output = output

def compute_EER(sim_matrix):
    """
    Compute EER, FAR, FRR and the threshold at which EER occurs.

    Args:
    - sim_matrix (torch.Tensor): A similarity matrix of shape
      (num of speakers, num of utterances, num of speakers).

    Returns:
    - EER (float): Equal error rate.
    - threshold (float): The threshold at which EER occurs.
    - FAR (float): False acceptance rate at EER.
    - FRR (float): False rejection rate at EER.
    """
    num_of_speakers, num_of_utters, _ = sim_matrix.shape

    # Initialize values
    diff = float('inf')
    EER = 0.0
    threshold = 0.0
    EER_FAR = 0.0
    EER_FRR = 0.0

    # Iterate over potential thresholds
    for thres in torch.linspace(0.5, 1.0, 501):
        sim_matrix_thresh = sim_matrix > thres

        # Compute FAR and FRR
        FAR = sum([(sim_matrix_thresh[i].sum() - sim_matrix_thresh[i, :, i].sum()).float()
                    for i in range(num_of_speakers)]) / (num_of_speakers - 1.0) / (num_of_utters) / num_of_speakers

        FRR = sum([(num_of_utters - sim_matrix_thresh[i, :, i].sum()).float()
                   for i in range(num_of_speakers)]) / (num_of_utters) / num_of_speakers

        # Update if this is the closest FAR and FRR we've seen so far
        if diff > abs(FAR - FRR):
            diff = abs(FAR - FRR)
            EER = ((FAR + FRR) / 2).item()
            threshold = thres.item()
            EER_FAR = FAR.item()
            EER_FRR = FRR.item()

    return EER, threshold, EER_FAR, EER_FRR

def train_and_test_model(device, models, loss_func, data_set, optimizer, train_batch_size, test_batch_size,
                         n_fft=512, hop_length=256, win_length=512, window_fn = torch.hann_window, power=None,
                         num_epochs=2000, train_ratio=0.8, contain_sicnet=False, contain_2_extractor=False):
    # writer = SummaryWriter()
    # number of user for train and test in each epoch. Train 4 users, test 2 users.
    data_size = len(data_set)
    train_size = int(data_size * train_ratio)
    test_size = data_size - train_size
    train_tmp_set, test_tmp_set = torch.utils.data.random_split(data_set, [train_size, test_size])
    train_loader = DataLoader(train_tmp_set, batch_size=train_batch_size, shuffle=True, drop_last=True)
    test_loader = DataLoader(test_tmp_set, batch_size=test_batch_size, shuffle=True, drop_last=True)

    if contain_2_extractor:
        sinc_a, sinc_p, extractor_a, extractor_p = models
    else:
        sinc_a, sinc_p, extractor = models

    # functions for STFT matrix calculation
    spectrogram = torchaudio.transforms.Spectrogram(
                                                    n_fft=n_fft,
                                                    win_length=win_length,
                                                    hop_length=hop_length,
                                                    window_fn=window_fn,
                                                    power=power,  # For power spectrogram, use 2. For complex spectrogram, use None.
                                                    # batch_first=True,
                                                    # sample_rate=16000
                                                ).to(device)
    amplitude_to_db = torchaudio.transforms.AmplitudeToDB(stype='magnitude', top_db=80).to(device)

    for epoch in range(num_epochs):
        print(f'Epoch {epoch + 1}/{num_epochs}')
        print('-' * 10)

        # train and test model
        for phase in ['train', 'test']:
            if phase == 'train':
                # set model to training
                sinc_a.train()
                sinc_p.train()
                if contain_2_extractor:
                    extractor_a.train()
                    extractor_p.train()
                else:
                    extractor.train()
                dataloader = train_loader
            else:
                # set model to test
                sinc_a.eval()
                sinc_p.eval()
                if contain_2_extractor:
                    extractor_a.eval()
                    extractor_p.eval()
                else:
                    extractor.eval()
                dataloader = test_loader

            # train each batch
            loss_avg_batch_all = 0.0
            loss_avg_batch_all_piezo = 0.0
            loss_avg_batch_all_audio = 0.0
            acc_audio = 0.0
            acc_piezo = 0.0

            for batch_id, (piezo_clips, audio_clips, ids) in enumerate(dataloader):
                # get shape of input
                batch_size, _ = piezo_clips.shape

                with torch.set_grad_enabled(phase == 'train') and torch.autograd.set_detect_anomaly(True):

                    if phase == 'train':
                        # training
                        # (batch_size, samples) (8000 raw audio signal, 16000 sample rate)
                        # process data using sincnet
                        # (batch_size, samples) -> (batch_size, sinc_out_channel, samples)
                        piezo_clips = piezo_clips.unsqueeze(1).to(device)
                        audio_clips = audio_clips.unsqueeze(1).to(device)
                        if contain_sicnet:
                            piezo_clips = sinc_p(piezo_clips)
                            # audio_clips = sinc_a(audio_clips)
                        _, channel_clip, _ = piezo_clips.shape
                        piezo_clips.contiguous()
                        piezo_clips = piezo_clips.view(batch_size * channel_clip, -1)
                        audio_clips.contiguous()
                        audio_clips = audio_clips.view(batch_size * 1, -1) # changed to 1

                        piezos = amplitude_to_db(spectrogram(piezo_clips).abs())
                        _, f_len, t_len = piezos.shape
                        audios = amplitude_to_db(spectrogram(audio_clips).abs())
                        _, f_len_a, t_len_a = audios.shape

                        # audio-based normalization
                        # piezos, audios = normalize_tensors_based_audio(tensor_p=piezos, tensor_a=audios)
                        #reshape
                        piezos.contiguous()
                        piezos = piezos.view(batch_size, channel_clip, f_len, t_len)
                        audios.contiguous()
                        audios = audios.view(batch_size, 1, f_len_a, t_len_a) ## changed to 1

                        if contain_2_extractor:
                            pred_ids_piezo = extractor_p(piezos)
                            pred_ids_audio = extractor_a(audios)
                        else:
                            pred_ids_piezo = extractor(piezos)
                            pred_ids_audio = extractor(audios)

                        ids_gpu = ids.to(device)
                        loss_class_piezo = loss_func(pred_ids_piezo, ids_gpu)
                        loss_class_audio = loss_func(pred_ids_audio, ids_gpu)
                        acc_audio += np.sum(np.argmax(pred_ids_audio.cpu().data.numpy(), axis=1) == ids_gpu.cpu().data.numpy())
                        acc_piezo += np.sum(np.argmax(pred_ids_piezo.cpu().data.numpy(), axis=1) == ids_gpu.cpu().data.numpy())
                        loss_extractor = loss_class_piezo + loss_class_audio
                        loss_avg_batch_all += loss_extractor.item()
                        loss_avg_batch_all_audio += loss_class_audio.item()
                        loss_avg_batch_all_piezo += loss_class_piezo.item()

                        optimizer.zero_grad()
                        loss_extractor.backward()
                        torch.nn.utils.clip_grad_norm_(sinc_a.parameters(), 10.0)
                        torch.nn.utils.clip_grad_norm_(sinc_p.parameters(), 10.0)
                        if contain_2_extractor:
                            torch.nn.utils.clip_grad_norm_(extractor_p.parameters(), 10.0)
                            torch.nn.utils.clip_grad_norm_(extractor_a.parameters(), 10.0)
                        else:
                            torch.nn.utils.clip_grad_norm_(extractor.parameters(), 10.0)
                        optimizer.step()

                    if phase == 'test':
                        # tesing
                        # (batch_size, samples) (8000 raw audio signal, 16000 sample rate)
                        # process data using sincnet
                        # (batch_size, samples) -> (batch_size, sinc_out_channel, samples)
                        piezo_clips = piezo_clips.unsqueeze(1).to(device)
                        audio_clips = audio_clips.unsqueeze(1).to(device)
                        if contain_sicnet:
                            piezo_clips = sinc_p(piezo_clips)
                            # audio_clips = sinc_a(audio_clips)
                        _, channel_clip, _ = piezo_clips.shape
                        piezo_clips.contiguous()
                        piezo_clips = piezo_clips.view(batch_size * channel_clip, -1)
                        audio_clips.contiguous()
                        audio_clips = audio_clips.view(batch_size * 1, -1)

                        piezos = amplitude_to_db(spectrogram(piezo_clips).abs())
                        _, f_len, t_len = piezos.shape
                        audios = amplitude_to_db(spectrogram(audio_clips).abs())
                        _, f_len_a, t_len_a = audios.shape

                        # audio-based normalization
                        # piezos, audios = normalize_tensors_based_audio(tensor_p=piezos, tensor_a=audios)
                        #reshape
                        piezos.contiguous()
                        piezos = piezos.view(batch_size, channel_clip, f_len, t_len)
                        audios.contiguous()
                        audios = audios.view(batch_size, 1, f_len_a, t_len_a)

                        if contain_2_extractor:
                            pred_ids_piezo = extractor_p(piezos)
                            pred_ids_audio = extractor_a(audios)
                        else:
                            pred_ids_piezo = extractor(piezos)
                            pred_ids_audio = extractor(audios)

                        ids_gpu = ids.to(device)
                        loss_class_piezo = loss_func(pred_ids_piezo, ids_gpu)
                        loss_class_audio = loss_func(pred_ids_audio, ids_gpu)
                        acc_audio += np.sum(np.argmax(pred_ids_audio.cpu().data.numpy(), axis=1) == ids_gpu.cpu().data.numpy())
                        acc_piezo += np.sum(np.argmax(pred_ids_piezo.cpu().data.numpy(), axis=1) == ids_gpu.cpu().data.numpy())
                        loss_extractor = loss_class_piezo + loss_class_audio
                        loss_avg_batch_all += loss_extractor.item()
                        loss_avg_batch_all_audio += loss_class_audio.item()
                        loss_avg_batch_all_piezo += loss_class_piezo.item()

            epoch_loss_all = loss_avg_batch_all / len(dataloader)
            epoch_loss_audio = loss_avg_batch_all_audio / len(dataloader)
            epoch_loss_piezo = loss_avg_batch_all_piezo / len(dataloader)

            if phase == 'train':
                epoch_acc_audio = acc_audio / (len(dataloader) * train_batch_size)
                epoch_acc_piezo = acc_piezo / (len(dataloader) * train_batch_size)
                print(f'{phase} Loss Extractor: {epoch_loss_all:.4f}')
                print(f'{phase} Acc Audio: {epoch_acc_audio:.4f}')
                print(f'{phase} Acc Piezo: {epoch_acc_piezo:.4f}')
                wandb.log({'epoch': epoch, f'Loss/{phase}_all': epoch_loss_all})
                wandb.log({'epoch': epoch, f'Loss/{phase}_audio': epoch_loss_audio})
                wandb.log({'epoch': epoch, f'Loss/{phase}_piezo': epoch_loss_piezo})

                wandb.log({'epoch': epoch, f'Acc/{phase}_audio': epoch_acc_audio})
                wandb.log({'epoch': epoch, f'Acc/{phase}_piezo': epoch_acc_piezo})
            if phase == 'test':
                epoch_acc_audio = acc_audio / (len(dataloader) * test_batch_size)
                epoch_acc_piezo = acc_piezo / (len(dataloader) * test_batch_size)
                print(f'{phase} Loss Extractor: {epoch_loss_all:.4f}')
                print(f'{phase} Acc Audio: {epoch_acc_audio:.4f}')
                print(f'{phase} Acc Piezo: {epoch_acc_piezo:.4f}')
                wandb.log({'epoch': epoch, f'Loss/{phase}_extractor': epoch_loss_all})
                wandb.log({'epoch': epoch, f'Loss/{phase}_audio': epoch_loss_audio})
                wandb.log({'epoch': epoch, f'Loss/{phase}_piezo': epoch_loss_piezo})

                wandb.log({'epoch': epoch, f'Acc/{phase}_audio': epoch_acc_audio})
                wandb.log({'epoch': epoch, f'Acc/{phase}_piezo': epoch_acc_piezo})

    return (sinc_a, sinc_p, extractor)

if __name__ == "__main__":

    device = "cuda:1" if torch.cuda.is_available() else "cpu"

    data_file_dir = './processed_data/wav_clips/piezobuds/' # folder where stores the data for training and test
    pth_store_dir = './pth_model/'
    os.makedirs(pth_store_dir, exist_ok=True)

    # set the params of each train
    # ----------------------------------------------------------------------------------------------------------------
    # Be sure to go through all the params before each run in case the models are saved in wrong folders!
    # ----------------------------------------------------------------------------------------------------------------

    lr = 0.001
    n_user = 51
    train_ratio = 0.8
    num_of_epoches = 300
    train_seperate = False
    sinc_out_channel = 4
    sinc_kernel_size = 512
    contain_sincnet = True
    contain_2_extractor = True
    train_batch_size = 512
    test_batch_size = 128

    n_fft = 512  # Size of FFT, affects the frequency granularity
    hop_length = 256  # Typically n_fft // 4 (is None, then hop_length = n_fft // 2 by default)
    win_length = n_fft  # Typically the same as n_fft
    window_fn = torch.hann_window # Window function

    if contain_sincnet:
        comment = 'sincnet_out_{}_ksize_{}_mobilenetv3small_nfft_{}_hop_{}_classification'.format(sinc_out_channel, sinc_kernel_size, n_fft, hop_length)
    else:
        comment = 'nfft_{}_hop_{}_classification'.format(sinc_out_channel, sinc_kernel_size, n_fft, hop_length)
    if contain_2_extractor:
        comment += '_2_extractor'
    # comment = 'mobilenetv3large1d_960_hop_256_t_16_class_pwr_spec_49u' # simple descriptions of specifications of this model, for example, 't_f' means we use the model which contains time and frequency nn layers

    if contain_sincnet == False:
        sinc_out_channel = 1
    # SincNet initialization
    sinc_p = SincConv_fast(out_channels=sinc_out_channel, kernel_size=sinc_kernel_size).to(device)
    sinc_a = SincConv_fast(out_channels=sinc_out_channel, kernel_size=sinc_kernel_size).to(device)

    # extractor initialization
    if contain_2_extractor == False:
        extractor = torchvision.models.mobilenet_v3_large(pretrained=True)
        extractor.features[0][0] = nn.Conv2d(sinc_out_channel, 16, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
        extractor.classifier[3] = nn.Linear(extractor.classifier[3].in_features, n_user)
        # hook for feature layer output
        hook = extractor.avgpool.register_forward_hook(hook_fn)
        extractor.to(device)
    else:
        extractor_a = torchvision.models.mobilenet_v3_large(pretrained=True)
        extractor_a.features[0][0] = nn.Conv2d(1, 16, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False) ## changed to 1s
        extractor_a.classifier[3] = nn.Linear(extractor_a.classifier[3].in_features, n_user)
        # hook for feature layer output
        hook_a = extractor_a.avgpool.register_forward_hook(hook_fn)
        extractor_a.to(device)

        extractor_p = torchvision.models.mobilenet_v3_large(pretrained=True)
        extractor_p.features[0][0] = nn.Conv2d(sinc_out_channel, 16, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
        extractor_p.classifier[3] = nn.Linear(extractor_p.classifier[3].in_features, n_user)
        # hook for feature layer output
        hook = extractor_p.avgpool.register_forward_hook(hook_fn)
        extractor_p.to(device)

    if contain_sincnet:
        if contain_2_extractor:
            optimizer = torch.optim.Adam([
            {'params': extractor_a.parameters()},
            {'params': extractor_p.parameters()},
            {'params': sinc_a.parameters()},
            {'params': sinc_p.parameters()}
        ], lr=lr)
        else:
            optimizer = torch.optim.Adam([
            {'params': extractor.parameters()},
            {'params': sinc_a.parameters()},
            {'params': sinc_p.parameters()}
        ], lr=lr)
    else:
        if contain_2_extractor:
            optimizer = torch.optim.Adam([
            {'params': extractor_a.parameters()},
            {'params': extractor_p.parameters()},
        ], lr=lr)
        else:
            optimizer = torch.optim.Adam([
            {'params': extractor.parameters()},
        ], lr=lr)

    # create the folder to store the model
    model_struct = 'model_' + comment
    # initialize the wandb configuration
    time_stamp = time.strftime("%Y_%m_%d_%H_%M", time.localtime())
    wandb.init(
        # team name
        entity="piezobuds",
        # set the project name
        project="PiezoBuds",
        # params of the task
        name=model_struct+'_'+time_stamp
    )
    model_store_pth = pth_store_dir + model_struct + '/'
    os.makedirs(model_store_pth, exist_ok=True)
    model_final_path = model_store_pth + time_stamp + '/'
    os.makedirs(model_final_path, exist_ok=True)

    # load the data
    data_set = WavDataset('./processed_data/wav_clips/piezobuds/', n_user=n_user, is_multi_moda=True)
    print(len(data_set))

    loss_func = nn.CrossEntropyLoss()

    if contain_2_extractor:
        models = (sinc_a, sinc_p, extractor_a, extractor_p)
        (sinc_a, sinc_p, extractor_a, extractor_p) = train_and_test_model(device=device, models=models, loss_func=loss_func, data_set=data_set, optimizer=optimizer,
                                                        train_batch_size=train_batch_size, test_batch_size=test_batch_size, n_fft=n_fft,
                                                        hop_length=hop_length, win_length=win_length, window_fn=window_fn, power=None,
                                                        num_epochs=num_of_epoches, train_ratio=train_ratio, contain_sicnet=True, contain_2_extractor=contain_2_extractor)

        torch.save(extractor_a.state_dict(), model_final_path+'extractor_a.pth')
        torch.save(extractor_a.state_dict(), model_final_path+'extractor_p.pth')
        torch.save(sinc_a.state_dict(), model_final_path+'sinc_a.pth')
        torch.save(sinc_p.state_dict(), model_final_path+'sinc_p.pth')
    else:
        models = (sinc_a, sinc_p, extractor)
        (sinc_a, sinc_p, extractor) = train_and_test_model(device=device, models=models, loss_func=loss_func, data_set=data_set, optimizer=optimizer,
                                                        train_batch_size=train_batch_size, test_batch_size=test_batch_size, n_fft=n_fft,
                                                        hop_length=hop_length, win_length=win_length, window_fn=window_fn, power=None,
                                                        num_epochs=num_of_epoches, train_ratio=train_ratio, contain_sicnet=True, contain_2_extractor=contain_2_extractor)

        torch.save(extractor.state_dict(), model_final_path+'extractor.pth')
        torch.save(sinc_a.state_dict(), model_final_path+'sinc_a.pth')
        torch.save(sinc_p.state_dict(), model_final_path+'sinc_p.pth')