from pathlib import Path
import math

def parse_input(path: str):
    return Path(path).read_text().split("\n")

def get_sounding_heights(x: int, y: int, floor_map: list):
    try:
        right = int(floor_map[y][x + 1])
    except IndexError:
        right = 9
    try:
        left = int(floor_map[y][x - 1]) if 0 <= x - 1 else 9
    except IndexError:
        left = 9
    try:
        up = int(floor_map[y - 1][x]) if 0 <= y - 1 else 9
    except IndexError:
        up = 9
    try:
        down = int(floor_map[y + 1][x])
    except IndexError:
        down = 9
    return int(right), int(left), int(up), int(down)

def find_basin(x: int, y: int, floor_map: list, included_positions):

    included_positions.add((x, y))
    right, left, up, down = get_sounding_heights(x, y, floor_map)

    if right != 9 and (x + 1, y) not in included_positions:
        included_positions.union(find_basin(x + 1, y, floor_map, included_positions))
    if left != 9 and (x - 1, y) not in included_positions:
        included_positions.union(find_basin(x - 1, y, floor_map, included_positions))
    if up != 9 and (x, y - 1) not in included_positions:
        included_positions.union(find_basin(x, y - 1, floor_map, included_positions))
    if down != 9 and (x, y + 1) not in included_positions:
        included_positions.union(find_basin(x, y + 1, floor_map, included_positions))
    return included_positions

def main():
    floor_map = parse_input("advent_of_code/2021/advent_of_code_day9/input.txt")
    total_risk = 0
    basins = []
    basin_sizes = []
    for y in range(len(floor_map)):
        for x in range(len(floor_map[0])):
            right, left, up, down = get_sounding_heights(x, y, floor_map)
            current_pos_height = int(floor_map[y][x])
            if (
                current_pos_height < right
                and current_pos_height < left
                and current_pos_height < up
                and current_pos_height < down
            ):
                total_risk += current_pos_height + 1
                basin = find_basin(x, y, floor_map, set())
                basin_sizes.append(len(basin))
                basins.append(basin)

    print(f"Part 1: {total_risk}")

    basin_sizes.sort()
    print(f"Part 2: {math.prod(basin_sizes[-3:])}")

    symbols = "#?*.-&%|<>!'"
    for i, basin in enumerate(basins):
        for x, y in basin:
            floor_map[y] = (
                floor_map[y][:x] + symbols[i % len(symbols)] + floor_map[y][x + 1 :]
            )
    with open("advent_of_code/2021/advent_of_code_day9/output.txt", "w") as out_file:
        for row in floor_map:
            out_file.write(row.replace("9", " ") + "\n")

if __name__ == "__main__":
    main()


import json
import time
import sys
import cv2
import numpy
import math
from enum import Enum

from cscore import CameraServer, VideoSource, UsbCamera, MjpegServer
from networktables import NetworkTablesInstance, NetworkTables
import ntcore

#   JSON format:
#   {
#       "team": <team number>,
#       "ntmode": <"client" or "server", "client" if unspecified>
#       "cameras": [
#           {
#               "name": <camera name>
#               "path": <path, e.g. "/dev/video0">
#               "pixel format": <"MJPEG", "YUYV", etc>   // optional
#               "width": <video mode width>              // optional
#               "height": <video mode height>            // optional
#               "fps": <video mode fps>                  // optional
#               "brightness": <percentage brightness>    // optional
#               "white balance": <"auto", "hold", value> // optional
#               "exposure": <"auto", "hold", value>      // optional
#               "properties": [                          // optional
#                   {
#                       "name": <property name>
#                       "value": <property value>
#                   }
#               ],
#               "stream": {                              // optional
#                   "properties": [
#                       {
#                           "name": <stream property name>
#                           "value": <stream property value>
#                       }
#                   ]
#               }
#           }
#       ]
#       "switched cameras": [
#           {
#               "name": <virtual camera name>
#               "key": <network table key used for selection>
#               // if NT value is a string, it's treated as a name
#               // if NT value is a double, it's treated as an integer index
#           }
#       ]
#   }

configFile = "/boot/frc.json"

class CameraConfig: pass

team = None
server = False
cameraConfigs = []
switchedCameraConfigs = []
cameras = []

def parseError(str):
    """Report parse error."""
    print("config error in '" + configFile + "': " + str, file=sys.stderr)

def readCameraConfig(config):
    """Read single camera configuration."""
    cam = CameraConfig()

    # name
    try:
        cam.name = config["name"]
    except KeyError:
        parseError("could not read camera name")
        return False

    # path
    try:
        cam.path = config["path"]
    except KeyError:
        parseError("camera '{}': could not read path".format(cam.name))
        return False

    # stream properties
    cam.streamConfig = config.get("stream")

    cam.config = config

    cameraConfigs.append(cam)
    return True

def readSwitchedCameraConfig(config):
    """Read single switched camera configuration."""
    cam = CameraConfig()

    # name
    try:
        cam.name = config["name"]
    except KeyError:
        parseError("could not read switched camera name")
        return False

    # path
    try:
        cam.key = config["key"]
    except KeyError:
        parseError("switched camera '{}': could not read key".format(cam.name))
        return False

    switchedCameraConfigs.append(cam)
    return True

def readConfig():
    """Read configuration file."""
    global team
    global server

    # parse file
    try:
        with open(configFile, "rt", encoding="utf-8") as f:
            j = json.load(f)
    except OSError as err:
        print("could not open '{}': {}".format(configFile, err), file=sys.stderr)
        return False

    # top level must be an object
    if not isinstance(j, dict):
        parseError("must be JSON object")
        return False

    # team number
    try:
        team = j["team"]
    except KeyError:
        parseError("could not read team number")
        return False

    # ntmode (optional)
    if "ntmode" in j:
        str = j["ntmode"]
        if str.lower() == "client":
            server = False
        elif str.lower() == "server":
            server = True
        else:
            parseError("could not understand ntmode value '{}'".format(str))

    # cameras
    try:
        cameras = j["cameras"]
    except KeyError:
        parseError("could not read cameras")
        return False
    for camera in cameras:
        if not readCameraConfig(camera):
            return False

    # switched cameras
    if "switched cameras" in j:
        for camera in j["switched cameras"]:
            if not readSwitchedCameraConfig(camera):
                return False

    return True

def startCamera(config):
    """Start running the camera."""
    print("Starting camera '{}' on {}".format(config.name, config.path))
    inst = CameraServer.getInstance()
    camera = UsbCamera(config.name, config.path)
    server = inst.startAutomaticCapture(camera=camera, return_server=True)

    camera.setConfigJson(json.dumps(config.config))
    camera.setConnectionStrategy(VideoSource.ConnectionStrategy.kKeepOpen)

    if config.streamConfig is not None:
        server.setConfigJson(json.dumps(config.streamConfig))

    return camera

def startSwitchedCamera(config):
    """Start running the switched camera."""
    print("Starting switched camera '{}' on {}".format(config.name, config.key))
    server = CameraServer.getInstance().addSwitchedCamera(config.name)

    def listener(fromobj, key, value, isNew):
        if isinstance(value, float):
            i = int(value)
            if i >= 0 and i < len(cameras):
              server.setSource(cameras[i])
        elif isinstance(value, str):
            for i in range(len(cameraConfigs)):
                if value == cameraConfigs[i].name:
                    server.setSource(cameras[i])
                    break

    NetworkTablesInstance.getDefault().getEntry(config.key).addListener(
        listener,
        ntcore.constants.NT_NOTIFY_IMMEDIATE |
        ntcore.constants.NT_NOTIFY_NEW |
        ntcore.constants.NT_NOTIFY_UPDATE)

    return server

class Pipeline:

    """
    An OpenCV pipeline generated by GRIP.
    """

    def __init__(self):
        """initializes all values to presets or None if need to be set
        """

        self.__hsv_threshold_hue = [38.84891468843968, 92.96928718228388]
        self.__hsv_threshold_saturation = [175.80937019569413, 255.0]
        self.__hsv_threshold_value = [122.30215535747061, 255.0]

        self.hsv_threshold_output = None

        self.__cv_dilate_src = self.hsv_threshold_output
        self.__cv_dilate_kernel = None
        self.__cv_dilate_anchor = (-1, -1)
        self.__cv_dilate_iterations = 6.0
        self.__cv_dilate_bordertype = cv2.BORDER_CONSTANT
        self.__cv_dilate_bordervalue = (-1)

        self.cv_dilate_output = None

        self.__cv_erode_src = self.cv_dilate_output
        self.__cv_erode_kernel = None
        self.__cv_erode_anchor = (-1, -1)
        self.__cv_erode_iterations = 4.0
        self.__cv_erode_bordertype = cv2.BORDER_CONSTANT
        self.__cv_erode_bordervalue = (-1)

        self.cv_erode_output = None

        self.__find_contours_input = self.cv_erode_output
        self.__find_contours_external_only = False

        self.find_contours_output = None

        self.__filter_contours_contours = self.find_contours_output
        self.__filter_contours_min_area = 1500.0
        self.__filter_contours_min_perimeter = 0.0
        self.__filter_contours_min_width = 0.0
        self.__filter_contours_max_width = 1000.0
        self.__filter_contours_min_height = 0.0
        self.__filter_contours_max_height = 1000.0
        self.__filter_contours_solidity = [0, 100.0]
        self.__filter_contours_max_vertices = 1000000.0
        self.__filter_contours_min_vertices = 0.0
        self.__filter_contours_min_ratio = 0.0
        self.__filter_contours_max_ratio = 1000.0

        self.filter_contours_output = None

    def process(self, source0):
        """
        Runs the pipeline and sets all outputs to new values.
        """
        # Step HSV_Threshold0:
        self.__hsv_threshold_input = source0
        (self.hsv_threshold_output) = self.__hsv_threshold(self.__hsv_threshold_input, self.__hsv_threshold_hue, self.__hsv_threshold_saturation, self.__hsv_threshold_value)

        # Step CV_dilate0:
        self.__cv_dilate_src = self.hsv_threshold_output
        (self.cv_dilate_output) = self.__cv_dilate(self.__cv_dilate_src, self.__cv_dilate_kernel, self.__cv_dilate_anchor, self.__cv_dilate_iterations, self.__cv_dilate_bordertype, self.__cv_dilate_bordervalue)

        # Step CV_erode0:
        self.__cv_erode_src = self.cv_dilate_output
        (self.cv_erode_output) = self.__cv_erode(self.__cv_erode_src, self.__cv_erode_kernel, self.__cv_erode_anchor, self.__cv_erode_iterations, self.__cv_erode_bordertype, self.__cv_erode_bordervalue)

        # Step Find_Contours0:
        self.__find_contours_input = self.cv_erode_output
        (self.find_contours_output) = self.__find_contours(self.__find_contours_input, self.__find_contours_external_only)

        # Step Filter_Contours0:
        self.__filter_contours_contours = self.find_contours_output
        (self.filter_contours_output) = self.__filter_contours(self.__filter_contours_contours, self.__filter_contours_min_area, self.__filter_contours_min_perimeter, self.__filter_contours_min_width, self.__filter_contours_max_width, self.__filter_contours_min_height, self.__filter_contours_max_height, self.__filter_contours_solidity, self.__filter_contours_max_vertices, self.__filter_contours_min_vertices, self.__filter_contours_min_ratio, self.__filter_contours_max_ratio)

    @staticmethod
    def __hsv_threshold(input, hue, sat, val):
        """Segment an image based on hue, saturation, and value ranges.
        Args:
            input: A BGR numpy.ndarray.
            hue: A list of two numbers the are the min and max hue.
            sat: A list of two numbers the are the min and max saturation.
            lum: A list of two numbers the are the min and max value.
        Returns:
            A black and white numpy.ndarray.
        """
        out = cv2.cvtColor(input, cv2.COLOR_BGR2HSV)
        return cv2.inRange(out, (hue[0], sat[0], val[0]),  (hue[1], sat[1], val[1]))

    @staticmethod
    def __cv_dilate(src, kernel, anchor, iterations, border_type, border_value):
        """Expands area of higher value in an image.
        Args:
           src: A numpy.ndarray.
           kernel: The kernel for dilation. A numpy.ndarray.
           iterations: the number of times to dilate.
           border_type: Opencv enum that represents a border type.
           border_value: value to be used for a constant border.
        Returns:
            A numpy.ndarray after dilation.
        """
        return cv2.dilate(src, kernel, anchor, iterations = (int) (iterations +0.5),
                            borderType = border_type, borderValue = border_value)

    @staticmethod
    def __cv_erode(src, kernel, anchor, iterations, border_type, border_value):
        """Expands area of lower value in an image.
        Args:
           src: A numpy.ndarray.
           kernel: The kernel for erosion. A numpy.ndarray.
           iterations: the number of times to erode.
           border_type: Opencv enum that represents a border type.
           border_value: value to be used for a constant border.
        Returns:
            A numpy.ndarray after erosion.
        """
        return cv2.erode(src, kernel, anchor, iterations = (int) (iterations +0.5),
                            borderType = border_type, borderValue = border_value)

    @staticmethod
    def __find_contours(input, external_only):
        """Sets the values of pixels in a binary image to their distance to the nearest black pixel.
        Args:
            input: A numpy.ndarray.
            external_only: A boolean. If true only external contours are found.
        Return:
            A list of numpy.ndarray where each one represents a contour.
        """
        if(external_only):
            mode = cv2.RETR_EXTERNAL
        else:
            mode = cv2.RETR_LIST
        method = cv2.CHAIN_APPROX_SIMPLE
        im2, contours, hierarchy =cv2.findContours(input, mode=mode, method=method)
        return contours

    @staticmethod
    def __filter_contours(input_contours, min_area, min_perimeter, min_width, max_width,
                        min_height, max_height, solidity, max_vertex_count, min_vertex_count,
                        min_ratio, max_ratio):
        """Filters out contours that do not meet certain criteria.
        Args:
            input_contours: Contours as a list of numpy.ndarray.
            min_area: The minimum area of a contour that will be kept.
            min_perimeter: The minimum perimeter of a contour that will be kept.
            min_width: Minimum width of a contour.
            max_width: MaxWidth maximum width.
            min_height: Minimum height.
            max_height: Maximimum height.
            solidity: The minimum and maximum solidity of a contour.
            min_vertex_count: Minimum vertex Count of the contours.
            max_vertex_count: Maximum vertex Count.
            min_ratio: Minimum ratio of width to height.
            max_ratio: Maximum ratio of width to height.
        Returns:
            Contours as a list of numpy.ndarray.
        """
        output = []
        for contour in input_contours:
            x,y,w,h = cv2.boundingRect(contour)
            if (w < min_width or w > max_width):
                continue
            if (h < min_height or h > max_height):
                continue
            area = cv2.contourArea(contour)
            if (area < min_area):
                continue
            if (cv2.arcLength(contour, True) < min_perimeter):
                continue
            hull = cv2.convexHull(contour)
            solid = 100 * area / cv2.contourArea(hull)
            if (solid < solidity[0] or solid > solidity[1]):
                continue
            if (len(contour) < min_vertex_count or len(contour) > max_vertex_count):
                continue
            ratio = (float)(w) / h
            if (ratio < min_ratio or ratio > max_ratio):
                continue
            output.append(contour)
        return output

