import attractiveness
import util
import numpy as np
from keras.models import Sequential
from keras.layers import Dense
from keras.layers import Reshape
import tkinter as tk
from tkinter import filedialog as fd
import cv2

model = Sequential()

# This trainable layer represents an image that we will be modifying
# It always receives exactly 1 as its input, and so produces its weights as its output
# The weights are just the flattened image
image_layer = Dense(util.shape_volume(attractiveness.image_shape), use_bias = False, input_shape = (1,))
model.add(image_layer)
model.add(Reshape(attractiveness.image_shape))

# Now take the image and estimate its attractiveness without changing the estimator
attractiveness.model.trainable = False
model.add(attractiveness.model)
model.compile(optimizer = 'adam', loss = 'mse')

def get_image(dims = None):
    weights = image_layer.get_weights()
    reshaped = np.reshape(weights, attractiveness.image_shape)
    return util.cut_resize_image(reshaped, dims)

def set_image(image):
    cut_resized = util.cut_resize_image(image, attractiveness.image_dims)
    reshaped = np.reshape(cut_resized, (1, 1, util.shape_volume(attractiveness.image_shape)))
    image_layer.set_weights(reshaped)

def get_attractiveness():
    batch_prediction = model.predict_on_batch([1.0])
    return batch_prediction[0][0]

def train_attractiveness(target):
    model.train_on_batch([1.0], [target])

set_image(util.test_faces[0].get_image(attractiveness.image_dims))

def manual_interface():
    preview_name = 'Image'
    slider_length = 256
    target_attractiveness = get_attractiveness()

    control_window = tk.Tk()
    control_window.title('Control')
    current_slider = tk.Scale(control_window, from_ = 0, to = 1, resolution = 0.01,
    orient = tk.HORIZONTAL, length = slider_length, label = 'Current attractiveness')
    current_slider.pack()
    target_slider = tk.Scale(control_window, from_ = 0, to = 1, resolution = 0.01,
    orient = tk.HORIZONTAL, length = slider_length, label = 'Target attractiveness')
    target_slider.set(target_attractiveness)
    target_slider.pack()

    closed = False
    while not closed:
        try:
            current_slider.set(get_attractiveness())
            target_attractiveness = target_slider.get()
            train_attractiveness(target_attractiveness)
            control_window.update()
        except:
            closed = True
        preview(window_name = preview_name, blocking = False)
        if cv2.getWindowProperty(preview_name, 0) < 0:
            closed = True

    try:
        control_window.destroy()
    except:
        pass
    try:
        cv2.destroyWindow(preview_name)
    except:
        pass

def preview(window_name = 'Image', dims = (512, 512), blocking = True):
    util.preview([window_name], [get_image()], dims = dims,
    destroy = blocking, milliseconds = 0 if blocking else 1)


import numpy as np
import cv2 as cv
from matplotlib import pyplot as plt
import argparse

ap = argparse.ArgumentParser()
ap.add_argument("-i", "--image", type=str, required=True,
	help="path input image that we'll detect blur in")
args = vars(ap.parse_args())

img = cv.imread(args['image'],0)
dft = cv.dft(np.float32(img),flags = cv.DFT_COMPLEX_OUTPUT)
dft_shift = np.fft.fftshift(dft)
magnitude_spectrum = 20*np.log(cv.magnitude(dft_shift[:,:,0],dft_shift[:,:,1]))
plt.subplot(121),plt.imshow(img, cmap = 'gray')
plt.title('Input Image'), plt.xticks([]), plt.yticks([])
plt.subplot(122),plt.imshow(magnitude_spectrum, cmap = 'gray')
plt.title('Magnitude Spectrum'), plt.xticks([]), plt.yticks([])
plt.show()

rows, cols = img.shape
crow,ccol = rows//2 , cols//2
# create a mask first, center square is 1, remaining all zeros
mask = np.zeros((rows,cols,2),np.uint8)
mask[crow-30:crow+30, ccol-30:ccol+30] = 1
# apply mask and inverse DFT
fshift = dft_shift*mask
f_ishift = np.fft.ifftshift(fshift)
img_back = cv.idft(f_ishift)
img_back = cv.magnitude(img_back[:,:,0],img_back[:,:,1])
plt.subplot(121),plt.imshow(img, cmap = 'gray')
plt.title('Input Image'), plt.xticks([]), plt.yticks([])
plt.subplot(122),plt.imshow(img_back, cmap = 'gray')
plt.title('Magnitude Spectrum'), plt.xticks([]), plt.yticks([])
plt.show()


#!/usr/bin/env python3
"""
Read class averages mrc file and save it to jpg.
Automatically remove the edges.
INPUT: mrcs file of 2D class averages
OUTPUT: a dir for the jpg output
The name of the jpg file would be "particlename_diamxxkxx_classnumber.jpg"
"""

import os
import mrcfile
import numpy as np
from PIL import Image
import argparse
import shutil

def setupParserOptions():
    ap = argparse.ArgumentParser()
    ap.add_argument('-i', '--input',
                    help="Input mrcs file of 2D class averages.")
    ap.add_argument('-n', '--name', default='particle',
                    help="Name of the particle")
    ap.add_argument('-o', '--output', default='2DAssess',
                    help="Output jpg dir.")
    args = vars(ap.parse_args())
    return args

def cutbyradius(img):
    h = img.shape[0]
    w = img.shape[1]
    # empty_val = img[0,0] # because the image is already masked (2d class avg), the [0,0] point must be empty
    edge_l = 0
    for i in range(w):
        if np.sum(img[i,:]) > 1e-7 or np.sum(img[:,i]) < -1e-7:
            edge_l = i
            break
    edge_r = 0
    for ii in range(w):
        if np.sum(img[-ii,:]) > 1e-7 or np.sum(img[:,-ii]) < -1e-7:
            edge_r = ii
            break
    edge_t = 0
    for j in range(h):
        if np.sum(img[:,j]) > 1e-7 or np.sum(img[:,j]) < -1e-7:
            edge_t = j
            break
    edge_b = 0
    for jj in range(h):
        if np.sum(img[:,-jj]) > 1e-7 or np.sum(img[:,-jj]) < -1e-7:
            edge_b = jj
            break
    edge = min(edge_l, edge_r, edge_t, edge_b)
    new_img = img[edge:h-edge+1, edge:w-edge+1]
    return new_img

def save_mrcs(wkdir, **args):
    print('Converting mrcs to jpg....')
    os.chdir(wkdir) # navigate to the par dir of input file
    # try:
    #     shutil.rmtree(args['output'])
    # except OSError:
    #     pass
    # os.makedirs(args['output'], exist_ok=True)
    os.mkdir(os.path.join(args['output'], 'data'))

    # avg_mrc = mrcfile.open(os.path.basename(args['input'])).data
    avg_mrc = mrcfile.open(args['input']).data
    if len(avg_mrc.shape) == 3:
        num_part = avg_mrc.shape[0]
    elif len(avg_mrc.shape) == 2:
        num_part = 1

    for i in range(num_part):
        new_img = avg_mrc[i,:,:]
        if np.sum(new_img) > 1e-7 or np.sum(new_img) < -1e-7:
            new_img = cutbyradius(new_img)
            new_img = ((new_img-new_img.min())/((new_img.max()-new_img.min())+1e-7)*255).astype('uint8')
            new_img = Image.fromarray(new_img)
            new_img = new_img.convert("L")
            new_img.save(os.path.join(args['output'], 'data', (args['name'] + '_' + str(i+1) + '.jpg')))

if __name__ == '__main__':
    args = setupParserOptions()
    save_mrcs(wkdir, **args)


import urllib.request as urlreq
from typing import List, Dict

from bs4 import BeautifulSoup
import os
import zipfile
import optparse
import asyncio
import shutil

redditurl: str = 'https://old.reddit.com/r/%s'      # the url for reddit with %s to insert the subreddit name
dl_dir: str = './.cache/'  # Format must be ./      # the directory where files are cached. Will be created if it doesn't exist
img_ext: List[str] = ['jpg', 'png', 'bmp']          # file extensions that are images
blacklist: List[str] = ['b.thumbs.redditmedia.com', 'reddit.com']   # where images shouldn't be downloaded from
hdr: Dict[str, str] = {                             # request header
    'User-Agent': """Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.11 (KHTML, like Gecko)
                     Chrome/23.0.1271.64 Safari/537.11""",
    'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8',
    'Accept-Charset': 'ISO-8859-1,utf-8;q=0.7,*;q=0.3', 'Accept-Encoding': 'none', 'Accept-Language': 'en-US,en;q=0.8',
    'Connection': 'keep-alive'}