if __name__ == "__main__":
    if len(sys.argv) >= 2:
        configFile = sys.argv[1]

    # read configuration
    if not readConfig():
        sys.exit(1)

    # start NetworkTables
    ntinst = NetworkTablesInstance.getDefault()
    if server:
        print("Setting up NetworkTables server")
        ntinst.startServer()
    else:
        print("Setting up NetworkTables client for team {}".format(team))
        ntinst.startClientTeam(team)

    # start cameras
    for config in cameraConfigs:
        cameras.append(startCamera(config))

    # start switched cameras
    for config in switchedCameraConfigs:
        startSwitchedCamera(config)

    # prefix of vision processing
    inst = CameraServer.getInstance() # 取得 Camera Server
    pipeline = Pipeline() # 定義 pipeline 為 Pipeline() 類別
    videoCap = inst.getVideo(name='cam 0') # 定義 videoCap 為 cam 0 的影像
    output = inst.putVideo('vs', 320, 240) # 定義 output 為 Dashboard 中新的影像來源 名稱為 vs 寬度 320 高度240

    # preallocate new images
    img = numpy.zeros(shape=(240, 320, 3), dtype=numpy.uint8) # 預先建立 img 圖像 內容全為0
    frame = videoCap.grabFrame(None, timeout=0.5)[1] # 預先從 videoCap 中取得一次畫面

    # NetworkTables init
    table = NetworkTables.getTable('/vision/vs') # 預先取得 Network Table 中的 /vision/vs 內容
    table.setDefaultNumber('x', -1)
    table.setDefaultNumber('y', -1)
    table.setDefaultNumber('area', -1) # 設定 x,y,area 值預設為 -1 避免誤判

    # loop forever
    while True:
        # get image
        frame = videoCap.grabFrame(frame, timeout=0.5)[1] # 從 videoCap 中取得影像
        img = frame # img 定為鏡頭原像

        # processing
        pipeline.process(frame) # 使用 GRIP 的流程進行處理

        # process with any result
        biggest = {'x': -1, 'y': -1, 'area': -1} # 定義最大的球
        for contour in pipeline.filter_contours_output: # 針對每一個找到的結果判別
            # get point
            x, y, w, h = cv2.boundingRect(contour) # 尋找端點

           # compare with biggest
            a = cv2.contourArea(contour)
            if a > biggest['area']:
                biggest['x'] = x+w/2 # 中心點 x 座標
                biggest['y'] = y # 中心點 y 座標
                biggest['area'] = a

            # draw on output img
            img = cv2.line(img, (x, y), (x+w, y), (0, 255, 0), 2) # 標出六角形中心點(十字)
            img = cv2.line(img, (int(x+w/2), y-h), (int(x+w/2), y+h), (0, 255, 0), 2)
            img = cv2.putText(img, 'vs', (x, y-30), cv2.FONT_HERSHEY_DUPLEX, 0.5, (0, 0, 0))
            text = 'x: ' + str(x+w/2) + 'y: ' + str(y) # 對六角形標示中心點座標
            img = cv2.putText(img, text, (x, y-30), cv2.FONT_HERSHEY_DUPLEX, 0.5, (0, 255, 0))

        # push value onto Table
        table.putNumber('x', biggest['x']) # 將最大的球的座標值送上 Network Table 方便從機台的程式讀取
        table.putNumber('y', biggest['y'])
        table.putNumber('area', biggest['area'])

        # output image onto dashboard
        output.putFrame(img) # 將繪圖過的圖片輸出至 Dashboard

        # sleep to wait update
        # time.sleep(10)

2020 raspberry pi vision written by Brian Chen.py


from flask_script import Manager
from flask import Flask

def create_app():
    app = Flask(__name__)
    app.config['JSON_AS_ASCII'] = False
    return app

application = create_app()

if __name__ == "__main__":
    application.run()


import torchvision.models as models
import _pickle as cPickle

#reshaping weights to load to modified resnet18

resnet18 = models.resnet18(pretrained=True)

state_dict_new={}
for key in resnet18.state_dict():
    print(key, resnet18.state_dict()[key].shape)
    var=resnet18.state_dict()[key]
    if (key.find('conv') != -1 or key.find('downsample.0') != -1):
        x = var.unsqueeze(0)
        x=x.permute(1,2,0,3,4)
        state_dict_new[key]=x
    else:
         state_dict_new[key]=resnet18.state_dict()[key]

output = open('resnet18_weights.pkl', 'wb')
cPickle.dump(state_dict_new, output, 2)
output.close()


import matplotlib.pyplot as plt
import pyspark
from pyspark import SparkContext
from pyspark.sql import SparkSession, SQLContext, HiveContext
from pyspark.sql.functions import *
from fpdf import FPDF
from PIL import Image
import glob
import matplotlib.cm as cm
import numpy as np

import pymongo
import pandas as pd
import json

WIDTH = 210
HEIGHT = 297

#HELPER FUNCTIONS
def read_from_database(mongo_ip):
    iris = sqlC.read.format("com.mongodb.spark.sql.DefaultSource").option("uri", mongo_ip + "Iris").load()
    iris.createOrReplaceTempView("iris")
    iris = sqlC.sql("SELECT * FROM iris")
    #iris = iris.withColumn("Processor", iris["Processor"].cast("string"))
    return iris

def load_to_database(path_to_xlsx, database_name):
    client = pymongo.MongoClient("mongodb://localhost:27017") #connect to local server
    df = pd.read_excel(path_to_xlsx) #read from xlsx file
    data = df.to_dict(orient = "records") #insert data into a dictionary and choose the type of values as records
    db = client[database_name] #construct db
    db.Iris.insert_many(data) #insert data into db

#by creating RDD's manually, optimize the memory used and shorten execution time.
def create_partition(dataframe, sparkcontext):
    dataframe.repartition(6).createOrReplaceTempView('sf_view')
    sparkcontext.catalog.cacheTable('sf_view')
    dataframe = sparkcontext.table('sf_view')
    return dataframe

#Extract a dataframe for every client using _Total values for Processor
def extract_dataframes(dataframe):
    clients = dataframe.withColumn('Server_Name', dataframe['Server_Name'].substr(0, 3)).select('Server_Name').distinct().toPandas()["Server_Name"].values.tolist()
    client_dataframes = {}
    for client in clients:
        client_dataframes[client] = dataframe.filter(dataframe.Server_Name.contains(client) & dataframe.Processor.contains('_Total')) #take every row for a client which includes total processor usage data.
    return clients, client_dataframes

class Client:
    def __init__(self, client_dataframe, sparkcontext, client_name):
        self.client_dataframe = client_dataframe
        self.sparkcontext = sparkcontext

        self.client_name = client_name
        self.hostnames = self.client_dataframe.withColumn('Server_Name', self.client_dataframe['Server_Name'].substr(5, 1)).select('Server_Name').distinct().toPandas()["Server_Name"].values.tolist()
        self.hostnames = sorted(self.hostnames) #sort hostnames alphabetically

        self.hosts_dataframes = {} #Holds dataframes for every host of a client. Keys are hostnames.

        for host in self.hostnames:
            self.hosts_dataframes[host] = self.client_dataframe.filter(self.client_dataframe.Server_Name.contains(host)) #take every row includes the name of the host

        #mesaj gidiyo motor donmuyo hero kodu degmis olabilir

#Report template for CPU usage
class ClientReport:
    def __init__(self, client):

        #client object will be passed and graphs related to that client will be printed as a seperate pdf.
        self.client = client

        self.chart_row = 1
        self.figure_num = 0
        self.page_num = 1
        self.minus = 0
        self.figure_per_page_num = 0

        self.create_pdf()

    def add_figure(self, path, pdf, width, is_last):

        if(self.figure_num % 2 == 0 or width > WIDTH/2-5):

            if(width > WIDTH/2-5):
                pdf.image(path, 40, 80*self.chart_row-self.minus, width)
                self.figure_num += 2
                self.figure_per_page_num +=2
            else:
                pdf.image(path, 5, 80*self.chart_row-self.minus, width)
                self.figure_per_page_num += 1
                self.figure_num += 1
            #self.y = 80*chart_row-minus
            #self.x = WIDTH/2+5
        else:
            pdf.image(path, WIDTH/2+5, 80*self.chart_row-self.minus, width)
            self.figure_per_page_num +=1
            self.chart_row += 1
            self.figure_num += 1
            #self.y = 80*self.chart_row-self.minus
            #self.x = 5

        if (not is_last):
            if(self.page_num > 1 and self.figure_per_page_num == 6):
                pdf.add_page()
                self.footer(pdf)
                self.line_header(pdf)
                self.figure_per_page_num = 0
                self.chart_row = 1
                self.page_num += 1

            elif(self.page_num == 1 and self.figure_per_page_num == 4):
                pdf.add_page()
                self.footer(pdf)
                self.line_header(pdf)
                self.figure_per_page_num = 0
                self.chart_row = 1
                self.page_num += 1
                self.minus = 50

    def usage_for_every_host(self, pdf):
        i = 0
        is_last = False
        for host in self.client.hostnames:
            i+=1
            if(i == len(self.client.hostnames)):
                is_last == True
            # if(self.page_num == 1):
            #     offset = self.chart_row*50
            # else:
            offset = 75
            if(self.figure_per_page_num==0 and self.page_num > 1):
                offset = 23
            elif(self.figure_per_page_num != 0):
                offset += 7

            #pdf.cell(w=0, txt = "Charts for host "+host, ln = 1, align = 'L')
            pdf.ln(offset)
            pdf.set_font('Arial', '',10)
            pdf.write(3,"            Charts for host "+host);
            title = "Hourly CPU Usage for " + self.client.client_name + " with host " + host
            title2 = "Weekly CPU Usage for " + self.client.client_name + " with host " + host
            self.hourly_line_chart(title, self.client.hosts_dataframes[host], host, self.client.client_name)
            self.weekly_bar_chart(title2, self.client.hosts_dataframes[host], host, self.client.client_name)
            path = "/home/basan/internship-codes/charts/"+self.client.client_name+host+'.jpg'
            path2 = "/home/basan/internship-codes/charts/weekly"+self.client.client_name+host+'.jpg'
            width = WIDTH/2-5
            self.add_figure(path, pdf, width, is_last)
            self.add_figure(path2, pdf, width, is_last)

    #function for PDF output.
    def traingle_header(self, pdf, day):
        pdf.image("/home/basan/internship-codes/letterhead.png",0,0,WIDTH)
        pdf.image("/home/basan/internship-codes/ibmlogo.png", 170, 20, 25)
        pdf.set_font("Arial", size = 24)
        pdf.ln(35)
        pdf.cell(w=0, txt = "CPU Report", ln = 5, align = 'L')
        pdf.set_font("Arial", size = 15)
        pdf.ln(10)
        pdf.cell(w=0, txt = f'{day}', ln = 4, align = 'L')
        #pdf.cell(200, 30, txt = "Report Example 1", ln = 1, align = 'C')

    def line_header(self, pdf):
        #pdf.image("/home/basan/internship-codes/arkaheader.png",0,0,WIDTH)
        pdf.image("/home/basan/internship-codes/ibmlogo.png", 183, 20, 15)

    def footer(self, pdf):
        pdf.image("/home/basan/internship-codes/altheader.png", 0, 285, WIDTH)
        pdf.set_margins(0,0,0)
        pdf.set_y(-8.3)
        pdf.cell(w= 0, h=0, txt = str(pdf.page_no()) + '        ', ln = 3, align = 'R')
        pdf.set_y(0)

    def create_pdf(self):

        pdf = FPDF()
        pdf.set_auto_page_break(False, 0)
        pdf.add_page()
        pdf.set_margins(13,0,0)
        day = "01/09/2021"
        self.traingle_header(pdf, day)
        self.footer(pdf)
        #pdf.set_auto_page_break(True)
        self.usage_for_every_host(pdf)
        path, width = self.all_hosts_pie_chart(pdf)
        self.add_figure(path, pdf, width, True)

        pdf.output("cpu-reports/" + self.client.client_name+"-example1.pdf")

    #a chart that shows weekly CPU usage for an host of a single client. (7 days)
    #Hangi gunler daha cok kullanilmis
    def weekly_bar_chart(self, title, df, hostname, clientname):
        plt.clf()
        client_days = df.select(dayofweek('HR_Time')).distinct().orderBy('dayofweek(HR_Time)')
        #empty arrays to be filled (will be used in matplotlib graphs)
        labels = ["Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun"]
        x =[]
        y = []
        temp = []
        weeks = 4*24

        for aday in client_days.collect():
            x.append(aday[0])
            #her saat icin avg processor time columni topla
            df_for_day = df.filter(dayofweek('HR_Time') == lit(aday[0])).groupBy().sum('AVG_%_Processor_Time')
            temp.append(df_for_day.toPandas()["sum(AVG_%_Processor_Time)"].values.tolist())

        for i in range(0,len(temp)):
            y.append(temp[i][0])

        y1 = [value / weeks for value in y]

        plt.rcParams['axes.edgecolor']='#333F4B'
        plt.rcParams['axes.linewidth']=0.8
        plt.rcParams['xtick.color']='#333F4B'
        plt.rcParams['ytick.color']='#333F4B'

        plt.xticks(x, labels)
        plt.title(title)
        plt.bar(x,y1,color=(0.2, 0.4, 0.6, 0.6))
        #ax = sns.barplot(y= "Deaths", x = "Causes", data = deaths_pd, palette=("Blues_d"))
        #sns.set_context("poster")
        plt.savefig('./charts/weekly'+clientname+hostname+'.jpg')

    #a chart that shows hourly CPU usage for an host of a single client (24 hours).
    #Hangi saatler daha cok kullanilmis
    def hourly_line_chart(self, title, df, hostname, clientname):
        plt.clf()
        client_hours = df.select(hour('HR_Time')).distinct().orderBy('hour(HR_Time)')
        #empty arrays to be filled (will be used in matplotlib graphs)
        x = []
        y = []
        temp = []
        days = 31

        for anhour in client_hours.collect():
            x.append(anhour[0]) #insert hour names into array to be used in x axis
            #her saat icin avg processor time columni topla
            df_for_hour = df.filter(hour('HR_Time') == lit(anhour[0])).groupBy().sum('AVG_%_Processor_Time')
            temp.append(df_for_hour.toPandas()["sum(AVG_%_Processor_Time)"].values.tolist())

        for i in range(0,len(temp)):
            y.append(temp[i][0])

        y1 = [value / days for value in y]

        plt.rcParams['axes.edgecolor']='#333F4B'
        plt.rcParams['axes.linewidth']=0.8
        plt.rcParams['xtick.color']='#333F4B'
        plt.rcParams['ytick.color']='#333F4B'

        plt.xticks(x)
        plt.title(title)
        plt.plot(x,y1)
        plt.grid(True)
        plt.savefig('./charts/'+clientname+hostname+'.jpg')

    #a chart that shows total percentage usage of CPU of every host
    #Hangi host daha cok kullanmis
    def all_hosts_pie_chart(self, pdf):
        plt.clf()
        labels = self.client.hostnames
        sizes = []
        d2_sizes = []
        for host in self.client.hostnames:
            df = self.client.hosts_dataframes[host].select('AVG_%_Processor_Time').groupBy().sum()
            d2_sizes.append(df.toPandas()["sum(AVG_%_Processor_Time)"].values.tolist())

        for i in range(0,len(d2_sizes)):
            sizes.append(d2_sizes[i][0])

        #fig1, ax1 = plt.subplots(figsize=(3, 3))
        #fig1, ax1 = plt.subplots()
        #fig1.subplots_adjust(0.1,0,1,1)

        colors = cm.rainbow(np.linspace(0, 1, len(sizes)))
        plt.gca().axis("equal")
        plt.pie(sizes, labels=labels, colors=colors, autopct = '%1.1f%%', pctdistance=1.25, labeldistance=0.9, textprops={'fontsize': 8})
        plt.title("Total Percentage Usage per Host in 1 month")
        path = "/home/basan/internship-codes/charts/"+self.client.client_name+'allhosts.jpg'
        #legend_labels = ['%s, %1.1f %%' % (l, s) for l, s in zip(labels, sizes)]
        #plt.legend(pie[0], labels=labels, bbox_to_anchor=(0.5,0.5), loc='center right', fontsize=8)
        #plt.subplots_adjust(left=0.1, bottom=0.1, right=0.11)
        plt.savefig(path)
        width = 3*WIDTH/5
        return path, width

class OverallReport:
    def __init__(self, dataframe):
        self.dataframe = dataframe

    def client_usage(self):
        client_hours = []
        cli = []
        client_dataframes = []

        for i in self.clients:
            client_dataframes.append(self.dataframe.filter(self.dataframe.Server_Name.contains(i) & self.dataframe.Processor.contains('_Total')))

        for client in client_dataframes:
            df = client.select('AVG_%_Processor_Time').groupBy().sum()
            client_hours.append(df.toPandas()["sum(AVG_%_Processor_Time)"].values.tolist())

        for c in client_hours:
            cli.append(c[0])

        sizes = cli
        labels = ['aig', 'eti', 'tuv', 'aho', 'zor']

        fig1, ax1 = plt.subplots(figsize=(3, 3))
        fig1.subplots_adjust(0.1,0,1,1)

        theme = plt.get_cmap('flag')
        ax1.set_prop_cycle("color", [theme(1. * i / len(sizes)) for i in range(len(sizes))])

        ax1.pie(sizes, labels=labels, autopct='%1.1f%%', shadow=True, startangle=90, radius=100)
        ax1.axis('equal')  # Equal aspect ratio ensures that pie is drawn as a circle.
        plt.title("Total Client Usage")
        plt.savefig('client_usage.jpg')

if __name__ == "__main__":
    conf = pyspark.SparkConf().set("spark.jars.packages", "org.mongodb.spark:mongo-spark-connector_2.12:3.0.1").setMaster("local").setAppName("newApp").setAll([("spark.driver.memory", "15g"), ("spark.executer.memory", "20g")])
    sc = SparkContext(conf=conf)

    sqlC = SQLContext(sc)

    spark = SparkSession.builder \
    .master('local[*]') \
    .config("spark.driver.memory", "15g") \
    .appName('newApp') \
    .getOrCreate()

    mongo_ip = "mongodb://localhost:27017/basanto."
    iris = read_from_database(mongo_ip)
    dataframe = create_partition(iris, spark)

    clients = [] #names of clients
    client_dataframes = {} #dataframes of clients
    clients, client_dataframes = extract_dataframes(dataframe)
    client_objects = {} #dictionary of client objects

    for client in clients:
        client_objects[client] = Client(client_dataframes[client], spark, client)
        print(client)
        ClientReport(client_objects[client])

    OverallReport(dataframe)


import math

import mesa

from .model import State, VirusOnNetwork, number_infected

def network_portrayal(G):
    # The model ensures there is always 1 agent per node

    def node_color(agent):
        return {State.INFECTED: "#FF0000", State.SUSCEPTIBLE: "#008000"}.get(
            agent.state, "#808080"
        )

    def edge_color(agent1, agent2):
        if State.RESISTANT in (agent1.state, agent2.state):
            return "#000000"
        return "#e8e8e8"

    def edge_width(agent1, agent2):
        if State.RESISTANT in (agent1.state, agent2.state):
            return 3
        return 2

    def get_agents(source, target):
        return G.nodes[source]["agent"][0], G.nodes[target]["agent"][0]

    portrayal = {}
    portrayal["nodes"] = [
        {
            "size": 6,
            "color": node_color(agents[0]),
            "tooltip": f"id: {agents[0].unique_id}
state: {agents[0].state.name}",
        }
        for (_, agents) in G.nodes.data("agent")
    ]

    portrayal["edges"] = [
        {
            "source": source,
            "target": target,
            "color": edge_color(*get_agents(source, target)),
            "width": edge_width(*get_agents(source, target)),
        }
        for (source, target) in G.edges
    ]

    return portrayal

network = mesa.visualization.NetworkModule(network_portrayal, 500, 500)
chart = mesa.visualization.ChartModule(
    [
        {"Label": "Infected", "Color": "#FF0000"},
        {"Label": "Susceptible", "Color": "#008000"},
        {"Label": "Resistant", "Color": "#808080"},
    ]
)

def get_resistant_susceptible_ratio(model):
    ratio = model.resistant_susceptible_ratio()
    ratio_text = "∞" if ratio is math.inf else f"{ratio:.2f}"
    infected_text = str(number_infected(model))

    return "Resistant/Susceptible Ratio: {}
Infected Remaining: {}".format(
        ratio_text, infected_text
    )

model_params = {
    "num_nodes": mesa.visualization.Slider(
        "Number of agents",
        10,
        10,
        100,
        1,
        description="Choose how many agents to include in the model",
    ),
    "avg_node_degree": mesa.visualization.Slider(
        "Avg Node Degree", 3, 3, 8, 1, description="Avg Node Degree"
    ),
    "initial_outbreak_size": mesa.visualization.Slider(
        "Initial Outbreak Size",
        1,
        1,
        10,
        1,
        description="Initial Outbreak Size",
    ),
    "virus_spread_chance": mesa.visualization.Slider(
        "Virus Spread Chance",
        0.4,
        0.0,
        1.0,
        0.1,
        description="Probability that susceptible neighbor will be infected",
    ),
    "virus_check_frequency": mesa.visualization.Slider(
        "Virus Check Frequency",
        0.4,
        0.0,
        1.0,
        0.1,
        description="Frequency the nodes check whether they are infected by a virus",
    ),
    "recovery_chance": mesa.visualization.Slider(
        "Recovery Chance",
        0.3,
        0.0,
        1.0,
        0.1,
        description="Probability that the virus will be removed",
    ),
    "gain_resistance_chance": mesa.visualization.Slider(
        "Gain Resistance Chance",
        0.5,
        0.0,
        1.0,
        0.1,
        description="Probability that a recovered agent will become "
        "resistant to this virus in the future",
    ),
}