# prints a progress bar
def print_progress(iteration, total, prefix='', suffix='', decimals=1, length=100, fill='█'):
    percent = ("{0:." + str(decimals) + "f}").format(100 * (iteration / float(total)))
    filled_length = int(length * iteration // total)
    bar = fill * filled_length + '-' * (length - filled_length)
    print('\r%s |%s| %s%% %s' % (prefix, bar, percent, suffix), end='\r')
    # Print New Line on Complete
    if iteration == total:
        print()

# returns a soup for the given url
async def request_soup(url):
    req = urlreq.Request(url, headers=hdr)
    html = None
    for x in range(0, 10):
        try:
            html = urlreq.urlopen(req).read()
            break
        except Exception as e:
            print('[-]', e)
            await asyncio.sleep(1)
    soup = BeautifulSoup(html, "lxml")
    return soup

# returns all images for the given url
async def get_img_as(url):
    soup = await  request_soup(url)
    ret = []
    for t in soup.find_all(has_source):
        if 'redditmedia' not in t['src']:
            try:
                ret.append(t['src'])
            except KeyError:
                pass
    return ret

# returns the last post id in the given reddit page
async def get_next(url):
    ids = []
    soup = await request_soup(url)
    for t in soup.find_all(has_source):
        if 'redditmedia' not in t['src']:
            try:
                fname = t['data-fullname']
                ids.append(fname)
            except KeyError:
                pass
    return [_id for _id in ids if _id][-1]

# returns if the given tag has a source attribute that is an image
def has_source(tag):
    if tag.has_attr('src'):
        try:
            return tag['src'].split('.')[-1].lower() in img_ext
        except IndexError or KeyError:
            return False
    elif tag.has_attr('data-url'):
        try:
            tag['src'] = tag['data-url']
            return tag['src'].split('.')[-1].lower() in img_ext
        except KeyError or KeyError:
            return False
    else:
        return False

# downloads all images for the given url and puts them in a zipfile
async def download_async(url, zfile=None, test=False):
    images = await get_img_as(url)
    print('[+] Found %s images' % len(images))
    logmsg = ""
    imgcount = len(images)
    savedcount = 0
    count = 0
    print_progress(count, imgcount, prefix="Downloading: ", suffix="Complete")
    for img in images:
        print_progress(count+1, imgcount, prefix="Downloading: ", suffix="Complete")
        count += 1
        if test:
            continue
        try:
            if 'http' not in img.split('/')[0] and '//' not in img.split('.')[0]:
                img = url + img
            if 'http' not in img.split('/')[0]:
                img = 'http:' + img
            if img.strip('http://').strip('https://').split('/')[0] in blacklist:
                img = None
                continue
            imgname = img.split('/')[-1]
            name = dl_dir + imgname
            if os.path.isfile(name):
                continue
            f = open(name, "wb")
            req = urlreq.Request(img, headers=hdr)
            image = urlreq.urlopen(req)
            f.write(image.read())
            f.close()
            zfile.write(name, imgname, zipfile.ZIP_DEFLATED)
            try:
                os.remove(name)
            except FileNotFoundError or PermissionError:
                pass
            savedcount += 1
            await asyncio.sleep(0.25)
        except Exception as error:
            logmsg += '[-] Failed with %s %s\n' % (img, error)
    print('[+] %s images downloaded | %s finished %s' % (savedcount, logmsg, url))

# loops over reddit-pages until no more images are found
async def dl_loop(section, zfile, loop, chaos=False, test=False):
    baseurl = redditurl % section
    url = baseurl
    if chaos:
        loop.create_task(download_async(url, zfile, test))
    else:
        await loop.create_task(download_async(url, zfile, test))
    while True:
        print('[*] Getting Images from %s' % url)
        try:
            after = await get_next(url)
            url = '{}/?after={}'.format(baseurl, after)
            if chaos:
                loop.create_task(download_async(url, zfile, test))
            else:
                await loop.create_task(download_async(url, zfile, test))
        except Exception as ex:
            print('[-]', ex)
            zfile.close()
            break
        finally:
            await asyncio.sleep(0.1)

# the main function
def main(sections, opts):
    chaos = opts.chaos
    if not os.path.exists(dl_dir):
        os.makedirs(dl_dir)
    zfiles = {}
    for sect in sections:
        mode = 'w'
        if os.path.isfile(sect + '.zip'):
            mode = 'a'
        zfiles[sect] = zipfile.ZipFile('%s.zip' % sect, mode)
    loop = asyncio.get_event_loop()
    try:
        for sect in sections:
            if chaos:
                loop.create_task(loop.create_task(
                    dl_loop(sect, zfiles[sect], loop, chaos=True, test=opts.test)))
            else:
                loop.run_until_complete(loop.create_task(
                    dl_loop(sect, zfiles[sect], loop, test=opts.test)))
        if chaos:
            loop.run_forever()
    except KeyboardInterrupt:
        for sect in sections:
            try:
                zfiles[sect].close()
            except Exception as error:
                print(error)
    finally:
        shutil.rmtree(dl_dir)

if __name__ == '__main__':
    parser = optparse.OptionParser(usage="usage: %prog [options] [subreddits]")
    parser.add_option('-c', '--chaos', dest='chaos',
                      action='store_true', default=False,
                      help=""" Doesn't wait for previous downloads to finish and doesn't exit when no more
                      images can be found. Do only activate this if you want to download a lot of images
                      from multiple subreddits at the same time. Only option to exit is CTRL + C.""")
    parser.add_option('-t', '--test', dest='test',
                      action='store_true', default=False,
                      help='Tests the functions of the script')
    options, sects = parser.parse_args()
    print('[~] Recieved subreddits %s' % ', '.join(sects))
    main(sects, opts=options)


import sys
import pyqtgraph as pg
from pyqtgraph.Qt import QtCore, QtWidgets
import numpy as np
import cv2
import math
import os

VERSION = (0, 2, 0)

class ControlWindow(QtWidgets.QWidget):
    def __init__(self):
        QtWidgets.QWidget.__init__(self)
        self.cap = None
        self.filename=None
        self.line1_x=300
        self.line2_x=600
        self.magnification=1

        self.resize(400, 800)

        self.videoSize = QtWidgets.QComboBox(self)
        self.videoSize.addItem("1X")
        self.videoSize.addItem("1.5X")
        self.videoSize.addItem("2X")
        self.videoSize.addItem("3X")
        self.videoSize.addItem("4X")
        self.videoSize.setCurrentIndex(0)
        self.videoSize.setEnabled(True)
        self.videoSize.activated.connect(self.magnific)
        self.videoSize.setGeometry(QtCore.QRect(675, 30, 50, 20))

        self.LablePath = QtWidgets.QLineEdit(self)
        self.LablePath.setEnabled(False)
        self.LablePath.setGeometry(QtCore.QRect(90, 10, 641, 20))
        self.LablePath.setFrame(False)

        self.LableCamera = QtWidgets.QLineEdit(self)
        self.LableCamera.setEnabled(False)
        self.LableCamera.setGeometry(QtCore.QRect(675, 10, 50, 20))
        self.LableCamera.setFrame(False)
        self.LableCamera.setText("Camera:")

        self.Camera = QtWidgets.QComboBox(self)
        self.Camera.addItem("Left")
        self.Camera.addItem("Right")
        self.Camera.setEnabled(True)
        self.Camera.setGeometry(QtCore.QRect(675, 30, 50, 20))

        self.LableRatL = QtWidgets.QLineEdit(self)
        self.LableRatL.setEnabled(False)
        self.LableRatL.setGeometry(QtCore.QRect(735, 10, 50, 20))
        self.LableRatL.setFrame(False)
        self.LableRatL.setText("Rat L:")

        self.LableRat1 = QtWidgets.QLineEdit(self)
        self.LableRat1.setEnabled(True)
        self.LableRat1.setText("0")
        self.LableRat1.setGeometry(QtCore.QRect(735, 30, 40, 20))
        self.LableRat1.setFrame(True)

        self.LableRatM = QtWidgets.QLineEdit(self)
        self.LableRatM.setEnabled(False)
        self.LableRatM.setGeometry(QtCore.QRect(785, 10, 50, 20))
        self.LableRatM.setFrame(False)
        self.LableRatM.setText("Rat M:")

        self.LableRat2 = QtWidgets.QLineEdit(self)
        self.LableRat2.setEnabled(True)
        self.LableRat2.setText("0")
        self.LableRat2.setGeometry(QtCore.QRect(785, 30, 40, 20))
        self.LableRat2.setFrame(True)

        self.LableRatR = QtWidgets.QLineEdit(self)
        self.LableRatR.setEnabled(False)
        self.LableRatR.setGeometry(QtCore.QRect(835, 10, 50, 20))
        self.LableRatR.setFrame(False)
        self.LableRatR.setText("Rat R:")

        self.LableRat3 = QtWidgets.QLineEdit(self)
        self.LableRat3.setEnabled(True)
        self.LableRat3.setText("0")
        self.LableRat3.setGeometry(QtCore.QRect(835, 30, 40, 20))
        self.LableRat3.setFrame(True)

        self.Table = QtWidgets.QTableWidget(self)
        self.Table.setGeometry(QtCore.QRect(675, 60,235, 560))
        self.Table.setColumnCount(5)
        self.Table.resizeRowsToContents()
#        self.Table.keyPressEvent=self.keyPressEvent
        self.Table.itemClicked.connect(self.tableClick)
#        self.Table.mousePressEvent=self.tableClick
        self.Table.setHorizontalHeaderLabels(["Frame","Time(ms)", "Rat "+self.LableRat1.text(),"Rat "+self.LableRat2.text(),"Rat "+self.LableRat3.text()])
        self.Table.resizeColumnsToContents()
        self.Table.verticalHeader().setVisible(False)

        self.TableSumary = QtWidgets.QTableWidget(self)
        self.TableSumary.setGeometry(QtCore.QRect(300, 530,340,110))
        self.TableSumary.setEnabled(False)
        self.TableSumary.setColumnCount(8)
        self.TableSumary.setRowCount(3)
        self.TableSumary.resizeRowsToContents()
        self.TableSumary.setVerticalHeaderLabels(["Rat "+self.LableRat1.text(),"Rat "+self.LableRat2.text(),"Rat "+self.LableRat3.text()])
        self.TableSumary.setHorizontalHeaderLabels(["X0", "T", "K", "D", "S", "I", "Attempts", "Trials"])
        self.TableSumary.resizeColumnsToContents()
        self.TableSumary.resizeRowsToContents()

        for i in range (0, 3):
            for j in range (0,8):
                newitem = QtWidgets.QTableWidgetItem("0")
                self.TableSumary.setItem(i, j, newitem)

        self.open_button = QtWidgets.QPushButton(self)
        self.open_button.setText("Open Video")
        self.open_button.clicked.connect(self.getfile)
        self.open_button.setGeometry(QtCore.QRect(9, 9, 75, 23))

        self.save_button = QtWidgets.QPushButton(self)
        self.save_button.setText("Save data")
        self.save_button.clicked.connect(self.handleSave)
        self.save_button.setGeometry(QtCore.QRect(575, 9, 75, 23))

        self.play_button = QtWidgets.QPushButton(self)
        self.play_button.setText("Play")
        self.play_button.clicked.connect(self.playBtn)
        self.play_button.setGeometry(QtCore.QRect(9, 530, 75, 23))

        self.nplay_button = QtWidgets.QRadioButton(self)
        self.nplay_button.setText("Normal")
        self.nplay_button.setGeometry(QtCore.QRect(90, 530, 75, 23))

        self.fplay_button = QtWidgets.QRadioButton(self)
        self.fplay_button.setText("Fast")
        self.fplay_button.setChecked(True)
        self.fplay_button.setGeometry(QtCore.QRect(150, 530, 75, 23))

        self.LableRat1.textChanged.connect(self.tableHeader)
        self.LableRat2.textChanged.connect(self.tableHeader)
        self.LableRat3.textChanged.connect(self.tableHeader)

        layout = QtWidgets.QGridLayout()
        self.setLayout(layout)

        layout.addWidget(self.open_button,0,0)
        layout.addWidget(self.videoSize,0,5)
        layout.addWidget(self.LablePath,0,1)
        layout.addWidget(self.save_button,10,5)
        layout.addWidget(self.play_button,10,0)
        layout.addWidget(self.nplay_button,10,1)
        layout.addWidget(self.fplay_button,10,2)
        layout.addWidget(self.LableCamera, 10,3)
        layout.addWidget(self.Camera,10,4)
        layout.addWidget(self.LableRatL,1,0)
        layout.addWidget(self.LableRat1,1,1)
        layout.addWidget(self.LableRatM,1,2)
        layout.addWidget(self.LableRat2,1,3)
        layout.addWidget(self.LableRatR,1,4)
        layout.addWidget(self.LableRat3,1,5)
        layout.addWidget(self.TableSumary,2,0,2,6)
        layout.addWidget(self.Table,4,0,6,6)

        self.setWindowTitle('Control Panel')
        self.show()

    def startCapture(self):
        if not self.cap:
            if self.filename:
                self.cap = cv2.VideoCapture(str(self.filename))

                if not self.cap.isOpened():
                    print("could not open: %s" % self.filename)
#                    self.getfile()
                    return

                self.length = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
                self.fps=self.cap.get(cv2.CAP_PROP_FPS)
                self.fps2=self.fps
#                self.LableFrames.setText("0/"+str(self.length))
                self.Table.setRowCount(self.length)

                self.videoFrame=videoWindow(window=self)

                for i in range(0, self.length):
                    newitem = QtWidgets.QTableWidgetItem(str(i))
                    self.Table.setItem(i, 0, newitem)

                    newitem = QtWidgets.QTableWidgetItem(str(self.fps*i))
                    self.Table.setItem(i, 1, newitem)

                    for j in range(2, 5):
                       newitem = QtWidgets.QTableWidgetItem("-")
                       self.Table.setItem(i, j, newitem)
                       0
                for i in range(0, 3):
                    for j in range(0,8):
                        newitem = QtWidgets.QTableWidgetItem("0")
                        self.TableSumary.setItem(i, j, newitem)

#        ret, frame = self.cap.read()
#        self.videoFrame.video_frame.setGeometry(QtCore.QRect(10, 40, frame.shape[1], frame.shape[0]))
#        self.video_screen.setGeometry(QtCore.QRect(10, 40, frame.shape[1], frame.shape[0]))
#        self.LableFrames.setGeometry(QtCore.QRect(self.video_screen.x()+self.video_screen.width()-100,self.video_screen.y()+self.video_screen.height()-30 , 81, 20))
#        self.LableFrames.setVisible(True)

        self.timer = QtCore.QTimer()
        self.timer.timeout.connect(self.nextFrameSlot)
        self.nextFrameSlot()

    def endCapture(self):
        #self.stop()
        self.cap.release()
        self.cap = None

    def getfile(self):
        if self.cap:
            self.endCapture()

        filename = QtWidgets.QFileDialog.getOpenFileName(self, 'Open file',
                        'c:\\',"Video files (*.avi *.mp4 *.MOD)")[0]
        if os.path.exists(filename):
            self.filename = os.path.abspath(filename)
        else:
            return

        self.startCapture()
        self.LablePath.setText(self.filename)

    def refresh(self):
        self.nextFrameSlot()
        if (self.play_button.text()=="Stop"):
            self.start()

    def magnific(self):
        if self.videoSize.currentText()=='1X':self.magnification=1
        if self.videoSize.currentText()=='1.5X':self.magnification=1.5
        if self.videoSize.currentText()=='2X':self.magnification=2
        if self.videoSize.currentText()=='3X':self.magnification=3
        if self.videoSize.currentText()=='4X':self.magnification=4
        self.nextFrameSlot()
        self.backwards(1)

    def nextFrameSlot(self):
        nframe = self.cap.get(cv2.CAP_PROP_POS_FRAMES)

        if nframe >= self.length - 1:
            self.stop()
            return

        ret, frame = self.cap.read()
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        frame = cv2.resize(frame,None,fx=self.magnification, fy=self.magnification, interpolation = cv2.INTER_NEAREST)
        img = QtWidgets.QImage(frame, frame.shape[1], frame.shape[0], QtWidgets.QImage.Format_RGB888)
        pix = QtWidgets.QPixmap.fromImage(img)
        self.videoFrame.video_frame.setPixmap(pix)
        self.videoFrame.video_screen.resize(frame.shape[1], frame.shape[0])
        self.videoFrame.video_screen.setGeometry(QtCore.QRect(10, 40, frame.shape[1], frame.shape[0]))
        self.Table.selectRow(nframe)

    def backwards(self, countBack):
        nframe = self.Table.currentRow() - countBack
        self.cap.set(cv2.CAP_PROP_POS_FRAMES,nframe)
        self.nextFrameSlot()

    def playBtn(self):
        if self.fplay_button.isChecked():
            self.fps2=300
        else:
            self.fps2=self.fps

        if (self.play_button.text()=="Play"):
            self.start()
        else:
            self.stop()

    def stop(self):
        self.play_button.setText("Play")
        self.timer.stop()

    def start(self):
        self.play_button.setText("Stop")
        self.timer.start(1000. / self.fps2)

    def openright(self):
        popMenu = QtWidgets.QMenu()

        if (self.play_button.text()=="Stop"):
            self.gtn()

        popMenu.addAction(QtWidgets.QAction("X0", self,  enabled=True, triggered= self.X0))
        popMenu.addAction(QtWidgets.QAction("T", self,  enabled=True, triggered= self.T))
        popMenu.addAction(QtWidgets.QAction("K", self,  enabled=True,triggered= self.K))
        popMenu.addAction(QtWidgets.QAction("D", self,  enabled=True,triggered= self.D))
        popMenu.addAction(QtWidgets.QAction("S", self,  enabled=True,triggered= self.S))
        popMenu.addAction(QtWidgets.QAction("I", self,  enabled=True,triggered= self.I))
        popMenu.addAction(QtWidgets.QAction("----------", self,  enabled=False))
        popMenu.addAction(QtWidgets.QAction("Set Line 1", self,  enabled=True,triggered= self.setLine1))
        popMenu.addAction(QtWidgets.QAction("Set Line 2", self,  enabled=True,triggered= self.setLine2))
        popMenu.addAction(QtWidgets.QAction("----------", self,  enabled=False))
        popMenu.addAction(QtWidgets.QAction("Erese", self,  enabled=True, triggered= self.Erese))
        popMenu.exec_(QtWidgets.QCursor.pos())

    def keyPressEvent(self, event):
         key = event.key()
         if key == QtCore.Qt.Key_Up:
             nframe=self.Table.currentRow()+50
             self.cap.set(cv2.CAP_PROP_POS_FRAMES,nframe)
             self.nextFrameSlot()
         if key == QtCore.Qt.Key_Right:
             self.nextFrameSlot()
         if key == QtCore.Qt.Key_Down:
             self.backwards(52)
         if key == QtCore.Qt.Key_Left:
             self.backwards(2)
         if key == QtCore.Qt.Key_Shift:
             self.playBtn()

    def X0(self): self.menuAction("X0")
    def T(self): self.menuAction("T")
    def K(self): self.menuAction("K")
    def D(self): self.menuAction("D")
    def S(self): self.menuAction("S")
    def I(self): self.menuAction("I")
    def Erese(self): self.menuAction("-")

    def setLine1(self, event):
        self.line1_x=self.x
        self.videoFrame.video_screen.update()
#        self.nextFrameSlot()
#        self.backwards(1)

    def setLine2(self, event):
        self.line2_x=self.x
        self.videoFrame.video_screen.update()
#        self.nextFrameSlot()
#        self.backwards(1)

    def getPos(self , event):
        self.x = event.pos().x()
        self.y = event.pos().y()

    def tableClick(self):
        self.cap.set(1,self.Table.currentRow())
        self.refresh()

    def tableHeader(self):
        self.Table.setHorizontalHeaderLabels(["Frame","Time (ms)", "Rat "+self.LableRat1.text(),"Rat "+self.LableRat2.text(),"Rat "+self.LableRat3.text()])
        self.TableSumary.setVerticalHeaderLabels(["Rat "+self.LableRat1.text(),"Rat "+self.LableRat2.text(),"Rat "+self.LableRat3.text()])

    def menuAction(self, value):
        currentFrame=self.cap.get(cv2.CAP_PROP_POS_FRAMES)

        if self.x<=self.line1_x: j=2
        if self.x<=self.line2_x and self.x>=self.line1_x: j=3
        if self.x>=self.line2_x: j=4

        olditem=self.Table.item(currentFrame, j).text()
        newitem = QtWidgets.QTableWidgetItem(value)
        self.Table.setItem(currentFrame, j, newitem)
        self.Table.selectRow(currentFrame)

        if not olditem==value:
            self.Summary(value, j, 1)

        if not olditem=="-":
            if not olditem==value:
                self.Summary(olditem, j,-1)

        self.playBtn()

    def Summary(self, value, Rat, operator):
        Rat=Rat-2
        suma0=int(self.TableSumary.item(Rat,0).text())
        suma1=int(self.TableSumary.item(Rat,1).text())
        suma2=int(self.TableSumary.item(Rat,2).text())
        suma3=int(self.TableSumary.item(Rat,3).text())
        suma4=int(self.TableSumary.item(Rat,4).text())
        suma5=int(self.TableSumary.item(Rat,5).text())
        suma6=int(self.TableSumary.item(Rat,6).text())
        suma7=int(self.TableSumary.item(Rat,7).text())

        if value=="X0":
            suma0=int(self.TableSumary.item(Rat,0).text())+operator
            newitem = QtWidgets.QTableWidgetItem(str(suma0))
            self.TableSumary.setItem(Rat,0, newitem)
        if value=="T":
            suma1=int(self.TableSumary.item(Rat,1).text())+operator
            newitem = QtWidgets.QTableWidgetItem(str(suma1))
            self.TableSumary.setItem(Rat,1, newitem)
        if value=="K":
            suma2=int(self.TableSumary.item(Rat,2).text())+operator
            newitem = QtWidgets.QTableWidgetItem(str(suma2))
            self.TableSumary.setItem(Rat,2, newitem)
        if value=="D":
            suma3=int(self.TableSumary.item(Rat,3).text())+operator
            newitem = QtWidgets.QTableWidgetItem(str(suma3))
            self.TableSumary.setItem(Rat,3, newitem)
        if value=="S":
            suma4=int(self.TableSumary.item(Rat,4).text())+operator
            newitem = QtWidgets.QTableWidgetItem(str(suma4))
            self.TableSumary.setItem(Rat,4, newitem)
        if value=="I":
            suma5=int(self.TableSumary.item(Rat,5).text())+operator
            newitem = QtWidgets.QTableWidgetItem(str(suma5))
            self.TableSumary.setItem(Rat,5, newitem)

        suma6=suma0+suma1+suma2+suma3+suma4+suma5
        newitem = QtWidgets.QTableWidgetItem(str(suma6))
        self.TableSumary.setItem(Rat,6, newitem)

        suma7=suma2+suma3+suma4+suma5
        newitem = QtWidgets.QTableWidgetItem(str(suma7))
        self.TableSumary.setItem(Rat,7, newitem)

    def handleSave(self):
        s = self.filename
        end = s.find(self.filename)-4
        s[0:end]
        filename=s[0:end]
        with open(filename+".txt", "w") as text_file:
            text_file.write("Frame")
            text_file.write("\t")
            text_file.write("Time (ms)")
            text_file.write("\t")
            text_file.write("Rat "+self.LableRat1.text())
            text_file.write("\t")
            text_file.write("Rat "+self.LableRat2.text())
            text_file.write("\t")
            text_file.write("Rat "+self.LableRat3.text())
            text_file.write("\n")
            for row in range(self.Table.rowCount()):
                for column in range(self.Table.columnCount()):
                    item = self.Table.item(row, column)
                    text_file.write(item.text())
                    text_file.write("\t")
                text_file.write("\n")

        filename2=filename+"_S"
        Rat=[self.LableRat1.text(),self.LableRat2.text(),self.LableRat3.text()]
        with open(filename2+".txt", "w") as text_file:
            text_file.write("Animal")
            text_file.write("\t")
            text_file.write("Camera")
            text_file.write("\t")
            text_file.write("X0")
            text_file.write("\t")
            text_file.write("T")
            text_file.write("\t")
            text_file.write("K")
            text_file.write("\t")
            text_file.write("D")
            text_file.write("\t")
            text_file.write("S")
            text_file.write("\t")
            text_file.write("I")
            text_file.write("\t")
            text_file.write("Attempts")
            text_file.write("\t")
            text_file.write("Trials")
            text_file.write("\n")

            for row in range(self.TableSumary.rowCount()):
                text_file.write(Rat[row])
                text_file.write("\t")
                text_file.write(self.Camera.currentText())
                text_file.write("\t")
                for column in range(self.TableSumary.columnCount()):
                    item = self.TableSumary.item(row, column)
                    text_file.write(item.text())
                    text_file.write("\t")
                text_file.write("\n")

class videoWindow(QtWidgets.QWidget):
        def __init__(self, window):
            self.window = window
            QtWidgets.QWidget.__init__(self)

            self.resize(500, 500)
            layout = QtWidgets.QGridLayout()
            self.setLayout(layout)
            self.keyPressEvent= self.window.keyPressEvent

            self.video_frame=QtWidgets.QLabel()
            self.video_frame.setGeometry(QtCore.QRect(10, 40, 640, 480))
            self.video_frame.setFrameShape(QtWidgets.QFrame.Box)

            self.video_screen=QtWidgets.QLabel()
            self.video_screen.setGeometry(QtCore.QRect(10, 40, 640, 480))
            self.video_screen.paintEvent= self.PaintVerticalLine
            self.video_screen.mousePressEvent= self.window.getPos
            self.video_screen.setContextMenuPolicy(QtCore.Qt.CustomContextMenu)
            self.video_screen.customContextMenuRequested.connect(self.openright)

            self.sl_frame = QtWidgets.QSlider()
            self.sl_frame.setOrientation(QtCore.Qt.Horizontal)
            self.sl_frame.setMinimum(1)
            self.sl_frame.setMaximum(255)
            self.sl_frame.setValue(1)
            self.sl_frame.setTickPosition(QtWidgets.QSlider.TicksBelow)
            self.sl_frame.setTickInterval(1000)
            self.sl_frame.setGeometry(QtCore.QRect(9, 530, 640, 20))
            self.sl_frame.valueChanged.connect(self.window.nextFrameSlot)
            self.sl_frame.setEnabled(False)

            layout.addWidget(self.video_frame, 0,0)
            layout.addWidget(self.video_screen, 0,0)
            layout.addWidget(self.sl_frame, 1,0)
            self.show()

        def PaintVerticalLine(self, event):
            painter = QtWidgets.QPainter(self.video_screen)
            pen = QtWidgets.QPen()
            pen.setColor(QtWidgets.QColor(0,255,0))
            pen.setWidth(2)
            painter.setPen(pen)
            painter.drawLine(self.window.line1_x, 0, self.window.line1_x, self.video_screen.height())
            painter.drawLine(self.window.line2_x, 0, self.window.line2_x, self.video_screen.height())
            painter.drawText(self.window.line1_x/2-5, 15, "Rat L")
            painter.drawText((self.window.line2_x + self.window.line1_x)/2-5, 15, "Rat M")
            painter.drawText((self.video_screen.width() + self.window.line2_x)/2-5, 15, "Rat R")

        def openright(self):
            popMenu = QtWidgets.QMenu()

            if self.window.play_button.text() == "Stop":
                self.window.playBtn()
            popMenu.addAction(QtWidgets.QAction("X0", self,  enabled=True, triggered= self.window.X0))
            popMenu.addAction(QtWidgets.QAction("T", self,  enabled=True, triggered= self.window.T))
            popMenu.addAction(QtWidgets.QAction("K", self,  enabled=True,triggered= self.window.K))
            popMenu.addAction(QtWidgets.QAction("D", self,  enabled=True,triggered= self.window.D))
            popMenu.addAction(QtWidgets.QAction("S", self,  enabled=True,triggered= self.window.S))
            popMenu.addAction(QtWidgets.QAction("I", self,  enabled=True,triggered= self.window.I))
            popMenu.addAction(QtWidgets.QAction("----------", self,  enabled=True))
            popMenu.addAction(QtWidgets.QAction("Set Line 1", self,  enabled=True,triggered= self.window.setLine1))
            popMenu.addAction(QtWidgets.QAction("Set Line 2", self,  enabled=True,triggered= self.window.setLine2))
            popMenu.addAction(QtWidgets.QAction("----------", self,  enabled=False))
            popMenu.addAction(QtWidgets.QAction("Erese", self,  enabled=True, triggered= self.window.Erese))
            popMenu.exec_(QtWidgets.QCursor.pos())

def main():
    app = QtWidgets.QApplication(sys.argv)
    window = ControlWindow()
    sys.exit(app.exec_())

if __name__ == '__main__':
    main()


from random import random, randint, sample
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import os
import shutil
from DQN import DQN
from tetris import Tetris
from collections import deque, namedtuple

BATCH_SIZE = 512
GAMMA = 0.99
EPS_START = 1
EPS_END = 0.001
EPS_DECAY = 2000
SAVE_INTERVAL = 10
LEARNING_RATE = 0.001
EPOCHS = 2500
MEMORY_SIZE = 30000
SAVE_PATH = "trained_models"
SAVE_INTERVAL = 500

Transition = namedtuple('Transition', ('state', 'action', 'next_state', 'reward'))

class ReplayMemory(object):

    def __init__(self, capacity):
        self.memory = deque([],maxlen=capacity)

    def push(self, *args):
        """Save a transition"""
        self.memory.append(Transition(*args))

    def sample(self, batch_size):
        return sample(self.memory, batch_size)

    def __len__(self):
        return len(self.memory)

def select_action(epoch, results, model):
  sample = random()
  if results:
    next_actions, next_states = zip(*results.items())
    next_states = torch.stack(next_states)
    if torch.cuda.is_available():
      next_states = next_states.cuda()
  eps_threshold = EPS_END + (max(EPS_DECAY - epoch, 0) * (EPS_START - EPS_END) / EPS_DECAY)
  model.eval()
  random_action = sample <= eps_threshold
  if random_action or not results:
    if not results:
      action_index = randint(0, 8)
    else:
      action_index = randint(0, len(results)-1)
  else:
    with torch.no_grad():
      predictions = model(next_states)[:, 0]
      action_index = torch.argmax(predictions).item()
  model.train()
  next_state = next_states[action_index, :]
  action = next_actions[action_index]
  return action, next_state

def optimize_model(memory, model, optimizer, criterion):
  if len(memory) < MEMORY_SIZE / 10:
    return
  batch = memory.sample(min(len(memory), BATCH_SIZE))
  state_batch, reward_batch, next_state_batch, done_batch = zip(*batch)
  state_batch = torch.stack(tuple(state for state in state_batch))
  reward_batch = torch.from_numpy(np.array(reward_batch, dtype=np.float32)[:, None])
  next_state_batch = torch.stack(tuple(state for state in next_state_batch))
  if torch.cuda.is_available():
    state_batch = state_batch.cuda()
    reward_batch = reward_batch.cuda()
    next_state_batch = next_state_batch.cuda()

  state_action_values = model(state_batch)
  model.eval()
  with torch.no_grad():
    prediction_batch = model(next_state_batch)
  model.train()
  expected_state_action_values = torch.cat(
            tuple(reward if done else reward + GAMMA * prediction for reward, done, prediction in
                  zip(reward_batch, done_batch, prediction_batch)))[:, None]
  optimizer.zero_grad()
  loss = criterion(state_action_values, expected_state_action_values)
  loss.backward()
  optimizer.step()

def train():
  torch.cuda.seed()
  env = Tetris(20, 10)
  # env = Tetris(width = 10, height = 20, block_size = 30)
  model = DQN()
  optimizer = optim.Adam(model.parameters(), lr = LEARNING_RATE)
  criterion = nn.MSELoss()
  memory = ReplayMemory(MEMORY_SIZE)
  state = env.get_new_state()
  # state = env.reset()

  if torch.cuda.is_available():
    model.cuda()
    state = state.cuda()
  epoch = 0
  while epoch < EPOCHS:
    results = env.get_all_states()
    # results = env.get_next_states()
    if results:
      action, next_state = select_action(epoch, results, model)
    else:
      action = (randint(0,8), 0)
      next_state = state
    if torch.cuda.is_available():
      next_state = next_state.cuda()
    x, rotations = action
    reward, done = env.step(rotations, x)
    # reward, done = env.step(action, render=False)
    memory.push(state, reward, next_state, done)

    if done:
      score = env.score
      pieces = env.pieces
      # pieces = env.tetrominoes
      cleared = env.cleared
      # cleared = env.cleared_lines
      state = env.get_new_state()
      # state = env.reset()
    else:
      state = next_state
      continue
    epoch += 1
    optimize_model(memory, model, optimizer, criterion)

    print("Epoch: {}/{}, Score: {}, Pieces {}, Cleared lines: {}".format(
    epoch,
    EPOCHS,
    score,
    pieces,
    cleared))
    if epoch > 0 and epoch%SAVE_INTERVAL == 0:
      torch.save(model, "{}/tetris_{}".format(SAVE_PATH, epoch))

  torch.save(model, "{}/tetris".format(SAVE_PATH))

train()


from utilities import ParameterException, get_table_ref
import os
import boto3
import json
from datetime import datetime
from uuid import uuid4 as uuid
from boto3.dynamodb.conditions import Key

#HTTP Actions
def create_action(actionData):
  print("Creating Action")

  if not 'actionName' in actionData or not 'actionCommands' in actionData:
    raise ParameterException(400, "Invalid Parameters: missing actionName or actionCommands")
    #Validate all commands before create
  for command in actionData['actionCommands']:
    if not 'paramID' in command or not 'paramValue' in command:
      raise ParameterException(400, "Invalid actionCommand: missing paramID or paramValue")

  uid = uuid().hex
  nowtime = datetime.now().isoformat()
  action = {
      'uuid': uid,
      'actionName': actionData['actionName'],
      'actionCommands': actionData['actionCommands'],
      'created_at': nowtime,
      'updated_at': nowtime
  }

  actions_table().put_item(Item=action)

  response = {
      "isBase64Encoded": "false",
      "statusCode": 200,
      "body": json.dumps(action)
  }

  return response

def get_action(actionID):
  print("Getting Action")

  action = load_action(actionID)
  response = {
      "isBase64Encoded": "false",
      "statusCode": 200,
      "body": json.dumps(action)
  }

  return response

def get_actions():
  print("Getting Actions")
  table = actions_table()
  result = table.scan()
  actions = result['Items']
  while 'LastEvaluateKey' in result:
          result = table.scan(ExclusiveStartKey=result['LastEvaluateKey'])
          actions += result['Items']

  response = {
      "isBase64Encoded": "false",
      "statusCode": 200,
      "body": json.dumps(actions)
  }

  return response

def delete_action(actionID):
  print("Deleting Action")
  action = load_action(actionID) #makes sure action exists
  actions_table().delete_item(Key={'uuid': actionID})

  response = {
      "isBase64Encoded": "false",
      "statusCode": 200,
      "body": "{\"message\": \"Action Deleted.\"}"
  }
  return response

def update_action(actionID, actionData):
  print("Updating Action")
  if not actionID:
    raise ParameterException(400, "Invalid Parameters: Missing actionID")
  old_action = load_action(actionID) #Makes sure action exists.

  updateExpressions=[]
  attributeValues={}
  if 'actionName' in actionData:
    updateExpressions.append("actionName = :n")
    attributeValues[':n'] = actionData['actionName']
  if 'actionCommands' in actionData:
    if not validateParamters(actionData['actionCommands']):
      #new parameters do not validate
      raise ParameterException(400, "Invalid Action Parameters, missing paramID or paramValue")
    updateExpressions.append("actionCommands = :c")
    attributeValues[':c'] = actionData['actionCommands']

  if len(updateExpressions) < 1:
    #error if not updating anything
    raise ParameterException(400, "Invalid Parameters, nothing to update")

  #update time
  updateExpressions.append("updated_at = :u")
  attributeValues[':u'] = datetime.now().isoformat()

  updateExpressionStr = "set "+(",".join(updateExpressions))
  result = actions_table().update_item(
        Key={'uuid': actionID},
        UpdateExpression=updateExpressionStr,
        ExpressionAttributeValues=attributeValues)
  response = {
      "isBase64Encoded": "false",
      "statusCode": 200,
      "body": "{\"message\": \"Action updated\"}"
  }

  return response

def fire_action(actionID):
  print("Firing Action")
  action = load_action(actionID)
  update_parameters(action['actionCommands'])
  response = {
      "isBase64Encoded": "false",
      "statusCode": 200,
      "body": "{\"Message\": \"Action fired.\"}"
  }

  return response

#Helper Methods

def load_action(actionID):
  if not actionID:
    raise ParameterException(400, "Invalid parameters: missing actionID")
  response = actions_table().query(KeyConditionExpression=Key('uuid').eq(actionID))
  if len(response['Items'])!=1:
    raise ParameterException(404, "Action not found")
  return response['Items'][0]

def validateParamters(parameters):
  for parameter in parameters:
    if not 'paramID' in parameter or not 'paramValue' in parameter:
      return False

  return True

def update_parameters(params):
  nowtime = datetime.now().isoformat()

  for parameter in params:
      params_table().update_item(
        Key={'uuid': parameter['paramID']},
        UpdateExpression="set paramValue = :v, updated_at = :t",
        ExpressionAttributeValues={':v': parameter['paramValue'], ':t': nowtime}
      )

def get_actionID(event):
  if event['pathParameters'] and 'actionID' in event['pathParameters']:
    return event['pathParameters']['actionID']
  else:
    return None

def actions_table():
  return get_table_ref('ACTIONS')

def params_table():
  return get_table_ref('PARAMETERS')

def lambda_handler(event, context):
  print("Starting Devices Lambda Function")
  try:
    if event['httpMethod'] == "GET":
      actionID = get_actionID(event)
      if actionID:
        return get_action(actionID)
      else:
        return get_actions()
    elif event['httpMethod'] == "POST":
      actionID = get_actionID(event)
      if actionID:
        return fire_action(actionID)
      else:
        return create_action(json.loads(event['body']))
    elif event['httpMethod'] == "DELETE":
      actionID = get_actionID(event)
      return delete_action(actionID)
    elif event['httpMethod'] == 'PATCH':
      actionID = get_actionID(event)
      return update_action(actionID,json.loads(event['body']))
  except ParameterException as e:
    response = {
        "isBase64Encoded": "false",
        "statusCode": e.args[0],
        "body": "{\"errorMessage\": \""+e.args[1]+".\"}"
    }
    return response
  except json.JSONDecodeError as e:
    response = {
        "isBase64Encoded": "false",
        "statusCode": 400,
        "body": "{\"errorMessage\": \"Malformed JSON: "+e.args[0]+"\"}"
    }
    return response


#!/usr/bin/env python
# coding: utf-8

# In[1]:

#!pip install nltk
#nltk.download("stopwords")

import re
import numpy as np
import math
import nltk
import pandas as pd
from nltk.stem import PorterStemmer
from pathlib import Path
from scipy import stats
TOTOAL_FILES = 1095
nltk_stopwords = nltk.corpus.stopwords.words("english")

# In[4]:

with open("training_set.txt") as training_data:
    all_data = training_data.readlines()
    data_dict = {}
    for i in all_data:
        cur_list = list(map(int, i.split()))
        data_dict[cur_list[0]] = cur_list[1:]
        print(cur_list[0], data_dict[cur_list[0]])

# In[10]:

class MultinomialNaiveBayes(object):
    def __init__(self, if_filt = False, filtered_num = 5, method = "chisquare"):
        self.if_filt = if_filt
        self.filtered_num = filtered_num
        self.method = method

        self.data_dict = {}
        self.total_word_num = 0 #總共有多少字
        self.total_word_freq = 0 #字出現次數的總和

        self.all_term_dict = {} #訓練集的所有字的freq
        self.class_term_dict = {} #訓練集的各cls的字freq
        self.class_term_present_dict ={} #紀錄該class檔案中出現該字與否的次數
        self.training_data_splited= {}

        self.filtered_all_term_dict = {}
        self.filtered_class_term_dict = {}

    def count_training_terms(self, data_loc = ".\\IRTM\\"):
        porter = PorterStemmer()

        for key in self.data_dict.keys():
            cur_class_files = self.data_dict[key]
            self.class_term_dict[key] = {} #每個類別各自的字
            cur_cls_file =[]
            for cur_num in cur_class_files:

                #print(f"cur_num {cur_num}")
                #class_term_dict[key][cur_num] = {}
                with open(data_loc+str(cur_num)+".txt", mode ="r") as file:
                    text = file.readlines()
                    #print(text)
                    cls_doc_split_text =[]
                    cur_doc ={} #用dict[term]來表示document
                    for i in text:
                        sub_text = re.sub(r'[^A-Za-z ]', '', i) #去掉非英文字
                        sub_text = sub_text.split()
                        #print(sub_text)
                        for term in sub_text:
                            word = term.lower()
                            word = porter.stem(word)

                            if not len(word) or word in nltk_stopwords:
                                continue #被切完就pass
                            else:
                                cur_doc[word] = 1
                                self.total_word_freq +=1
                                if self.all_term_dict.get(word):
                                    self.all_term_dict[word]+=1

                                else:
                                    self.all_term_dict[word]=1
                                    self.total_word_num += 1

                                if self.class_term_dict[key].get(word):
                                    self.class_term_dict[key][word]+=1

                                else:
                                    self.class_term_dict[key][word] = 1
                    #print(cur_doc)
                    #cls_doc_split_text.append(cur_doc)
                cur_cls_file.append(cur_doc)
            self.training_data_splited[key]= cur_cls_file

    def build_term_present_freq_dict(self):

        N = self.total_doc_num

        cls_list = [key for key in self.data_dict.keys()]

        for term in self.all_term_dict.keys():

            self.class_term_present_dict[term] ={}
            for cls in cls_list:
                onTP=0
                offTP= 0
                onTA=0
                offTA= 0
                docs = self.training_data_splited[cls]
                #print(doc)
                for doc in docs:
                    if doc.get(term):
                        onTP+=1
                    else:
                        offTP+=1
                for other_cls in cls_list:
                    if other_cls != cls:
                        other_cls_docs = self.training_data_splited[other_cls]
                        for other_cls_doc in other_cls_docs:
                            if other_cls_doc.get(term):
                                onTA+=1
                            else:
                                offTA+=1
                #print(other_cls_doc)
                self.class_term_present_dict[term][cls] =[onTP, offTP, onTA, offTA]
                #break

        return self.class_term_present_dict

    def get_chisquare(self):
        term_chi = {}

        N = 0 # training set len
        for cls in self.data_dict.keys():
            N+= len(self.data_dict[cls]) #training set len = 195
            term_chi[cls] = []
        for term in self.class_term_present_dict.keys():
            maximum_chi = 0
            max_cls = None
            for cls in sorted(list(self.data_dict.keys())):
                #[onTP, offTP, onTA, offTA]
                cur_table =self.class_term_present_dict[term][cls]
                cur_table = np.array(cur_table)
                cur_table = np.reshape(cur_table,(2,2))
                sum_row = np.sum(cur_table, axis=1)
                sum_col = np.sum(cur_table, axis=0)
                class_term_chi =0
                for i in range(len(cur_table)):
                    for j in range(len(cur_table[0])):
                        cur_slot = cur_table[i][j]
                        expect_val = (sum_col[j]*sum_row[i]/sum(sum_row))
                        cur_chi = (cur_slot-expect_val)**2/(expect_val)
                        class_term_chi+=cur_chi
                if class_term_chi>maximum_chi:
                    maximum_chi = class_term_chi
                    max_cls = cls
            term_chi[max_cls].append((term, maximum_chi))
        for cls in self.data_dict.keys():
            term_chi[cls].sort(key = lambda x: x[1], reverse=True)

        return term_chi

    def get_LLR(self):
        term_LLR = {}

        N = 0 # training set len
        for cls in self.data_dict.keys():
            N+= len(self.data_dict[cls]) #training set len = 195
            term_LLR[cls] = []
        for term in self.class_term_present_dict.keys():
            maximum_LLR = 0
            max_cls = None
            for cls in sorted(list(self.data_dict.keys())):
                #[onTP, offTP, onTA, offTA]
                cur_table =self.class_term_present_dict[term][cls]
                n11 = cur_table[0]
                n01 = cur_table[1]
                n10 = cur_table[2]
                n00 = cur_table[3]
                class_term_LLR =0
                upper = (((n11+n01)/N)**n11)*((1-(n11+n01)/N)**n10)*((((n11+n01)/N))**n01)*((1-(n11+n01)/N)**n00)
                upper = math.log(upper)
                lower = ((n11/(n11+n10))**n11)*((1-(n11/(n11+n10)))**n10)*((n01/(n01+n00))**n01)* (1-(n01/(n01+n00)))**n00
                lower = math.log(lower)
                CurLLR = -2*(upper-lower)

                if CurLLR>maximum_LLR:
                    maximum_LLR = CurLLR
                    max_cls = cls
            term_LLR[max_cls].append((term, maximum_LLR))
        for cls in self.data_dict.keys():
            term_LLR[cls].sort(key = lambda x: x[1], reverse=True)

        return term_LLR

    def set_filtered_terms(self, ): #using chisquare
        if not self.if_filt:
            self.filtered_all_term_dict = self.all_term_dict
            self.filtered_class_term_dict = self.class_term_dict

        else:
            if self.method =="chisquare":
                term_filtered = self.get_chisquare()

            if self.method == "LLR":
                term_filtered = self.get_LLR()

            for cls in term_filtered.keys():
                self.filtered_class_term_dict[cls] ={}
                for i in range(self.filtered_num):
                    term = term_filtered[cls][i][0]

                    self.filtered_all_term_dict[term] = self.all_term_dict[term] #filtered的all term dict
                    #for cls in range(1, self.categories_num+1):
                        #self.filtered_class_term_dict[cls] = {}
                    if self.class_term_dict[cls].get(term):
                        self.filtered_class_term_dict[cls][term] = self.class_term_dict[cls].get(term) #放進filtered的class dict

            self.all_term_dict = self.filtered_all_term_dict #訓練集的所有字的freq
            self.class_term_dict = self.filtered_class_term_dict

    def read_training_data(self,loc):
        with open(loc) as training_data:
            self.total_doc_num =0
            all_data = training_data.readlines()

            for i in all_data:
                cur_list = list(map(int, i.split()))
                self.data_dict[cur_list[0]] = cur_list[1:]
                self.total_doc_num += len(cur_list[1:])
                print(cur_list[0], self.data_dict[cur_list[0]])

        #self.data_dict = data_dict

        self.categories_num = len(list(self.data_dict.keys()))
        return self.data_dict

    def get_class_prob(self): #P(c)
        #train_data = self.read_training_data()
        #data_dict = train_data[0]
        #total_doc_num = train_data[1]

        class_prob = [] #P(c) 的list

        for cls in sorted(list(self.data_dict.keys())): #排個序避免順序跑掉
            cur_item_len = len(self.data_dict[cls])
            class_prob.append(math.log(cur_item_len)-math.log(self.total_doc_num)) #取機率log值

        return class_prob

    def get_term_prob(self): #P(X=t|c)
        #term_prob = np.array.zeros((self.total_word_num, len(data_dict.keys()))) #dim:(總共要放幾個字到矩陣 幾個class ex. 500,13)
        term_prob = {}
        all_term_list = list(self.filtered_all_term_dict.keys())
        cls_term_list = sorted(list(self.filtered_class_term_dict.keys()))

        cls_total_freq = []
        for key in self.class_term_dict.keys():
            cur_cls_freq = 0
            for word in self.class_term_dict[key].keys():
                cur_frq = self.class_term_dict[key][word]
                cur_cls_freq += cur_cls_freq

            cls_total_freq.append(cur_cls_freq)

        for j in range(len(all_term_list)):
            term = all_term_list[j]
            term_prob[term] = np.zeros(self.categories_num)
            cur_all_freq = self.all_term_dict[term]
            for i in range(len(cls_term_list)):
                cls = cls_term_list[i]
                cur_cls_total_freq = cls_total_freq[i]

                if self.class_term_dict[cls].get(term):
                    cur_class_term_freq = self.class_term_dict[cls][term]+1
                else:
                    cur_class_term_freq = 1

                term_prob[term][i] = math.log(cur_class_term_freq)-math.log(cur_cls_total_freq+self.total_word_num)#term先再放class查比較順

        return term_prob

    def read_test_data(self, test_data, data_loc = ".\\IRTM\\"):
        porter = PorterStemmer()
        test_data_set = set()
        #for data in test_data:
        with open(data_loc+test_data, mode ="r") as file:
            text = file.readlines()
            for i in text:
                sub_text = re.sub(r'[^A-Za-z ]', '', i) #去掉非英文字
                sub_text = sub_text.split()
                #print(sub_text)
                for term in sub_text:
                    word = term.lower()
                    word = porter.stem(word)

                    if not len(word) or word in nltk_stopwords:
                        continue #被切完就pass
                    else:
                        test_data_set.add(word)

        return test_data_set
    def train(self, loc = "training_set.txt"):
        self.read_training_data(loc)
        self.count_training_terms()
        self.build_term_present_freq_dict()

        self.set_filtered_terms()
        self.class_prob = self.get_class_prob()
        self.term_prob = self.get_term_prob()

    def predict(self, testing_data_list): #給檔案名list去開檔 斷詞

        class_prob = self.get_class_prob()
        ans = []
        #categories_num = len(list(self.data_dict.keys()))
        for data in testing_data_list:
            #print(data)
            cur_data_prob =np.array(class_prob) #init 一個含13個P(c)的list
            cur_test_set = self.read_test_data(data)
            for term in list(cur_test_set):
                #print(self.term_prob.get(term))
                if type(self.term_prob.get(term)) == np.ndarray:
                    #print(self.term_prob.get(term))
                    temp_term_prob = self.term_prob.get(term) #會有該詞的13 classes的機率
                    #print(temp_term_prob)
                    for cls in range(self.categories_num):
                        cur_data_prob[cls]+=temp_term_prob[cls]
            #print(cur_data_prob)
            index = np.argmax(cur_data_prob)

            ans.append(index+1)

        return ans

# In[11]:

test_data = []
training_set = []
for k in data_dict.keys():
    training_set.extend(data_dict[k])

for item in Path(".\\IRTM\\").iterdir():
    cur_num = int(item.name[:-4])

    if cur_num not in training_set:
        test_data.append(item.name)
print(len(test_data))

# In[12]:

NB_test =MultinomialNaiveBayes(if_filt=True, filtered_num=15, method ="chisquare")
NB_test.train(loc = "training_set.txt")

# In[13]:

filtered_prediction = NB_test.predict(test_data)

# In[14]:

filtered_prediction

# In[15]:

result_list = {}
for idx in range(len(test_data)):
    ide = test_data[idx][:-4]
    result = filtered_prediction[idx]
    result_list[ide] = result

# In[16]:

test_pd = pd.DataFrame.from_dict(result_list, orient="index", columns=['Value'])
test_pd.index.name = "Id"
test_pd.to_csv("R09725047_劉育志.csv")

# In[ ]:


'''
Created on 2022-11-16

@author: wf
'''

import os
import webbrowser
import sys
import traceback
from skg.version import Version
from argparse import ArgumentParser
from argparse import RawDescriptionHelpFormatter
from skg.wdsearch import WikidataSearch
from skg.wikidata import Wikidata
from skg.smw import SemWiki
from skg.kg import SKG_Def
from skg.graph import Node
from skg.paper import Paper
from skg.doi import DOI
from skg.orcid import ORCID
from skg.crossref import Crossref
from skg.skgbrowser import SkgBrowser
from skg.search import SearchOptions, SearchResult
from jpcore.justpy_app import JustpyServer

class SotSog():
    """
    Standing on the shoulders of giants
    """

    def __init__(self,args=None):
        """
        constructor

        Args:
            args(): command line Arguments(optional)
        """
        self.args=args
        debug=False
        if args:
            debug=args.debug
        self.debug=debug
        Node.debug=debug
        self.wikipedia_url="https://en.wikipedia.org/wiki/Standing_on_the_shoulders_of_giants"
        self.skg_def=SKG_Def()
        self.scholar_concept=self.skg_def.concepts["Scholar"]

    def getMarkups(self,item,options:SearchOptions)->dict:
        """
        get the markups for the given item and search options

        Args:
            item(Node): the item to get the markup for
            options(SearchOptions): the search options to apply
        """
        markups={}
        do_markup=len(options.markup_names)>0
        if do_markup:
            if item.concept.name=="Paper":
                doi=getattr(item, "doi",None)
                if doi is not None:
                    crossref=Crossref()
                    if "bibtex" in options.markup_names:
                        bibentry=crossref.doiBibEntry([doi])
                        markups["bibtex"]=bibentry
                    if "scite" in options.markup_names:
                        #meta_data=crossref.doiMetaData([doi])
                        #scite_entry=crossref.asScite(meta_data)
                        if not hasattr(item, "doi_obj"):
                            item.fromDOI(doi)
                        scite_entry=item.doi_obj.asScite()
                        markups["scite"]=scite_entry
            if item.concept.name=="Scholar":
                if "smw" in options.markup_names:
                    markups["smw"]=SemWiki.asMarkup(item)
        return markups

    def wd_search(self,wd:Wikidata,search_term:str,options)->list:
        """
        do a wikidata search
        """
        items=[]
        wds=WikidataSearch(language=options.lang,debug=self.debug)
        search_options=wds.searchOptions(search_term,limit=options.limit)
        qids=[]
        for qid,itemLabel,desc in search_options:
            qids.append(qid)
        class_map=wd.getClassQids(qids)
        for qid,itemLabel,desc in search_options:
            if qid in class_map:
                class_rows=class_map[qid]
                for class_row in class_rows:
                    class_qid=class_row["class_qid"]
                    concept=self.skg_def.conceptForQid(class_qid)
                    if concept is not None:
                        wd_items=concept.cls.from_wikidata_via_id(concept,"wikiDataId", qid, lang=options.lang)
                        if len(wd_items)>0:
                            item=wd_items[0]
                            items.append(item)
                            self.handleItem(item,qid,itemLabel,desc,options)
        return items

    def handleItem(self,item,item_id,itemLabel,desc,options):
        """
        handle the given item as a search result
        """
        if options.show:
            print(f"{itemLabel}({item_id}):{desc}✅")
            print(item)
        item.markups=self.getMarkups(item,options)
        if options.show:
            for markup_name,markup in item.markups.items():
                print(f"{markup_name} markup:")
                print(markup)
            pass
        if options.open_browser:
            browser_url=item.browser_url()
            if browser_url is not None:
                print(f"opening {browser_url} in browser")
                webbrowser.open(browser_url)

    def handleItems(self,items,options):
        """
        handle the given items
        """
        for item in items:
            item_id=item.wikiDataId
            itemLabel=item.label
            desc="?"
            self.handleItem(item, item_id, itemLabel, desc, options)

    def handleDoiItem(self,item,options:SearchOptions):
        item_id=item.doi
        itemLabel=item.title
        desc=item.title
        self.handleItem(item, item_id, itemLabel, desc, options)

    def search(self,search_list,options:SearchOptions)->SearchResult:
        """
        search with the given search list

        Args:
            search_list(list): a list of search terms
            options(SearchOptions): the search options to apply
        """
        search_result=SearchResult(search_list,options)
        search_term=' '.join(search_list)
        for prefix in ["https://doi.org"]:
            if search_term.startswith(prefix):
                search_term=search_term.replace(prefix,"")
        wd=Wikidata(debug=self.debug)
        if ORCID.isORCID(search_term):
            scholar_concept=self.skg_def.concepts["Scholar"]
            items=Node.from_wikidata_via_id(scholar_concept, "orcid", search_term, options.lang)
            self.handleItems(items,options)
        elif DOI.isDOI(search_term):
            # DOI may not be referencing paper but something else
            paper_concept=self.skg_def.concepts["Paper"]
            items=Paper.from_wikidata_via_id(paper_concept, "doi", search_term, options.lang)
            self.handleItems(items,options)
            dblp_items=Paper.from_dblp_via_id(paper_concept, "doi", search_term.lower())
            if len(dblp_items)==0:
                paper=Paper()
                paper.concept=paper_concept
                paper.fromDOI(search_term)
                paper.provenance="doi"
                dblp_items=[paper]
            for item in dblp_items:
                self.handleDoiItem(item,options)
            items.extend(dblp_items)
        else:
            items=self.wd_search(wd,search_term,options)
        search_result.items=items
        return search_result

__version__ = Version.version
__date__ = Version.date
__updated__ = Version.updated

def main(argv=None): # IGNORE:C0111
    '''main program.'''

    if argv is None:
        argv=sys.argv[1:]

    program_name = os.path.basename(__file__)
    program_shortdesc = Version.description

    program_version =f"v{__version__}"
    program_build_date = str(__updated__)
    program_version_message = f'{program_name} ({program_version},{program_build_date})'

    user_name="Wolfgang Fahl"
    program_license = '''%s

  Created by %s on %s.
  Copyright 2022 Wolfgang Fahl. All rights reserved.

  Licensed under the Apache License 2.0
  http://www.apache.org/licenses/LICENSE-2.0

  Distributed on an "AS IS" basis without warranties
  or conditions of any kind, either express or implied.

USAGE
''' % (program_shortdesc, user_name,str(__date__))
    try:
        # Setup argument parser
        parser = ArgumentParser(description=program_license, formatter_class=RawDescriptionHelpFormatter)
        parser.add_argument('search', action='store', nargs='*', help="search terms")
        parser.add_argument("--about",help="show about info",action="store_true")
        parser.add_argument("--bibtex",help="output bibtex format",action="store_true")
        parser.add_argument("-d", "--debug", dest="debug", action="store_true", help="show debug info")
        parser.add_argument('--host',default=JustpyServer.getDefaultHost())
        parser.add_argument("-la", "--lang",help="language code to use",default="en")
        parser.add_argument("-li", "--limit",help="limit the number of search results",type=int,default=9)
        parser.add_argument("-nb","--nobrowser",help="do not open browser",action="store_true")
        parser.add_argument('--port',type=int,default=8765)
        parser.add_argument("--scite",help="output #scite format",action="store_true")
        parser.add_argument("--smw",help="output Semantic MediaWiki (SMW) format",action="store_true")
        parser.add_argument("--serve",help="start webserver",action="store_true")
        parser.add_argument("--wikiId",help="the id of the SMW wiki to connect with",default="ceur-ws")
        parser.add_argument('-V', '--version', action='version', version=program_version_message)
        args = parser.parse_args(argv)
        if len(argv) < 1:
            parser.print_usage()
            sys.exit(1)
        sotsog=SotSog(args)
        markup_names=[]
        if args.bibtex: markup_names.append("bibtex")
        if args.scite: markup_names.append("scite")
        if args.smw: markup_names.append("smw")
        options=SearchOptions(limit=args.limit,lang=args.lang,
                      markup_names=markup_names,
                      open_browser=not args.nobrowser)
        if args.about:
            print(program_version_message)
            print(f"see {Version.doc_url}")
            webbrowser.open(Version.doc_url)
        elif args.serve:
            skgBrowser=SkgBrowser(version=Version,sotsog=sotsog,options=options)
            url=f"http://{args.host}:{args.port}"
            webbrowser.open(url)
            skgBrowser.start(args.host, args.port,debug=args.debug)
            pass
        else:
            sotsog.search(args.search,options)
        pass
    except KeyboardInterrupt:
        ### handle keyboard interrupt ###
        return 1
    except Exception as e:
        if DEBUG:
            raise(e)
        indent = len(program_name) * " "
        sys.stderr.write(program_name + ": " + repr(e) + "\n")
        sys.stderr.write(indent + "  for help use --help")
        if args.debug:
            print(traceback.format_exc())
        return 2

DEBUG = 1
if __name__ == "__main__":
    if DEBUG:
        sys.argv.append("-d")
    sys.exit(main())


import os
import pytest
import subprocess
import sys

def _check_binaries(dir_path):
    expected_endswith = (
        "win_amd64.pyd",
        "win32.pyd",
        "darwin.so",
        "i386-linux-gnu.so",
        "x86_64-linux-gnu.so",
    )

    binaries = list(p for p in os.listdir(dir_path) if p.endswith(expected_endswith))

    assert len(binaries) == len(expected_endswith)

@pytest.mark.skipif(
    sys.version_info[:2] != (3, 7),
    reason="DEBUGPY wheels shipped for Python 3.7 only",
)
def test_install_debugpy(tmpdir):
    import install_debugpy

    install_debugpy.main(str(tmpdir))
    dir_path = os.path.join(
        str(tmpdir), "debugpy", "_vendored", "pydevd", "_pydevd_bundle"
    )
    _check_binaries(dir_path)

    dir_path = os.path.join(
        str(tmpdir), "debugpy", "_vendored", "pydevd", "_pydevd_frame_eval"
    )
    _check_binaries(dir_path)


from conda.exceptions import EnvironmentFileEmpty, EnvironmentFileNotFound

from .. import env

class YamlFileSpec:
    _environment = None
    extensions = {".yaml", ".yml"}

    def __init__(self, filename=None, **kwargs):
        self.filename = filename
        self.msg = None

    def can_handle(self):
        try:
            self._environment = env.from_file(self.filename)
            return True
        except EnvironmentFileNotFound as e:
            self.msg = str(e)
            return False
        except EnvironmentFileEmpty as e:
            self.msg = e.message
            return False
        except TypeError:
            self.msg = f"{self.filename} is not a valid yaml file."
            return False

    @property
    def environment(self):
        if not self._environment:
            self.can_handle()
        return self._environment


from __future__ import division
from __future__ import print_function

import json
import os
import pipes
import shlex
import time

from cros_utils import command_executer

TEST_THAT_PATH = '/usr/bin/test_that'
TAST_PATH = '/usr/bin/tast'
SKYLAB_PATH = '/usr/local/bin/skylab'
GS_UTIL = 'src/chromium/depot_tools/gsutil.py'
AUTOTEST_DIR = '/mnt/host/source/src/third_party/autotest/files'
CHROME_MOUNT_DIR = '/tmp/chrome_root'

def GetProfilerArgs(profiler_args):
  # Remove "--" from in front of profiler args.
  args_list = shlex.split(profiler_args)
  new_list = []
  for arg in args_list:
    if arg[0:2] == '--':
      arg = arg[2:]
    new_list.append(arg)
  args_list = new_list

  # Remove "perf_options=" from middle of profiler args.
  new_list = []
  for arg in args_list:
    idx = arg.find('perf_options=')
    if idx != -1:
      prefix = arg[0:idx]
      suffix = arg[idx + len('perf_options=') + 1:-1]
      new_arg = prefix + "'" + suffix + "'"
      new_list.append(new_arg)
    else:
      new_list.append(arg)
  args_list = new_list

  return ' '.join(args_list)

def GetDutConfigArgs(dut_config):
  return 'dut_config={}'.format(pipes.quote(json.dumps(dut_config)))

class SuiteRunner(object):
  """This defines the interface from crosperf to test script."""

  def __init__(self,
               dut_config,
               logger_to_use=None,
               log_level='verbose',
               cmd_exec=None,
               cmd_term=None):
    self.logger = logger_to_use
    self.log_level = log_level
    self._ce = cmd_exec or command_executer.GetCommandExecuter(
        self.logger, log_level=self.log_level)
    # DUT command executer.
    # Will be initialized and used within Run.
    self._ct = cmd_term or command_executer.CommandTerminator()
    self.dut_config = dut_config

  def Run(self, cros_machine, label, benchmark, test_args, profiler_args):
    machine_name = cros_machine.name
    for i in range(0, benchmark.retries + 1):
      if label.skylab:
        ret_tup = self.Skylab_Run(label, benchmark, test_args, profiler_args)
      else:
        if benchmark.suite == 'tast':
          ret_tup = self.Tast_Run(machine_name, label, benchmark)
        else:
          ret_tup = self.Test_That_Run(machine_name, label, benchmark,
                                       test_args, profiler_args)
      if ret_tup[0] != 0:
        self.logger.LogOutput('benchmark %s failed. Retries left: %s' %
                              (benchmark.name, benchmark.retries - i))
      elif i > 0:
        self.logger.LogOutput(
            'benchmark %s succeded after %s retries' % (benchmark.name, i))
        break
      else:
        self.logger.LogOutput(
            'benchmark %s succeded on first try' % benchmark.name)
        break
    return ret_tup

  def RemoveTelemetryTempFile(self, machine, chromeos_root):
    filename = 'telemetry@%s' % machine
    fullname = os.path.join(chromeos_root, 'chroot', 'tmp', filename)
    if os.path.exists(fullname):
      os.remove(fullname)

  def GenTestArgs(self, benchmark, test_args, profiler_args):
    args_list = []

    if benchmark.suite != 'telemetry_Crosperf' and profiler_args:
      self.logger.LogFatal('Tests other than telemetry_Crosperf do not '
                           'support profiler.')

    if test_args:
      # Strip double quotes off args (so we can wrap them in single
      # quotes, to pass through to Telemetry).
      if test_args[0] == '"' and test_args[-1] == '"':
        test_args = test_args[1:-1]
      args_list.append("test_args='%s'" % test_args)

    args_list.append(GetDutConfigArgs(self.dut_config))

    if not (benchmark.suite == 'telemetry_Crosperf' or
            benchmark.suite == 'crosperf_Wrapper'):
      self.logger.LogWarning('Please make sure the server test has stage for '
                             'device setup.\n')
    else:
      args_list.append('test=%s' % benchmark.test_name)
      if benchmark.suite == 'telemetry_Crosperf':
        args_list.append('run_local=%s' % benchmark.run_local)
        args_list.append(GetProfilerArgs(profiler_args))

    return args_list

  # TODO(zhizhouy): Currently do not support passing arguments or running
  # customized tast tests, as we do not have such requirements.
  def Tast_Run(self, machine, label, benchmark):
    # Remove existing tast results
    command = 'rm -rf /usr/local/autotest/results/*'
    self._ce.CrosRunCommand(
        command, machine=machine, chromeos_root=label.chromeos_root)

    command = ' '.join(
        [TAST_PATH, 'run', '-build=False', machine, benchmark.test_name])

    if self.log_level != 'verbose':
      self.logger.LogOutput('Running test.')
      self.logger.LogOutput('CMD: %s' % command)

    return self._ce.ChrootRunCommandWOutput(
        label.chromeos_root, command, command_terminator=self._ct)

  def Test_That_Run(self, machine, label, benchmark, test_args, profiler_args):
    """Run the test_that test.."""

    # Remove existing test_that results
    command = 'rm -rf /usr/local/autotest/results/*'
    self._ce.CrosRunCommand(
        command, machine=machine, chromeos_root=label.chromeos_root)

    if benchmark.suite == 'telemetry_Crosperf':
      if not os.path.isdir(label.chrome_src):
        self.logger.LogFatal('Cannot find chrome src dir to '
                             'run telemetry: %s' % label.chrome_src)
      # Check for and remove temporary file that may have been left by
      # previous telemetry runs (and which might prevent this run from
      # working).
      self.RemoveTelemetryTempFile(machine, label.chromeos_root)

    # --autotest_dir specifies which autotest directory to use.
    autotest_dir_arg = '--autotest_dir=%s' % (
        label.autotest_path if label.autotest_path else AUTOTEST_DIR)

    # --fast avoids unnecessary copies of syslogs.
    fast_arg = '--fast'
    board_arg = '--board=%s' % label.board

    args_list = self.GenTestArgs(benchmark, test_args, profiler_args)
    args_arg = '--args=%s' % pipes.quote(' '.join(args_list))

    command = ' '.join([
        TEST_THAT_PATH, autotest_dir_arg, fast_arg, board_arg, args_arg,
        machine, benchmark.suite if
        (benchmark.suite == 'telemetry_Crosperf' or
         benchmark.suite == 'crosperf_Wrapper') else benchmark.test_name
    ])

    # Use --no-ns-pid so that cros_sdk does not create a different
    # process namespace and we can kill process created easily by their
    # process group.
    chrome_root_options = ('--no-ns-pid '
                           '--chrome_root={0} --chrome_root_mount={1} '
                           'FEATURES="-usersandbox" '
                           'CHROME_ROOT={1}'.format(label.chrome_src,
                                                    CHROME_MOUNT_DIR))

    if self.log_level != 'verbose':
      self.logger.LogOutput('Running test.')
      self.logger.LogOutput('CMD: %s' % command)

    return self._ce.ChrootRunCommandWOutput(
        label.chromeos_root,
        command,
        command_terminator=self._ct,
        cros_sdk_options=chrome_root_options)

  def DownloadResult(self, label, task_id):
    gsutil_cmd = os.path.join(label.chromeos_root, GS_UTIL)
    result_dir = 'gs://chromeos-autotest-results/swarming-%s' % task_id
    download_path = os.path.join(label.chromeos_root, 'chroot/tmp')
    ls_command = '%s ls %s' % (gsutil_cmd,
                               os.path.join(result_dir, 'autoserv_test'))
    cp_command = '%s -mq cp -r %s %s' % (gsutil_cmd, result_dir, download_path)

    # Server sometimes will not be able to generate the result directory right
    # after the test. Will try to access this gs location every 60s for
    # RETRY_LIMIT mins.
    t = 0
    RETRY_LIMIT = 10
    while t < RETRY_LIMIT:
      t += 1
      status = self._ce.RunCommand(ls_command, print_to_console=False)
      if status == 0:
        break
      if t < RETRY_LIMIT:
        self.logger.LogOutput('Result directory not generated yet, '
                              'retry (%d) in 60s.' % t)
        time.sleep(60)
      else:
        self.logger.LogOutput('No result directory for task %s' % task_id)
        return status

    # Wait for 60s to make sure server finished writing to gs location.
    time.sleep(60)

    status = self._ce.RunCommand(cp_command)
    if status != 0:
      self.logger.LogOutput('Cannot download results from task %s' % task_id)
    else:
      self.logger.LogOutput('Result downloaded for task %s' % task_id)
    return status

  def Skylab_Run(self, label, benchmark, test_args, profiler_args):
    """Run the test via skylab.."""
    options = []
    if label.board:
      options.append('-board=%s' % label.board)
    if label.build:
      options.append('-image=%s' % label.build)
    # TODO: now only put toolchain pool here, user need to be able to specify
    # which pool to use. Need to request feature to not use this option at all.
    options.append('-pool=toolchain')

    args_list = self.GenTestArgs(benchmark, test_args, profiler_args)
    options.append('-test-args=%s' % pipes.quote(' '.join(args_list)))

    dimensions = []
    for dut in label.remote:
      dimensions.append('-dim dut_name:%s' % dut.rstrip('.cros'))

    command = (('%s create-test %s %s %s') % \
              (SKYLAB_PATH, ' '.join(dimensions), ' '.join(options),
               benchmark.suite if
               (benchmark.suite == 'telemetry_Crosperf' or
                benchmark.suite == 'crosperf_Wrapper')
               else benchmark.test_name))

    if self.log_level != 'verbose':
      self.logger.LogOutput('Starting skylab test.')
      self.logger.LogOutput('CMD: %s' % command)
    ret_tup = self._ce.RunCommandWOutput(command, command_terminator=self._ct)

    if ret_tup[0] != 0:
      self.logger.LogOutput('Skylab test not created successfully.')
      return ret_tup

    # Std output of the command will look like:
    # Created request at https://ci.chromium.org/../cros_test_platform/b12345
    # We want to parse it and get the id number of the task, which is the
    # number in the very end of the link address.
    task_id = ret_tup[1].strip().split('b')[-1]

    command = ('skylab wait-task %s' % task_id)
    if self.log_level != 'verbose':
      self.logger.LogOutput('Waiting for skylab test to finish.')
      self.logger.LogOutput('CMD: %s' % command)

    ret_tup = self._ce.RunCommandWOutput(command, command_terminator=self._ct)

    # The output of `wait-task` command will be a combination of verbose and a
    # json format result in the end. The json result looks like this:
    # {"task-result":
    #   {"name":"Test Platform Invocation",
    #    "state":"", "failure":false, "success":true,
    #    "task-run-id":"12345",
    #    "task-run-url":"https://ci.chromium.org/.../cros_test_platform/b12345",
    #    "task-logs-url":""
    #    },
    #  "stdout":"",
    #  "child-results":
    #    [{"name":"graphics_WebGLAquarium",
    #      "state":"", "failure":false, "success":true, "task-run-id":"",
    #      "task-run-url":"https://chromeos-swarming.appspot.com/task?id=1234",
    #      "task-logs-url":"https://stainless.corp.google.com/1234/"}
    #    ]
    # }
    # We need the task id of the child-results to download result.
    output = json.loads(ret_tup[1].split('\n')[-1])
    output = output['child-results'][0]
    if output['success']:
      task_id = output['task-run-url'].split('=')[-1]
      if self.DownloadResult(label, task_id) == 0:
        result_dir = '\nResults placed in tmp/swarming-%s\n' % task_id
        return (ret_tup[0], result_dir, ret_tup[2])
    return ret_tup

  def CommandTerminator(self):
    return self._ct

  def Terminate(self):
    self._ct.Terminate()

class MockSuiteRunner(object):
  """Mock suite runner for test."""

  def __init__(self):
    self._true = True

  def Run(self, *_args):
    if self._true:
      return [0, '', '']
    else:
      return [0, '', '']


import os
import time
import json
import mimetypes
import subprocess
import argparse
import shutil

import boto3
import pystache
import requests
from bs4 import BeautifulSoup
from botocore.exceptions import ClientError
from markdown import markdown
from mdx_gfm import GithubFlavoredMarkdownExtension

# Older releases aren't identified as proper 'releases' in github. The tarballs
# were provided by New Vector from their archives (and uploaded directly into
# S3 manually). Because the dump of tarball releases from New Vector doesn't
# map precisely onto the meta persisted in github.com/vector-im/riot-web/releases
# (i.e. there is not an archived tarball for every tag in github), I have lazily
# recreated the relevant meta here manually.
OLDER_RELEASES = [
    {
        'name': '0.8.3',
        'body': '',
        'date': '2016-10-12'
    },
    {
        'name': '0.8.2-staging',
        'body': '',
        'date': '2016-10-05'
    },
    {
        'name': '0.8.1',
        'body': '',
        'date': '2016-09-21'
    },
    {
        'name': '0.8.0',
        'body': '',
        'date': '2016-09-21'
    },
    {
        'name': '0.7.5-r3',
        'body': '',
        'date': '2016-09-02'
    },
    {
        'name': '0.7.5-r2',
        'body': '',
        'date': '2016-09-01'
    },
    {
        'name': '0.7.5-r1',
        'body': '',
        'date': '2016-08-28'
    },
    {
        'name': '0.7.4',
        'body': '',
        'date': '2016-08-11'
    },
    {
        'name': '0.7.3',
        'body': '',
        'date': '2016-06-20'
    },
    {
        'name': '0.7.2',
        'body': '',
        'date': '2016-06-02'
    },
    {
        'name': '0.7.0',
        'body': '',
        'date': '2016-06-02'
    },
    {
        'name': '0.6.1',
        'body': '',
        'date': '2016-04-22'
    },
    {
        'name': '0.6.0',
        'body': '',
        'date': '2016-04-19'
    },
    {
        'name': '0.5.0',
        'body': '',
        'date': '2016-03-30'
    },
    {
        'name': '0.4.1',
        'body': '',
        'date': '2016-03-23'
    },
    {
        'name': '0.4.0',
        'body': '',
        'date': '2016-03-23'
    },
    {
        'name': '0.3.0',
        'body': '',
        'date': '2016-03-11'
    },
    {
        'name': '0.2.0',
        'body': '',
        'date': '2016-02-25'
    }
]

def compare_version_strings(a, b):
    a = [int(x) for x in a.split('.')]
    b = [int(x) for x in b.split('.')]

    for i in range(len(max(a,b))):
        a_part = a[i] if len(a) > i else 0
        b_part = b[i] if len(b) > i else 0
        if a_part != b_part:
            return (a_part > b_part) - (a_part < b_part)
    return 0

def get_releases(owner, repo, token):
    url = 'https://api.github.com/repos/%s/%s/releases' % (owner, repo)

    headers = {'Authorization': 'token %s' % token}
    response = requests.get(url, headers=headers)
    releases = response.json()

    while response.links.get('next'):
        response = requests.get(response.links.get('next').get('url'),
                                headers=headers)
        releases += response.json()

    return releases

def get_download_link(release, extension='.tar.gz'):
    tar_gzs = [asset for asset in release.get('assets')
               if asset.get('name').endswith(extension)]
    if len(tar_gzs) != 1:
        return None
    else:
        return tar_gzs[0].get('browser_download_url', None)

def get_name(release):
    return release.get('tag_name')[1:]

def is_version_uploaded(bucket, version):
    try:
        _ = bucket.Object(version + '/').last_modified
    except ClientError as e:
        if e.response['Error']['Code'] != '404':
            raise
        else:
            return False
    return True

def index(releases, bucket, older_releases=None):
    with open('site/index.mustache', 'r') as f:
        template = f.read()

    released = [
            {'name': release.get('name')[1:],
             'body': BeautifulSoup(markdown(release.get('body'), extensions=[GithubFlavoredMarkdownExtension()]), 'html.parser').prettify() if release.get('body') else '',
             'date': release.get('created_at')[:10]}
            for release in releases
            if is_version_uploaded(bucket, get_name(release))
            ]

    if older_releases is not None:
        released += older_releases

    renderer = pystache.Renderer(missing_tags='strict')
    rendered = renderer.render(template, {'releases': released})

    bucket.put_object(Key='index.html', Body=rendered,
            ACL='public-read',
            ContentType='text/html')

    site_files = [('site/style.css', 'style.css'),
                  ('site/privacy.html', 'privacy.html')]
    for source, destination in site_files:
        with open(source, 'rb') as data:
            mime_type = mimetypes.guess_type(source)[0]
            if not mime_type:
                mime_type = 'application/octet-stream'
            bucket.put_object(Key=destination,
                              Body=data,
                              ACL='public-read',
                              ContentType=mime_type)

def upload_directory(name, root, bucket):
    for subdir, dirs, files in os.walk(root):
        for file in files:
            full_path = os.path.join(subdir, file)
            destination_path = name + full_path[len(root):]
            with open(full_path, 'rb') as data:
                mime_type = mimetypes.guess_type(full_path)[0]
                if not mime_type:
                    mime_type = 'application/octet-stream'
                print('Putting %s (%s)' % (destination_path, mime_type))
                bucket.put_object(Key=destination_path,
                                  Body=data,
                                  ACL='public-read',
                                  ContentType=mime_type)
    # If we don't do this, we can't easily test whether the 'directory' exists in S3
    bucket.put_object(Key='%s/' % name, Body='')

def upload(releases, bucket, older_releases):
    for release in older_releases:
        name = release.get('name')
        if is_version_uploaded(bucket, name):
            print('%s: (older release) already hosted in bucket' % name)
        else:
            print('%s: (older release) uploading' % name)
            upload_directory(name, 'older_riots/%s' % name, bucket)
            print('%s: (older release) uploaded' % name)

    for release in releases:
        name = get_name(release)
        if name == '0.7.3':
            # We can't process any releases from 0.7.3 and older :( (yet)
            break
        if is_version_uploaded(bucket, name):
            print('%s: already hosted in bucket' % name)
        else:
            print('%s: not hosted in bucket; fetching...' % name)
            download_url = get_download_link(release)
            response = requests.get(download_url)
            os.makedirs('/tmp/downloads/%s' % name, exist_ok=True)
            if response.status_code == 200:
                filename = '/tmp/downloads/%s/%s.tar.gz' % (name, name)
                with open(filename, 'wb') as download:
                    for chunk in response.iter_content(chunk_size=128):
                        download.write(chunk)
            print('%s: tarball downloaded; exploding...' % name)
            os.makedirs('/tmp/exploded/%s' % name, exist_ok=True)
            subprocess.call(['tar', '-zxf', filename, '-C', '/tmp/exploded/%s/' % name])
            tar_root = '/tmp/exploded/%s/%s' % (name, os.listdir('/tmp/exploded/%s/' % name)[0])

            if name != '0.9.0':
                print('%s: tarball exploded; copying default config...' % name)
                shutil.copyfile(tar_root + '/config.sample.json', tar_root + '/config.json')
            else:
                print('%s: tarball exploded; patching in a config file for 0.9.0, released without config...' % name)
                shutil.copyfile('config.0.9.0.json', tar_root + '/config.json')

            print('%s: config inserted; uploading to s3...' % name)

            upload_directory(name, tar_root, bucket)
            print('%s upload complete; now available at https://riots.im/%s' % (name, name))

def invalidate_cloudfront_cache(cloudfrontClient, distributionId):
    paths = ['/', '/index.html']
    batch = {
        'Paths': {
            'Quantity': len(paths),
            'Items': paths
        },
        'CallerReference': str(time.time())
    }
    invalidation = cloudfrontClient.create_invalidation(
        DistributionId=distributionId,
        InvalidationBatch=batch,
    )
    return batch

def do_the_needful(aws_access_key_id, aws_secret_access_key, aws_bucket,
        aws_cloudfront_distribution_id, github_token, run_index, run_upload):

    session = boto3.Session(
        aws_access_key_id=aws_access_key_id,
        aws_secret_access_key=aws_secret_access_key
    )
    s3client = session.client('s3')
    s3resource = session.resource('s3')
    cloudfrontClient = session.client('cloudfront')

    releases = get_releases('vector-im', 'riot-web', token=github_token)
    bucket = s3resource.Bucket(aws_bucket)

    if run_index:
        index(releases, bucket, OLDER_RELEASES)
    if run_upload:
        upload(releases, bucket, OLDER_RELEASES)

    invalidate_cloudfront_cache(cloudfrontClient, aws_cloudfront_distribution_id)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Build a static site hosting historic Riot Web instances')
    parser.add_argument('--index', action="store_true")
    parser.add_argument('--upload', action="store_true")
    parser.add_argument('--aws-access-key-id', required=True)
    parser.add_argument('--aws-secret-access-key', required=True)
    parser.add_argument('--aws-bucket', required=True)
    parser.add_argument('--aws-cloudfront-distribution-id', required=True)
    parser.add_argument('--github-token', required=True)
    args = parser.parse_args()

    do_the_needful(
        aws_access_key_id=args.aws_access_key_id,
        aws_secret_access_key=args.aws_secret_access_key,
        aws_bucket=args.aws_bucket,
        aws_cloudfront_distribution_id=args.aws_cloudfront_distribution_id,
        github_token=args.github_token,
        run_index=args.index,
        run_upload=args.upload,
    )

def lambda_handler(who, cares):
    do_the_needful(
        aws_access_key_id=os.environ['aws_access_key_id'],
        aws_secret_access_key=os.environ['aws_secret_access_key'],
        aws_bucket=os.environ['aws_bucket'],
        aws_cloudfront_distribution_id=os.environ['aws_cloudfront_distribution_id'],
        github_token=os.environ['github_token'],
        run_index=True,
        run_upload=True
    )

    return {
        'statusCode': 200,
        'body': json.dumps('Hello from Lambda!')
    }