server = mesa.visualization.ModularServer(
    VirusOnNetwork,
    [network, get_resistant_susceptible_ratio, chart],
    "Virus Model",
    model_params,
)
server.port = 8521


import click
import json
from pyattck import Attck
from rich import print

from .download import download_json
from .database import load_database

attack = Attck()

class SharedOptions(click.Command):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.params.insert(
            0,
            click.core.Option(
                (
                    "-u",
                    "--url",
                ),
                envvar="ATTCK_BASEURL",
                default="https://raw.githubusercontent.com/mitre/cti/master/",
                help="URL to use when download ATT&CK data.",
            ),
        )

def load_matrices(base_url):
    matrices = [
        {
            "name": "enterprise",
            "url_ending": "enterprise-attack/enterprise-attack.json",
        },
        {
            "name": "mobile",
            "url_ending": "mobile-attack/mobile-attack.json",
        },
        {
            "name": "ics",
            "url_ending": "ics-attack/ics-attack.json",
        },
    ]
    for matrix in matrices:
        print(f"Downloading JSON for {matrix['name']}.")
        download_url = f"{base_url}{matrix['url_ending']}"
        matrix["download_url"] = download_url
        matrix["json"] = download_json(download_url)
        matrix["objects"] = matrix["json"]["objects"]

    return matrices

@click.group()
def cli():
    pass

@cli.command(cls=SharedOptions)
@click.option(
    "-d",
    "--data_dir",
    envvar="ATTCK_DATADIR",
    default="./data",
)
def download(url, data_dir):

    matrices = load_matrices(url)
    for matrix in matrices:
        with open(f"{data_dir}/{matrix['name']}-attack.json", "w") as output_jsonfile:
            json.dump(matrix["json"], output_jsonfile)

@cli.command(cls=SharedOptions)
def setup(url):

    matrices = load_matrices(url)
    load_database(matrices)

@cli.command()
@click.option("-m", "--matrix", envvar="ATTCKCLI_MATRIX", default="enterprise")
@click.option("--reference/--no-reference", default=False)
def list(matrix, reference):
    techniques = []

    for technique in getattr(attack, matrix).techniques:
        # for technique in attack.enterprise.techniques:
        technique_info = {
            "id": technique.id,
            "name": technique.name,
        }
        if reference:
            technique_info["reference"] = technique.reference
        if len(technique.subtechniques) > 0:
            technique_info["subtechniques"] = []
            for subtechnique in technique.subtechniques:
                subtechnique_info = {
                    "id": subtechnique.id,
                    "name": subtechnique.name,
                }
                if reference:
                    subtechnique_info["reference"] = subtechnique.reference

                technique_info["subtechniques"].append(subtechnique_info)

        techniques.append(technique_info)

    print(techniques)

if __name__ == "__main__":
    cli()


import numpy as np
from nltk.stem import WordNetLemmatizer
from nltk import word_tokenize
from nltk.corpus import stopwords
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder
from utils.hw2 import load_data, save_prediction, read_vocab
from utils.general import sigmoid, tanh, show_keras_model

from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Embedding, Conv1D, MaxPooling1D, Dense, GlobalMaxPooling1D

from sklearn.model_selection import train_test_split

data = load_data("train.txt")

X, y = data.text, data.target
X_train, X_dev, y_train, y_dev = train_test_split(X[:2000], y[:2000], test_size=0.3)

MAX_LENGTH = 100
wnet = WordNetLemmatizer()
vocab = read_vocab("vocab.txt")

unknown = vocab['__unknown__']
X_train2 = [[vocab.get(wnet.lemmatize(w), unknown) for w in word_tokenize(sent)] for sent in X_train]
X_dev2 = [[vocab.get(wnet.lemmatize(w), unknown) for w in word_tokenize(sent)] for sent in X_dev]

def trim_X(X, max_length=100, default=vocab['.']):
    for i in range(len(X)):
        if len(X[i]) > max_length:
            X[i] = X[i][:max_length]
        elif len(X[i]) < max_length:
            X[i] = X[i] + [default] * (max_length - len(X[i]))
    return np.array(X)

X_train2 = trim_X(X_train2, MAX_LENGTH)
X_dev2 = trim_X(X_dev2, MAX_LENGTH)

model = Sequential()
model.add(Embedding(input_dim=len(vocab), input_length=MAX_LENGTH, output_dim=1024, name="Embedding-1"))
model.add(Conv1D(filters=100, kernel_size=2, activation="tanh", name="Conv1D-1"))
model.add(GlobalMaxPooling1D(name="MaxPooling1D-1"))
model.add(Dense(1, activation="sigmoid", name="Dense-1"))

print(model.summary())
# show_keras_model(model)

model.compile(loss="binary_crossentropy", optimizer='adam', metrics=['accuracy'])
model.fit(X_train2, y_train, epochs=10, validation_data=[X_dev2, y_dev])

data = load_data("test.txt")

X, y = data.text, data.target
X = [[vocab.get(wnet.lemmatize(w), unknown) for w in word_tokenize(sent)] for sent in X]

X = trim_X(X, MAX_LENGTH)
y_predict = model.predict(X)

save_prediction(y_predict)


from flask import Flask
from googleplex_assistant.config import token
from logging.config import dictConfig

def create_app():
    dictConfig({
        'version': 1,
        'formatters': {'default': {
            'format': '[%(asctime)s] %(levelname)s in %(module)s: %(message)s',
        }},
        'handlers': {'wsgi': {
            'class': 'logging.StreamHandler',
            'stream': 'ext://flask.logging.wsgi_errors_stream',
            'formatter': 'default'
        }},
        'root': {
            'level': 'DEBUG',
            'handlers': ['wsgi']
        }
    })
    app = Flask(__name__, instance_relative_config=True)
    app.logger.debug("Parsing configuration...")

    app.logger.info("The token to include is")
    app.logger.info(token)
    from googleplex_assistant.controllers import hello, config, ifttt, dialogflow
    app.register_blueprint(hello.bp)
    app.register_blueprint(config.bp)
    app.register_blueprint(ifttt.bp)
    app.register_blueprint(dialogflow.bp)

    from googleplex_assistant.helpers import upnpclient
    upnpclient.init_app()
    upnpclient.map_port()

    return app

application = create_app()


from datetime import datetime as dt
from django.contrib.auth.decorators import login_required
from django.contrib.auth.models import Group
from django.core.urlresolvers import reverse
from django.http import HttpResponseRedirect,HttpResponseForbidden,\
    HttpResponseNotFound, HttpResponseNotAllowed, HttpResponse,\
    HttpResponseBadRequest
from lobber.userprofile.forms import CreateKeyForm
from lobber.settings import LOBBER_LOG_FILE
from django.contrib.auth.models import User, Group
from lobber.multiresponse import respond_to, make_response_dict, json_response
import lobber.log
from django.shortcuts import get_object_or_404, render_to_response
from lobber.userprofile.models import UserProfile, user_profile
from django.core.exceptions import ObjectDoesNotExist
import binascii
from lobber.share.models import Torrent
from django.utils.html import escape

logger = lobber.log.Logger("web", LOBBER_LOG_FILE)

def _list_keys(user):
        lst = []
        for p in UserProfile.objects.filter(creator=user).order_by('-id'):
            if p.expiration_date and p.expiration_date <= dt.now():
                continue
            if not p.user.username.startswith('key:'):
                continue
            lst.append(p)
        return lst

@login_required
def removekey(request,key):
    keyuser = get_object_or_404(User,username="key:%s" % key)
    keyprofile = user_profile(keyuser)
    if keyprofile.creator.username != request.user.username:
        return HttpResponseForbidden("That is not your key!")

    keyprofile.delete()
    keyuser.delete()
    return HttpResponseRedirect(reverse('lobber.userprofile.views.listkeys'))

@login_required
def addkey(request):
    if request.method == 'POST':
        form = CreateKeyForm(request.POST)
        form.fields['groups'].choices = [(g.id,g.name) for g in request.user.groups.all()]
        if form.is_valid():
            secret = create_key_user(request.user,
                                     form.cleaned_data['urlfilter'] or "^.*$",
                                     "", #TODO: tagconstraints
                                     form.cleaned_data['groups'],
                                     form.cleaned_data['expires'])
            return respond_to(request,
                              {'application/json': json_response(secret),
                               'text/html': HttpResponseRedirect("/user/key")})
    else:
        form = CreateKeyForm()
        form.fields['groups'].choices = [(g.id,g.name) for g in request.user.groups.all()]

    return respond_to(request,
                      {'application/json': HttpResponseBadRequest("invalid request"),
                       'text/html': 'userprofile/addkey.html'},{'form': form})

@login_required
def listkeys(request):
    return respond_to(request,
                      {"text/html": 'userprofile/keys.html'},{'keys': _list_keys(request.user)})

def create_key_user_profile(creator, urlfilter, tagconstraints, groups, expires=None):
    """
    Create a user profile named key:<random text>.

    Also, '$self' is substituted for the name of the newly created key-user.
    """
    creator_profile = None
    try:
        creator_profile = creator.profile.get()
    except ObjectDoesNotExist:
        pass
    if not creator_profile:
        return None

    secret = binascii.hexlify(open('/dev/urandom').read(13))
    username = 'key:%s' % secret

    user = User.objects.create_user(username, 'nomail@dev.null', username)
    urlfilter = ' '.join(urlfilter.split())
    tagconstraints = ' '.join(tagconstraints.split())

    profile = UserProfile(user=user,
                          creator=creator,
                          urlfilter=urlfilter,
                          tagconstraints=tagconstraints,
                          entitlements="",
                          expiration_date=expires)
    profile.save()

    for group in groups:
        g = Group.objects.get(pk=group[0])
        user.groups.add(g)

    return profile

def create_key_user(creator, urlfilter, tagconstraints, entitlements, expires=None):
    profile = create_key_user_profile(creator, urlfilter, tagconstraints, entitlements, expires)
    return profile.get_username()

@login_required
def user_self(req):
    lst = []
    for t in Torrent.objects.all().order_by('-creation_date')[:40]:
        if t.authz(req.user, 'r') and t.expiration_date > dt.now():
            lst.append(t)

    return render_to_response('userprofile/user.html',
                              make_response_dict(req,{'torrents': lst}))

@login_required
def ace_subjects(req):
    term = req.GET['term']
    subjects = []
    if term:
        subjects.append({'label': 'Anyone','id': 0,'type':'anyone'})
        for g in Group.objects.filter(name__contains=term):
            subjects.append({'label': "Members of '%s'" % g.name,'id':g.id,'type':'group'})
        for u in User.objects.filter(username__contains=term):
            subjects.append({'label': u.username,'id':u.id,'type':'group'})

    return json_response(subjects)

def get_profile(key):
    try:
        u = User.objects.get(username='key:%s' % key)
    except ObjectDoesNotExist:
        return None
    try:
        p = u.profile.get()
    except ObjectDoesNotExist:
        return None
    return p

def add_constraint(req, key, constraint, fun):
    p = get_profile(key)
    if not p:
        return HttpResponseNotFound("%s: key not found" % escape(key))
    if not fun(p, req.user, escape(constraint)):
        return HttpResponseNotAllowed("Permission denied.")
    p.save()
    return HttpResponse("Added %s to %s." % (escape(constraint), p))

def remove_constraint(req, key, constraint, fun):
    p = get_profile(key)
    if not p:
        return HttpResponse("%s: key not found" % escape(key))
    if not fun(p, req.user, escape(constraint)):
        return HttpResponse("Permission denied.", status=403)
    p.save()
    return HttpResponse("Removed %s from %s." % (escape(constraint), p))

@login_required
def add_urlfilter(req, key, pattern):
    return add_constraint(req, key, pattern, UserProfile.add_urlfilter)

@login_required
def remove_urlfilter(req, key, pattern):
    return remove_constraint(req, key, pattern, UserProfile.remove_urlfilter)

@login_required
def add_tagconstraint(req, key, tag):
    return remove_constraint(req, key, tag, UserProfile.add_tagconstraint)

@login_required
def remove_tagconstraint(req, key, tag):
    return remove_constraint(req, key, tag, UserProfile.remove_tagconstraint)


import os
import uuid

from pytest import fixture
from .helpers import setup_vfile
from ..api import VersionedHDF5File

# Run tests marked with @pytest.mark.slow last. See
# https://stackoverflow.com/questions/61533694/run-slow-pytest-commands-at-the-end-of-the-test-suite
def by_slow_marker(item):
    return bool(item.get_closest_marker('slow'))

def pytest_collection_modifyitems(items):
    items.sort(key=by_slow_marker)

@fixture
def filepath(tmp_path, request):
    file_name = os.path.join(tmp_path, 'file.hdf5')
    m = request.node.get_closest_marker('setup_args')
    if m is not None:
        if 'file_name' in m.kwargs.keys():
            file_name = m.kwargs['file_name']
    yield file_name

@fixture
def h5file(tmp_path, request):
    file_name = os.path.join(tmp_path, 'file.hdf5')
    version_name = None
    m = request.node.get_closest_marker('setup_args')
    if m is not None:
        if 'file_name' in m.kwargs.keys():
            file_name = m.kwargs['file_name']
        if 'name' in m.kwargs.keys():
            raise ValueError("The name argument is no longer used")
        if 'version_name' in m.kwargs.keys():
            version_name = m.kwargs['version_name']

    f = setup_vfile(file_name=file_name, version_name=version_name)
    yield f
    try:
        f.close()
    # Workaround upstream h5py bug. https://github.com/deshaw/versioned-hdf5/issues/162
    except ValueError as e:
        if e.args[0] == "Unrecognized type code -1":
            return
        raise
    except RuntimeError as e:
        if e.args[0] in ["Can't increment id ref count (can't locate ID)",
                         "Unspecified error in H5Iget_type (return value <0)",
                         "Can't retrieve file id (invalid data ID)"]:
            return
        raise

@fixture
def vfile(tmp_path, h5file):
    file = VersionedHDF5File(h5file)
    yield file
    file.close()

def generate_bad_data():
    """Generate versioned-hdf5 files with bad object dtype hash tables in them.

    See https://github.com/deshaw/versioned-hdf5/issues/256 for more information.

    Raises:
        ImportError: Raised if the user tries to generate bad data with newer versions
        of the library; you need an old version to replicate the hash table issue.
    """
    import numpy as np
    import h5py
    from versioned_hdf5 import VersionedHDF5File, __version__

    try:
        from versioned_hdf5.backend import DATA_VERSION  # noqa: F401
    except ImportError:
        DATA_VERSION = None

    if DATA_VERSION is not None:
        raise ImportError(
            f"versioned_hdf5=={__version__} installed; "
            "this file only generates bad data on versioned_hdf5 <= 1.3.14."
        )

    filename = "object_dtype_bad_hashtable_data.h5"
    with h5py.File(filename, mode="w") as f:
        vf = VersionedHDF5File(f)
        arr = np.array(
            [
                "abcd",
                "def",
                "ghi",
                "jkl",
            ],
            dtype=object
        )

        with vf.stage_version("r0") as group:
            group.create_dataset(
                "data_with_bad_hashtable",
                dtype=h5py.string_dtype(length=None),
                data=arr
            )

    filename = "object_dtype_bad_hashtable_data2.h5"
    with h5py.File(filename, mode="w") as f:
        vf = VersionedHDF5File(f)

        arr = np.array(
            [
                "abcd",
                "def",
                "ghi",
                "jkl",
            ],
            dtype=object
        )
        arr2 = np.array(
            [
                "abcd",
                "pqrs",
                "df",
                "tuvw",
                "xyz",
            ],
            dtype=object
        )
        arr3 = np.linspace(1, 100, 1000)

        with vf.stage_version("r0") as group:
            group.create_dataset(
                "data_with_bad_hashtable",
                dtype=h5py.string_dtype(length=None),
                data=arr
            )
            group.create_dataset(
                "data_with_bad_hashtable2",
                dtype=h5py.string_dtype(length=None),
                data=arr2
            )
            group.create_dataset("linspace", data=arr3)

        with vf.stage_version("r1") as group:
            group['data_with_bad_hashtable'][1] = "foo"
            group['linspace'][3] = 8

        with vf.stage_version("r2") as group:
            group['data_with_bad_hashtable'][2] = "bar"
            group['data_with_bad_hashtable2'][2] = "what"

    filename = "nested_data_old_data_version.h5"
    with h5py.File(filename, mode="w") as f:
        vf = VersionedHDF5File(f)

        with vf.stage_version('r0') as sv:
            data_group = sv.create_group('data')
            data_group.create_dataset('values',
                                      data=np.array(['1', '2', '3']),
                                      dtype=h5py.string_dtype(length=None))

    filename = "multiple_nested_data_old_data_version.h5"
    with h5py.File(filename, mode="w") as f:
        vf = VersionedHDF5File(f)

        arr = np.array(
            [
                "abcd",
                "def",
                "ghi",
                "jkl",
            ],
            dtype=object
        )

        with vf.stage_version('r0') as sv:
            g = sv.create_group('foo/bar/baz')
            g.create_dataset(
                "foo/bar/baz/values",
                dtype=h5py.string_dtype(length=None),
                data=arr
            )

    filename = "object_dtype_bad_hashtable_chunk_reuse.h5"
    with h5py.File(filename, mode="w") as f:
        vf = VersionedHDF5File(f)
        with vf.stage_version(str(0)) as sv:
            sv.create_dataset('values',
                              data=np.arange(0).astype(str).astype('O'),
                              dtype=h5py.string_dtype(length=None),
                              chunks=(10,))

    for i in range(1, 11):
        with h5py.File(filename, 'r+') as f:
            vf = VersionedHDF5File(f)
            with vf.stage_version(str(i)) as sv:
                sv['values'] = np.arange(i).astype(str).astype('O')

    filename = "object_dtype_bad_hashtable_chunk_reuse_unicode.h5"
    numbers = ['れい', 'いち', 'に', 'さん', 'し', 'ご', 'ろく', 'しち', 'はち', 'きゅう', 'じゅう']
    with h5py.File(filename, mode="w") as f:
        vf = VersionedHDF5File(f)
        with vf.stage_version(str(0)) as sv:
            sv.create_dataset('values',
                              data=np.array(numbers[:1], dtype='O'),
                              dtype=h5py.string_dtype(length=None),
                              chunks=(10,))

    for i in range(1, 11):
        with h5py.File(filename, 'r+') as f:
            vf = VersionedHDF5File(f)
            with vf.stage_version(str(i)) as sv:
                sv['values'] = np.array(numbers[:i + 1], dtype='O')

    filename = "object_dtype_bad_hashtable_chunk_reuse_multi_dim.h5"
    with h5py.File(filename, mode="w") as f:
        vf = VersionedHDF5File(f)
        with vf.stage_version(str(uuid.uuid4())) as sv:
            sv.create_dataset('values',
                              data=np.array([[chr(ord('a') + ((j + k) % 10)) * 3 for j in range(4)]
                                             for k in range(4)], dtype='O'),
                              dtype=h5py.string_dtype(length=None),
                              chunks=(2, 2))

    for i in range(1, 11):
        with h5py.File(filename, 'r+') as f:
            vf = VersionedHDF5File(f)
            with vf.stage_version(str(uuid.uuid4())) as sv:
                sv['values'] = np.array([[chr(ord('a') + ((i + j + k) % 10)) * 3 for j in range(4)]
                                         for k in range(4)], dtype='O')


import numpy as np
import matplotlib
matplotlib.rcParams['text.usetex'] = True
import matplotlib.pyplot as plt
import pickle

#from mpl_toolkits.axes_grid1.inset_locator import inset_axes

outfile = "plot_corr_triple_ce.pdf"; fbase =  "../cdataD3_4-10_" ;
ymax=(0.10, 0.10, 0.05); ylabel="$\langle a_1 a_2 a_3^* \\rangle $"

#-- canvas options --

fig, ax = plt.subplots(ncols=3, nrows=2, figsize=(10,7))

runs = ("e1", "c1", "c5")

for i in (0,1,2):

    run=runs[i]

    fname = fbase + run +".txt"

    print(fname)
    iax = ax[0][i]
    dat = np.loadtxt(fname, usecols=range(1,12))

    w    = dat[:,0]**2 + dat[:,1]**2
    corr = dat[:,2]
    bg   = corr/dat[:,3]
    phi   = dat[:,4]
    w2    = dat[:,5]**2 + dat[:,6]**2
    corr2 = dat[:,7]
    phi2  = dat[:,9]
    group = dat[:,10]

    ind = (group == 0 )

    iax.plot(w[ind], corr[ind],  'or',  mfc='none', ms=5,  mew=1, label="resonant")
    iax.plot(w[ind], bg[ind],    'og',  mfc='none', ms=5,  mew=0.5, label="background")
    iax.plot(w[ind], corr2[ind], 'db',  mfc='none', ms=5,  mew=0.5, label="non resonant")

    ind = (group == 1 )

    iax.plot(w[ind], corr[ind],  '+r',  mfc='none', ms=4,  mew=0.5)
    iax.plot(w[ind], corr2[ind], '+b',  mfc='none', ms=4,  mew=0.2)

    ind = (group == 2 )

    iax.plot(w[ind], corr[ind],  'xr',  mfc='none', ms=4,  mew=0.5)
    iax.plot(w[ind], corr2[ind], 'xb',  mfc='none', ms=4,  mew=0.2)

    ind = (group == 3 )

    iax.plot(w[ind], corr[ind],  '+r',  mfc='none', ms=4,  mew=0.5)
    iax.plot(w[ind], corr2[ind], '+b',  mfc='none', ms=4,  mew=0.2)

    ind = (group > 3 )

    iax.plot(w[ind], corr[ind],  'xr',  mfc='none', ms=4,  mew=0.5)
    iax.plot(w[ind], corr2[ind], 'xb',  mfc='none', ms=4,  mew=0.2)

    iax.set_title(run)
    iax.set(xlabel="$|k_1|^2 + |k_2|^2$", ylabel=ylabel)
    iax.legend() #frameon=False)

    iax.set_xlim(25, 175)
    iax.set_ylim(0, ymax[i])

    iax = ax[1][i]

    ind = (group == 0 )
    iax.plot(w[ind], phi[ind],  'or',  mfc='none', ms=5,  mew=1)
    iax.plot(w[ind], phi2[ind],  'db',  mfc='none', ms=5,  mew=0.5)

    ind = (group == 1 )
    iax.plot(w[ind], phi[ind],  '+r',  mfc='none', ms=4,  mew=0.5)
    iax.plot(w[ind], phi2[ind], '+b',  mfc='none', ms=4,  mew=0.2)

    ind = (group == 2 )
    iax.plot(w[ind], phi[ind],  'xr',  mfc='none', ms=4,  mew=0.5)
    iax.plot(w[ind], phi2[ind], 'xb',  mfc='none', ms=4,  mew=0.2)

    ind = (group == 3 )
    iax.plot(w[ind], phi[ind],  '+r',  mfc='none', ms=4,  mew=0.5)
    iax.plot(w[ind], phi2[ind], '+b',  mfc='none', ms=4,  mew=0.2)

    ind = (group > 3 )
    iax.plot(w[ind], phi[ind],  'xr',  mfc='none', ms=4,  mew=0.5)
    iax.plot(w[ind], phi2[ind], 'xb',  mfc='none', ms=4,  mew=0.2)

    iax.set(xlabel="$|k_1|^2 + |k_2|^2$", ylabel="$\\theta / \pi$")
    iax.set_xlim(25, 175)
    iax.set_ylim(-1, 1)

#-----------------------

plt.subplots_adjust(left=0.08, right=0.98, top=0.96, bottom=0.12, hspace=0.25, wspace=0.28)

plt.savefig(outfile, pad_inches=0)

plt.close()

#===========================================

#=========================================================


import pytest
import dash
import dash_html_components as html
import dash_extendable_graph as deg
from dash.dependencies import Input, Output, State
from dash.exceptions import PreventUpdate
import time

TIMEOUT = 20

# component responsive flag behaves as expected. This test is pulled from dash-core-components:
# https://github.com/plotly/dash-core-components/blob/dev/tests/integration/graph/test_graph_responsive.py

class WaitForTimeout(Exception):
    """This should only be raised inside the `wait_for` function."""

    pass

def wait_for(condition_function, get_message=lambda: "", *args, **kwargs):
    """
    Waits for condition_function to return True or raises WaitForTimeout.
    :param (function) condition_function: Should return True on success.
    :param args: Optional args to pass to condition_function.
    :param kwargs: Optional kwargs to pass to condition_function.
        if `timeout` is in kwargs, it will be used to override TIMEOUT
    :raises: WaitForTimeout If condition_function doesn't return True in time.
    Usage:
        def get_element(selector):
            # some code to get some element or return a `False`-y value.
        selector = '.js-plotly-plot'
        try:
            wait_for(get_element, selector)
        except WaitForTimeout:
            self.fail('element never appeared...')
        plot = get_element(selector)  # we know it exists.
    """

    def wrapped_condition_function():
        """We wrap this to alter the call base on the closure."""
        if args and kwargs:
            return condition_function(*args, **kwargs)
        if args:
            return condition_function(*args)
        if kwargs:
            return condition_function(**kwargs)
        return condition_function()

    if "timeout" in kwargs:
        timeout = kwargs["timeout"]
        del kwargs["timeout"]
    else:
        timeout = TIMEOUT

    start_time = time.time()
    while time.time() < start_time + timeout:
        if wrapped_condition_function():
            return True
        time.sleep(0.5)

    raise WaitForTimeout(get_message())

@pytest.mark.parametrize("responsive", [True, False])  # reduce iterations: None = False
@pytest.mark.parametrize("autosize", [True, False])  # reduce iterations -- None = True
@pytest.mark.parametrize("height", [600, None])
@pytest.mark.parametrize("is_responsive", [True, False, "auto"])
def test_extg006_responsive(dash_duo, responsive, autosize, height, is_responsive):
    width = 600  # reduce test iterations, just check vs height (ignore width)
    app = dash.Dash(__name__)

    header_style = dict(padding="10px", backgroundColor="yellow", flex="0 0 100px")

    graph_style = dict(padding="10px", backgroundColor="red", flex="1 0 0")

    card_style = dict(
        display="flex",
        flexFlow="column",
        backgroundColor="green",
        padding="10px",
        height="500px",
        width="1000px",
    )

    header = html.Div(
        id="header",
        style=header_style,
        children=[html.Button(id="resize", children=["Resize"])],
    )

    graph = html.Div(
        style=graph_style,
        children=[
            deg.ExtendableGraph(
                id="graph",
                responsive=is_responsive,
                style=dict(height="100%", width="100%"),
                config=dict(responsive=responsive),
                figure=dict(
                    layout=dict(autosize=autosize, height=height, width=width),
                    data=[
                        dict(x=[1, 2, 3, 4], y=[5, 4, 3, 6], line=dict(shape="spline"),)
                    ],
                ),
            )
        ],
    )

    app.layout = html.Div(
        [
            html.Div(
                [
                    "responsive: {}, ".format(responsive),
                    "autosize: {}, ".format(autosize),
                    "height: {}, ".format(height),
                    "width: {}, ".format(width),
                    "is_responsive: {}".format(is_responsive),
                ]
            ),
            html.Div(id="card", style=card_style, children=[header, graph]),
        ]
    )

    @app.callback(
        Output("header", "style"),
        [Input("resize", "n_clicks")],
        [State("header", "style")],
    )
    def resize(n_clicks, style):
        if n_clicks is None:
            raise PreventUpdate

        return dict(style, **dict(flex="0 0 200px"))

    dash_duo.start_server(app)

    # autosize=true|udefined will make the graph fit its parent on first render, responsive has no impact on that behavior
    #
    # responsive=true will make the graph resize only if autosize=true|undefined, interestingly enough, responsive=undefined
    # behaves the same as responsive=false (https://github.com/plotly/plotly.js/blob/master/src/plot_api/plot_config.js#L122)

    initial_responsive = is_responsive is True or (
        is_responsive == "auto"
        and autosize is not False
        and (height is None or width is None)
    )

    resize_responsive = is_responsive is True or (
        is_responsive == "auto"
        and responsive is True
        and autosize is not False
        and (height is None or width is None)
    )

    initial_height = (
        360
        if (initial_responsive and (height is None or is_responsive is True))
        else 450
        if height is None
        else height
    )

    resize_height = (
        260
        if (resize_responsive and (height is None or is_responsive is True))
        else initial_height
        if height is None
        else height
    )

    wait_for(
        # 500px card minus (100px header + 20px padding) minus (20px padding on container) -> 360px left
        lambda: dash_duo.wait_for_element("#graph svg.main-svg").size.get("height", -1)
        == initial_height,
        lambda: "initial graph height {}, expected {}".format(
            dash_duo.wait_for_element("#graph svg.main-svg").size.get("height", -1),
            initial_height,
        ),
    )
    dash_duo.wait_for_element("#resize").click()

    wait_for(
        # 500px card minus (200px header + 20px padding) minus (20px padding on container) -> 260px left
        lambda: dash_duo.wait_for_element("#graph svg.main-svg").size.get("height", -1)
        == resize_height,
        lambda: "resized graph height {}, expected {}".format(
            dash_duo.wait_for_element("#graph svg.main-svg").size.get("height", -1),
            resize_height,
        ),
    )


#This script is similar to context_retrieval and context_retrieval_probout.
#In addition, it facilitates two things
# 1) It randomly picks num_examples contexts out of topk most similar docs.
# 2) It calculates the average similarity/distance of context examples to the test doc.

from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM

import torch
from util.util_eval_probout import evaluate_rel, get_logger
from util.util_diverse_context_retrieval import PromptDatasetDiverseContext
import argparse

import pandas as pd
import os
import json
from pathlib import Path

def get_model_and_tokenizer(args, device):

    #Load and set the tokenizer
    tokenizer = AutoTokenizer.from_pretrained(args.model_name, padding_side="left")
    tokenizer.pad_token = tokenizer.eos_token

    #Load and send model to gpu
    #model = AutoModelForCausalLM.from_pretrained(args.model_name)
    model = AutoModelForCausalLM.from_pretrained(args.model_name, low_cpu_mem_usage=True, torch_dtype=torch.bfloat16)
    model.to(device)

    return model, tokenizer

def main(args):
    #First configure our logger
    log = get_logger(os.path.join(args.experiments_main_folder, args.experiment_folder, args.log))

    device = "cuda:"+args.device_cuda if torch.cuda.is_available() else "cpu"

    model,tokenizer = get_model_and_tokenizer(args,device)

    #queries, subjects, labels, wiki_dict = get_raw_data(args=args)
    #template = "The Wall (French: Le Mur) by Jean-Paul Sartre, a collection of short stories published in 1939 containing the eponymous story \"The Wall\", is considered one of the author's greatest existentialist works of fiction. Sartre dedicated the book to his companion Olga Kosakiewicz, a former student of Simone de Beauvoir.\nThe author of The Wall is Jean-Paul Sartre\n\nDragonsbane is a fantasy novel written by author Barbara Hambly and published by Del Rey Books in 1985.\nThe author of Dragonsbane is Barbara Hambly\n\n"
    #if args.out_triplet:
    #    template = "The Wall (French: Le Mur) by Jean-Paul Sartre, a collection of short stories published in 1939 containing the eponymous story \"The Wall\", is considered one of the author's greatest existentialist works of fiction. Sartre dedicated the book to his companion Olga Kosakiewicz, a former student of Simone de Beauvoir.\n(author of, The Wall, Jean-Paul Sartre)\n\nDragonsbane is a fantasy novel written by author Barbara Hambly and published by Del Rey Books in 1985.\n(author of, Dragonsbane, Barbara Hambly)\n\n"

    dataset = PromptDatasetDiverseContext(args=args)
    log("dataset is ready with size " + str(len(dataset)))

    log("Start predictions for the dataset")

    results_df = evaluate_rel(model, tokenizer, dataset, args, device=device, log=log)

    log("Done all the prediction! Now saving the prediction results")

    write_file = args.prediction_path
    write_file = os.path.join(args.experiments_main_folder, args.experiment_folder, write_file)
    results_df.to_csv(write_file, index=False)

if __name__ == "__main__":

    parser = argparse.ArgumentParser()

    parser.add_argument('--model_name', type=str, default="togethercomputer/GPT-JT-6B-v1", help='Currently supported model is togethercomputer/GPT-JT-6B-v1. It should work with other auto-regressive models with minor modifications.')
    parser.add_argument('-tmp', '--temperature', type=float, default=0.0001, help='The temperature of the model sampling')

    parser.add_argument('--path_data', type=str, default="DocRED/data/relation_docs_dev", help='Path to the DOCRED relation documents of evaluation set')
    parser.add_argument('--path_embeddings', type=str, default="DocRED/data/embeddings_dev.pkl", help='Path to the DOCRED document embeddings of evaluation set')

    parser.add_argument('--path_context', type=str, default="DocRED/data/relation_docs_distant", help='Path to the DOCRED relation documents to be used for in-context examples')
    parser.add_argument('--path_context_embeddings', type=str, default="DocRED/data/embeddings_train_distant.pkl", help='Path to the DOCRED in-context document embeddings')

    #Subset the documents that are extremely long. (to enable fitting into the memory)
    parser.add_argument('--doc_max_chars', type=int, default=1000, help='Remove paragraphs longer than doc_max_chars chars (for context)')

    #Below 2 lines are to subset the documents having relations between 'min_num_rels' and 'max_num_rels'. Default is no subsetting.
    parser.add_argument('--min_num_rels', type=int, default=0, help='Remove paragraphs having less relations than min_num_rels (for context)')
    parser.add_argument('--max_num_rels', type=int, default=99, help='Remove paragraphs having more relations than max_num_rels (for context)')

    #To control which relation to evaluate
    parser.add_argument('-rel', '--relation', type=str, default="P1001", help='enter the id of the relation to evaluate')

    #To control which relation to evaluate
    parser.add_argument('-dp', '--doc_prefix', type=str, default="Context: ", help='Enter the prefix to identify the context paragraph')
    #To control which relation to evaluate
    parser.add_argument('-rp', '--relation_prefix', type=str, default="Relation: ", help='Enter the prefix to identify the relation extracted')

    #To control the number of few-shot examples. NOTE: Refers to 'top-K' in paper
    parser.add_argument('-n_ex', '--num_examples', type=int, default=5, help='number of example contexts to extract triplets')
    #Total number of most-similar context examples, before random selection of num_examples. NOTE: Refers to 'top-N' in paper
    parser.add_argument('-t_k', '--topk', type=int, default=20, help='Total number of most-similar context examples considered to extract triplets')
    #In case of multiple relations exist (for the same context) use below argument to extract all relations
    parser.add_argument('-om', '--output_mode', type=str, default="multi", help='whether or not to output multiple relatons. Choose from {single, multi}')

    parser.add_argument('-bs', '--batch_size', type=int, default=4, help='Batch size')
    parser.add_argument('-mnt', '--max_new_tokens', type=int, default=200, help='max number of new tokens to be generated')
    parser.add_argument('--device_cuda', type=str, default="0", help='the cuda device to be used')

    #For adversarial experiments. NOT NEEDED FOR THE MAIN RESULTS
    parser.add_argument('--subset_context', type=float, default=1.0, help='What percentage of subset to keep. Used for noise experiments')

    parser.add_argument('-emf', '--experiments_main_folder', type=str, default="./experiments_gpt", help='path to the main experiment folder')
    parser.add_argument('-ef', '--experiment_folder', type=str, default="default_exp", help='name of subdirectory for the specific experiment')
    parser.add_argument('--prediction_path', type=str, default="predictions.csv", help='name of the prediction file')
    parser.add_argument('-l', '--log', type=str, default='train.log')

    #To control what is the separator
    parser.add_argument('-sep', '--separator', type=str, default=" <==>", help='special string to separate triplet elements')

    parser.add_argument('--seed', type=int, default=0, help='Randomness')

    args = parser.parse_args()

    full_experiment_path = os.path.join(args.experiments_main_folder, args.experiment_folder)
    full_experiment_path = Path(full_experiment_path)
    full_experiment_path.mkdir(parents=True, exist_ok=True)

    with open(os.path.join(args.experiments_main_folder, args.experiment_folder, 'commandline_args.txt'), 'w') as f:
        json.dump(args.__dict__, f, indent=2)

    main(args)


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

# 3rd party imports
import logging

import numpy as np

from .e_vxb import e_vxb

# Local imports
from .resample import resample

__author__ = "Louis Richard"
__email__ = "louisr@irfu.se"
__copyright__ = "Copyright 2020-2023"
__license__ = "MIT"
__version__ = "2.4.2"
__status__ = "Prototype"

def vht(e, b, no_ez: bool = False):
    r"""Estimate velocity of the De Hoffman-Teller frame from the velocity
    estimate the electric field eht=-vht x b

    Parameters
    ----------
    e : xarray.DataArray
        Time series of the electric field.
    b : xarray.DataArray
        Time series of the magnetic field.
    no_ez : boolean, Optional
        If True assumed no Ez. Default is False.

    Returns
    -------
    vht : numpy.ndarray
        De Hoffman Teller frame velocity [km/s].
    vht : xarray.DataArray
        Time series of the electric field in the De Hoffman frame.
    dv_ht : numpy.ndarray
        Error of De Hoffman Teller frame.

    """

    n_samples = len(e)

    # Resample magnetic field to electric field sampling (usually higher)

    if n_samples != len(b):
        b = resample(b, e)

    p = np.zeros(6)

    p[0] = np.sum(b[:, 0].data * b[:, 0].data) / n_samples  # Bx*Bx
    p[1] = np.sum(b[:, 0].data * b[:, 1].data) / n_samples  # Bx*By
    p[2] = np.sum(b[:, 0].data * b[:, 2].data) / n_samples  # Bx*Bz
    p[3] = np.sum(b[:, 1].data * b[:, 1].data) / n_samples  # By*By
    p[4] = np.sum(b[:, 1].data * b[:, 2].data) / n_samples  # By*Bz
    p[5] = np.sum(b[:, 2].data * b[:, 2].data) / n_samples  # Bz*Bz

    # assume only Ex and Ey
    if no_ez:
        e[:, 2] *= 0  # put z component to 0 when using only Ex and Ey

        k_mat = np.array(
            [[p[5], 0, -p[2]], [0, p[5], -p[4]], [-p[2], -p[4], p[0] + p[3]]],
        )
    else:
        k_mat = np.array(
            [
                [p[3] + p[5], -p[1], -p[2]],
                [-p[1], p[0] + p[5], -p[4]],
                [-p[2], -p[4], p[0] + p[3]],
            ],
        )

    exb = np.cross(e, b)

    ind_data = np.where(~np.isnan(exb[:, 0].data))[0]

    # revised by Wenya LI; 2015-11-21, wyli @ irfu
    exb_avg = np.sum(exb[ind_data, :], axis=0) / n_samples

    # averExB=sum(ExB(indData).data,1)/nSamples;
    # end revise.
    v_ht = np.linalg.solve(k_mat, exb_avg.T) * 1e3  # 9.12 in ISSI book

    v_ht_hat = v_ht / np.linalg.norm(v_ht, keepdims=True)

    logging.info(
        "v_ht =%(v_mag)7.4f * %(v_vec)s km/s",
        {"v_mag": np.linalg.norm(v_ht), "v_vec": np.array_str(v_ht_hat)},
    )

    # Calculate the goodness of the Hoffman Teller frame
    e_ht = e_vxb(v_ht, b)

    if no_ez:
        e_p, e_ht_p = [e[ind_data], e_ht[ind_data]]
        e_p.data[:, 2], e_ht_p.data[:, 2] = [0, 0]
    else:
        e_p, e_ht_p = [e[ind_data], e_ht[ind_data]]

    delta_e = e_p.data - e_ht_p.data

    poly_fit = np.polyfit(
        e_ht_p.data.reshape([len(e_ht_p) * 3]),
        e_p.data.reshape([len(e_p) * 3]),
        1,
    )
    corr_coeff = np.corrcoef(
        e_ht_p.data.reshape([len(e_ht_p) * 3]),
        e_p.data.reshape([len(e_p) * 3]),
    )

    logging.info(
        "slope = %(slope)6.4f, offs = %(offset)6.4f, cc = %(cc)6.4f",
        {"slope": poly_fit[0], "offset": poly_fit[1], "cc": corr_coeff[0, 1]},
    )

    dv_ht = np.sum(np.sum(delta_e**2)) / len(ind_data)
    s_mat = (dv_ht / (2 * len(ind_data) - 3)) / k_mat
    dv_ht = np.sqrt(np.diag(s_mat)) * 1e3

    dv_ht_hat = dv_ht / np.linalg.norm(dv_ht)

    logging.info(
        "dv_ht =%(dv_mag)7.4f * %(dv_vec)s km/s",
        {"dv_mag": np.linalg.norm(dv_ht), "dv_vec": np.array_str(dv_ht_hat)},
    )

    return v_ht, e_ht, dv_ht


import csv
import flask
import io
import sqlalchemy
from enum import Enum

import ruddock.validation_utils as vu

class PaymentType(Enum):
  CASH = 1
  CHECK = 2
  DEBIT = 3
  ONLINE = 4
  TRANSFER = 5
  OTHER = 6
  INCOME = 7
  PCARD_DEBIT = 8
  ACH_DIRECT_DEPOSIT_REIMBURSEMENT = 9
  CHECK_REIMBURSEMENT = 10

  @classmethod
  def has_value(cls, value):
    try:
      PaymentType(value)
      return True
    except ValueError:
      return False

  @classmethod
  def get_all(cls):
    values = {x.value: x.name.title() for x in PaymentType}
    for key, value in values.items():
      values[key] = value.replace("_", " ").replace("Reimbursement", "(Reimbursement)")
    return values

# ==== SQL QUERIES ====

def select_fyear_info(fyear_num):
  """
  Looks up the record for the given year, returning the record, and a boolean
  indicating if it is the current year.

  If None is given, returns the current year.
  """

  current = get_current_fyear()

  if fyear_num is None:
    return current, True

  fyear_num = int(fyear_num)
  record = [r for r in get_fyears() if r["fyear_num"] == fyear_num][0]

  return record, current["fyear_num"] == fyear_num

def get_current_fyear():
  """
  Looks up the record for the current year.
  If no matching row is found, returns None.
  """

  query = sqlalchemy.text("""
    SELECT fyear_id, fyear_num
    FROM budget_fyears
    WHERE CURDATE() BETWEEN start_date AND end_date
  """)
  rp = flask.g.db.execute(query)

  return rp.first()

def get_expenses():
  """Gets list of all expenses."""
  query = sqlalchemy.text("""
    SELECT expense_id, budget_id, budget_name, fyear_num, date_incurred,
        description, cost, payee_id, payee_name, payment_id
    FROM budget_expenses
      NATURAL JOIN budget_budgets
      NATURAL JOIN budget_fyears
      NATURAL LEFT JOIN budget_payees
    ORDER BY expense_id DESC
    """)

  return flask.g.db.execute(query).fetchall()

def get_expense(expense_id):
  """Gets a particular expense, or None"""
  query = sqlalchemy.text("""
    SELECT expense_id, budget_id, budget_name, fyear_id, fyear_num,
        date_incurred, description, cost, payee_id, payee_name, payment_id
    FROM budget_expenses
      NATURAL JOIN budget_budgets
      NATURAL JOIN budget_fyears
      NATURAL LEFT JOIN budget_payees
    WHERE expense_id = (:e)
    """)

  return flask.g.db.execute(query, e=expense_id).first()

def get_payment(payment_id):
  """Gets a particular payment, or None"""
  query = sqlalchemy.text("""
    SELECT payment_id, account_name, payment_type, amount, date_written,
        date_posted, payee_name, check_no
    FROM budget_payments
      NATURAL JOIN budget_accounts
      NATURAL LEFT JOIN budget_payees
    WHERE payment_id = (:e)
    """)

  return flask.g.db.execute(query, e=payment_id).first()

def get_payments():
  """Gets list of all payments."""
  query = sqlalchemy.text("""
    SELECT payment_id, account_name, payment_type, amount, date_written,
        date_posted, payee_name, check_no
    FROM budget_payments
      NATURAL JOIN budget_accounts
      NATURAL LEFT JOIN budget_payees
    ORDER BY payment_id DESC
    """)

  return flask.g.db.execute(query).fetchall()

def get_transactions():
  """Gets list of all expenses + payments."""
  # We need the USING clause because there's a payee_id in both tables
  # TODO change the column in one of the tables to a different name
  query = sqlalchemy.text("""
    SELECT *
    FROM budget_expenses LEFT JOIN budget_payments USING(payment_id)
      NATURAL JOIN budget_budgets
      NATURAL JOIN budget_fyears
      NATURAL LEFT JOIN budget_accounts
      LEFT JOIN budget_payees ON budget_expenses.payee_id = budget_payees.payee_id
    ORDER BY expense_id DESC
    """)

  return flask.g.db.execute(query).fetchall()

def get_fyears():
  """Gets list of all available fiscal years."""
  query = sqlalchemy.text("""
    SELECT fyear_id, fyear_num
    FROM budget_fyears
    ORDER BY fyear_id DESC
    """)

  return flask.g.db.execute(query).fetchall()

def get_budget_list(fyear_id):
  """Gets list of all budgets in the given year."""
  query = sqlalchemy.text("""
    SELECT budget_id, budget_name
    FROM budget_budgets
      NATURAL JOIN budget_fyears
    WHERE fyear_id = (:f)
    ORDER BY budget_name
    """)

  return flask.g.db.execute(query, f=fyear_id).fetchall()

def get_accounts():
  """Gets list of all accounts."""
  query = sqlalchemy.text("""
    SELECT account_id, account_name
    FROM budget_accounts
    ORDER BY account_name
    """)

  return flask.g.db.execute(query).fetchall()

def get_payees():
  """Gets list of all payees."""
  query = sqlalchemy.text("""
    SELECT payee_id, payee_name
    FROM budget_payees
    ORDER BY payee_name
    """)

  return flask.g.db.execute(query).fetchall()

def get_account_summary():
  """Gets the status of all accounts."""
  query = sqlalchemy.text("""
    SELECT account_name,
      initial_balance - IFNULL(SUM(amount), 0) AS bal
    FROM budget_accounts
      NATURAL LEFT JOIN (
        SELECT account_id, amount
        FROM budget_payments
        WHERE date_posted IS NOT NULL) AS t
    GROUP BY account_id
    ORDER BY account_name
    """)
  # you can't pull the WHERE outside the join, because if no payments
  # have posted, the account won't show up

  return flask.g.db.execute(query).fetchall()

def get_budget_summary(fyear_id):
  """Gets the status of all budgets for the given fiscal year."""
  query = sqlalchemy.text("""
    SELECT *,
      starting_amount - spent AS remaining
    FROM (
      SELECT budget_name, starting_amount, IFNULL(SUM(cost), 0) AS spent
      FROM budget_budgets
        NATURAL JOIN budget_fyears
        NATURAL LEFT JOIN budget_expenses
      WHERE fyear_id = (:f)
      GROUP BY budget_id
      ORDER BY budget_name
    ) AS t
    """)

  return flask.g.db.execute(query, f=fyear_id).fetchall()

def get_unpaid_expenses():
  """Gets all expenses without a corresponding payment."""

  # A list of all unpaid expenses
  list_query = sqlalchemy.text("""
    SELECT expense_id, budget_name, fyear_num, date_incurred, description,
      cost, payee_id
    FROM budget_expenses
      NATURAL JOIN budget_budgets
      NATURAL JOIN budget_fyears
      NATURAL JOIN budget_payees
    WHERE ISNULL(payment_id)
    ORDER BY payee_name, expense_id
  """)

  # A list of payees and the total amount they're owed
  total_query = sqlalchemy.text("""
    SELECT payee_id, payee_name, SUM(cost) AS total
    FROM budget_payees
      NATURAL JOIN budget_expenses
    WHERE ISNULL(payment_id)
    GROUP BY payee_id
    ORDER BY payee_name
  """)

  unpaid_expenses = flask.g.db.execute(list_query).fetchall()
  unpaid_payees = flask.g.db.execute(total_query).fetchall()

  def make_tuple(payee):
    payee_id = payee['payee_id']
    payee_name = payee['payee_name']
    total = payee['total']
    expenses = [r for r in unpaid_expenses if r['payee_id'] == payee_id]

    return payee_id, payee_name, total, expenses

  # We return a list of payees w/ addl data attached
  return list(map(make_tuple, unpaid_payees))

def get_unpaid_amount(payee_id):
  """Returns the amount owed to the given payee, or None if they have no
     outstanding expenses."""
  query = sqlalchemy.text("""
    SELECT SUM(cost) AS total
    FROM budget_expenses
    WHERE payee_id = (:p) AND ISNULL(payment_id)
    GROUP BY NULL
  """)

  result = flask.g.db.execute(query, p=payee_id).first()
  total = result["total"]

  return total

def get_unposted_payments():
  """Returns all payments that haven't been posted."""
  query = sqlalchemy.text("""
    SELECT payment_id, account_name, payment_type, amount, date_written,
      date_posted, payee_name, check_no
    FROM budget_payments
      NATURAL JOIN budget_accounts
      NATURAL JOIN budget_payees
    WHERE ISNULL(date_posted)
    ORDER BY payment_id DESC
  """)

  return flask.g.db.execute(query)

# ==== SQL UPDATES ====

def record_expense(budget_id, date_incurred, description, amount, payment_id,
    payee_id):
  """Inserts a new expense into the database."""

  query = sqlalchemy.text("""
    INSERT INTO budget_expenses
      (budget_id, date_incurred, description, cost, payment_id, payee_id)
    VALUES
      ((:b_id), (:d_inc), (:descr), (:cost), (:p_id), (:payee_id))
  """)

  flask.g.db.execute(
    query,
    b_id=budget_id,
    d_inc=date_incurred,
    descr=description,
    cost=amount,
    p_id=payment_id,
    payee_id=payee_id
  )

def edit_expense(expense_id, budget_id, date_incurred, description, amount, payee_id):
  """
  Changes the details of the given expense.
  """

  query = sqlalchemy.text("""
    UPDATE budget_expenses
    SET
      budget_id = (:budget_id),
      date_incurred = (:date_incurred),
      description = (:description),
      cost = (:amount),
      payee_id = (:payee_id)
    WHERE
      expense_id = (:expense_id)
  """)

  rp = flask.g.db.execute(
    query,
    expense_id=expense_id,
    budget_id=budget_id,
    date_incurred=date_incurred,
    description=description,
    amount=amount,
    payee_id=payee_id
  )

  return rp.rowcount != 0

def delete_expense(expense_id):
  """
  Deletes the given expense.
  """

  query = sqlalchemy.text("""
    DELETE FROM budget_expenses
    WHERE expense_id = (:expense_id)
  """)

  rp = flask.g.db.execute(
    query,
    expense_id=expense_id,
  )

  return rp.rowcount != 0

def edit_payment(payment_id, amount, date_written, payee_id, payment_type):
  """
  Changes the details of the given payment.
  """

  query = sqlalchemy.text("""
    UPDATE budget_payments
    SET
      date_written = (:date_written),
      amount = (:amount),
      payee_id = (:payee_id),
      payment_type = (:payment_type)
    WHERE
      payment_id = (:payment_id)
  """)

  rp = flask.g.db.execute(
    query,
    date_written=date_written,
    amount=amount,
    payee_id=payee_id,
    payment_id=payment_id,
    payment_type=payment_type
  )

  return rp.rowcount != 0

def delete_payment(payment_id):
  """
  Deletes the given payment.
  """

  query = sqlalchemy.text("""
    DELETE FROM budget_payments
    WHERE payment_id = (:payment_id)
  """)

  rp = flask.g.db.execute(
    query,
    payment_id=payment_id,
  )

  return rp.rowcount != 0

def record_payment(account_id, payment_type, amount, date_written, date_posted,
    payee_id, check_no):
  """Inserts a new payment into the database, returning its ID."""

  query = sqlalchemy.text("""
    INSERT INTO budget_payments
      (account_id, payment_type, amount, date_written, date_posted, payee_id,
      check_no)
    VALUES
      ((:a_id), (:t), (:amount), (:d_writ), (:d_post), (:payee_id), (:check_no))
  """)

  result = flask.g.db.execute(
    query,
    a_id=account_id,
    t=payment_type,
    amount=amount,
    d_writ=date_written,
    d_post=date_posted,
    payee_id=payee_id,
    check_no=check_no
  )

  return result.lastrowid

def record_new_payee(payee_name):
  """Inserts a new payee into the database, returning its ID."""
  query = sqlalchemy.text("""
    INSERT INTO budget_payees (payee_name)
    VALUES (:p)
  """)

  result = flask.g.db.execute(query, p=payee_name)
  return result.lastrowid

def mark_as_paid(payee_id, payment_id):
  """
  Assigns the given payment id to all unpaid expenses from the given payee.
  """
  query = sqlalchemy.text("""
    UPDATE budget_expenses
    SET payment_id = (:payment)
    WHERE ISNULL(payment_id) AND payee_id = (:payee)
  """)

  flask.g.db.execute(query, payment=payment_id, payee=payee_id)

def post_payment(payment_id, date_posted):
  """Marks the given payment as posted, with the given date."""
  query = sqlalchemy.text("""
    UPDATE budget_payments
    SET date_posted = (:dp)
    WHERE payment_id = (:pi)
  """)

  flask.g.db.execute(query, dp=date_posted, pi=payment_id)

def void_payment(payment_id):
  """Marks the given payment as posted, with the given date."""
  transaction = flask.g.db.begin()
  try:
    # Wipe payment ID from expenses
    query = sqlalchemy.text("""
      UPDATE budget_expenses
      SET payment_id = NULL
      WHERE payment_id = (:pi)
    """)

    # Delete payment
    query2 = sqlalchemy.text("""
      DELETE FROM budget_payments
      WHERE payment_id = (:pi)
    """)

    flask.g.db.execute(query, pi=payment_id)
    flask.g.db.execute(query2, pi=payment_id)
    transaction.commit()

  except Exception:
    transaction.rollback()
    flask.flash("An unexpected error occurred. Please find an IMSS rep.")

# ==== VALIDATION ====

def test_predicates(triplets, flash=True):
  """Makes validation methods more readable."""
  errs = [err for pred, guard, err in triplets if guard and not pred]

  if flash:
    for e in errs:
      flask.flash(e)

  return errs

def validate_expense(budget_id, date_incurred, amount, description):
  valid = True
  valid &= vu.validate_integer(budget_id, flash_errors=False)
  valid &= vu.validate_date(date_incurred, flash_errors=False)
  valid &= vu.validate_currency(amount, flash_errors=False)
  valid &= len(description) > 0

  return valid

def validate_payment(payment_type, account_id, check_no):
  valid = True
  valid &= vu.validate_integer(payment_type, flash_errors=False) and \
      PaymentType.has_value(int(payment_type))
  valid &= vu.validate_integer(account_id, flash_errors=False)
  if payment_type == PaymentType.CHECK:
    valid &= vu.validate_integer(check_no, flash_errors=False)

  return valid

def validate_payee(payee_id, payee_name):
  # TODO also check that the new payee isn't actually an old one...
  # This one's a bit unusual, we want a valid old payee XOR a valid new one
  existing_payee = vu.validate_integer(payee_id, flash_errors=False)
  new_payee = len(payee_name) > 0
  return (existing_payee != new_payee)

def download_expenses():
  expenses = get_transactions()
  ptypes = PaymentType.get_all()
  fields = ["expense_id", "budget_name", "fyear_num", "date_incurred",
            "description", "cost", "payment_id", "payee_name", "payment_type",
            "account_name", "check_no"]
  # If the first title is "ID" there is a formatting issue with Excel
  titles = ["Expense ID", "Budget", "FY", "Date Incurred", "Description",
            "Amount", "Payment ID", "Payee", "Type", "Account", "Check"]

  # References: stackoverflow.com/a/45111660 and stackoverflow.com/a/14614920
  proxy = io.StringIO()
  wr = csv.DictWriter(proxy, fields)
  wr.writerow(dict(zip(fields, titles)))
  for db_row in expenses:
    csv_row = {key: db_row[key] for key in fields}
    if csv_row["payment_type"] is not None:
      csv_row["payment_type"] = ptypes[csv_row["payment_type"]]
    wr.writerow(csv_row)

  download = io.BytesIO()
  download.write(proxy.getvalue().encode())
  download.seek(0)
  proxy.close()

  return flask.send_file(download, as_attachment=True,
                         attachment_filename="expenses.csv")

def download_summaries(fyear_id):
  query = sqlalchemy.text("""
    SELECT account_id, account_name, initial_balance,
      initial_balance - IFNULL(SUM(amount), 0) AS bal
    FROM budget_accounts
      NATURAL LEFT JOIN (
        SELECT account_id, amount
        FROM budget_payments
        WHERE date_posted IS NOT NULL) AS t
    GROUP BY account_id
    ORDER BY account_name
    """)

  accounts = flask.g.db.execute(query).fetchall()

  budgets = get_budget_summary(fyear_id)
  account_fields = ["account_id", "account_name", "initial_balance", "bal"]
  account_titles = ["Account ID", "Account", "Initial Balance", "Current Balance"]
  budget_fields = ["budget_name", "starting_amount", "spent", "remaining"]
  budget_titles = ["Budget", "Starting Amount", "Spent", "Remaining"]

  proxy = io.StringIO()

  wr = csv.DictWriter(proxy, account_fields)
  wr.writerow(dict(zip(account_fields, account_titles)))
  for db_row in accounts:
    csv_row = {key: db_row[key] for key in account_fields}
    wr.writerow(csv_row)

  wr.writerow({})

  wr = csv.DictWriter(proxy, budget_fields)
  wr.writerow(dict(zip(budget_fields, budget_titles)))
  for db_row in budgets:
    csv_row = {key: db_row[key] for key in budget_fields}
    wr.writerow(csv_row)

  download = io.BytesIO()
  download.write(proxy.getvalue().encode())
  download.seek(0)
  proxy.close()
  return flask.send_file(download, as_attachment=True, attachment_filename="budget_summaries.csv")


import pytest
from rest_framework.request import Request
from rest_framework.test import APIRequestFactory

from epic_app.models.epic_questions import (
    EvolutionQuestion,
    LinkagesQuestion,
    NationalFrameworkQuestion,
)
from epic_app.serializers.question_serializer import (
    EvolutionQuestionSerializer,
    LinkagesQuestionSerializer,
    NationalFrameworkQuestionSerializer,
)
from epic_app.tests.epic_db_fixture import epic_test_db

@pytest.fixture(autouse=True)
def question_serializer_fixture(
    epic_test_db: pytest.fixture,
):
    """
    Dummy fixture just to load a default db from dummy_db.

    Args:
        epic_test_db (pytest.fixture): Fixture to load for the whole file tests.
    """
    pass

def get_serializer():
    factory = APIRequestFactory()
    request = factory.get("/")

    return {
        "request": Request(request),
    }

serializer_context = get_serializer()

@pytest.mark.django_db
class TestNationalFrameworkSerializer:
    def test_given_valid_instances_when_to_representation_returns_expected_data(self):
        def validate_fields(dict_item: dict) -> bool:
            return "description" in list(dict_item.keys())

        serialized_data = list(
            NationalFrameworkQuestionSerializer(
                NationalFrameworkQuestion.objects.all(),
                many=True,
                context=serializer_context,
            ).data
        )

        assert len(serialized_data) == 2
        assert all(map(validate_fields, serialized_data))

@pytest.mark.django_db
class TestEvolutionQuestionSerializer:
    def test_given_valid_instances_when_to_representation_returns_expected_data(self):
        def validate_fields(dict_item: dict) -> bool:
            valid_fields = [
                "nascent_description",
                "engaged_description",
                "capable_description",
                "effective_description",
            ]
            return all(v_field in list(dict_item.keys()) for v_field in valid_fields)

        serialized_data = list(
            EvolutionQuestionSerializer(
                EvolutionQuestion.objects.all(), many=True, context=serializer_context
            ).data
        )

        assert len(serialized_data) == 2
        assert all(map(validate_fields, serialized_data))

@pytest.mark.django_db
class TestLinkagesQuestionSerializer:
    def test_given_valid_instances_when_to_representation_returns_expected_data(self):
        serialized_data = list(
            LinkagesQuestionSerializer(
                LinkagesQuestion.objects.all(), many=True, context=serializer_context
            ).data
        )

        assert len(serialized_data) == 1