import tensorflow as tf
import numpy as np
from typing import Union, Optional, Sequence
from pathlib import Path

from dui.datasets.hdf5datasetfactory import HDF5DatasetFactory
from dui.utils.signal import compress_db
from dui.layers.utils import get_channel_axis

def create_image_dataset(
        path: Union[str, Path],
        name: str,
        factor: Union[str, float] = '0db',
        # TODO: None as default or 1?
        signal_type: str = 'rf',
        # TODO: None or 'raw' as default?
        data_format: str = 'channels_last',
        # TODO: patch paddings typing elsewhere if validated
        # paddings: Optional[Union[Sequence[int], np.ndarray]] = None,
        paddings: Optional[Union[Sequence[Sequence[int]], np.ndarray]] = None,
        start: int = 0,
        stop: Optional[int] = None,
        step: int = 1,
        slicer: Optional[Sequence[slice]] = None,
        batch_size: int = 1,
        shuffle: bool = False,
        num_parallel_calls: Optional[int] = None,
        seed: Optional[int] = None,
) -> tf.data.Dataset:

    # Factory
    dataset_factory = HDF5DatasetFactory(
        path=path,
        name=name,
        start=start,
        stop=stop,
        step=step,
        slicer=slicer,
        shuffle=shuffle,
        seed=seed
    )

    # Check sample shape
    base_sample_shape = dataset_factory._output_sample_shape
    if len(base_sample_shape) != 2:
        raise ValueError(
            "Dataset sample must be a 2D array. Current shape: {}".format(
                base_sample_shape
            )
        )

    # Normalization factor
    if isinstance(factor, str):
        attr_key = factor
        factor = dataset_factory.attrs.get(attr_key)
        if factor is None:
            raise ValueError(
                "No attribute '{}' for dataset '{}' in '{}'".format(
                    attr_key, dataset_factory._dataset.name,
                    dataset_factory._dataset.file.filename
                )
            )
    elif type(factor) in (int, float):
        pass
    else:
        raise TypeError("Unsupported type for 'factor'")

    # Create dataset
    dataset = dataset_factory.create_dataset()
    # TODO: include factor directly and specialize the pre-processing
    #  for US-specific only?

    # Hack to avoid having an <unknown> shape (probably unsafe)
    # TODO: handle this in factory or by sub-classing tf.data.Dataset
    #  Note: Probably below some Dataset._element_structure properties
    #  Note: most probably not compatible with 1.15
    dataset._element_structure._shape = tf.TensorShape(base_sample_shape)

    # Pre-processing
    dataset = dataset.batch(batch_size=batch_size)
    # TODO: use `dataset.padded_batch` instead and remove following
    #  `paddings` option from following pre-processing
    # TODO: apply normalization factor before
    dataset = _preprocess_image_dataset(
        dataset=dataset,
        factor=factor,
        data_format=data_format,
        signal_type=signal_type,
        paddings=paddings,
        num_parallel_calls=num_parallel_calls
    )

    return dataset

def _preprocess_image_dataset(
        dataset: tf.data.Dataset,
        factor: Optional[float] = None,
        data_format: str = 'channels_last',
        signal_type: Optional[str] = None,
        paddings: Optional[Union[Sequence[int], np.ndarray]] = None,
        num_parallel_calls: Optional[int] = None
) -> tf.data.Dataset:

    # Specify pre-processing function as a mapping function
    def map_func(x: tf.Tensor) -> tf.Tensor:
        return _image_preproc_fun(
            x,
            factor=factor,
            data_format=data_format,
            signal_type=signal_type,
            paddings=paddings
        )

    return dataset.map(
        map_func=map_func,
        num_parallel_calls=num_parallel_calls
    )

def _image_preproc_fun(
        x: tf.Tensor,
        factor: Optional[float] = None,
        data_format: str = 'channels_last',
        signal_type: Optional[str] = None,
        paddings: Optional[Union[Sequence[int], np.ndarray]] = None,
) -> tf.Tensor:

    # TODO: check inputs
    x = tf.convert_to_tensor(x)

    # Normalization factor
    if factor:
        # TODO: apply factor before and keep this pre-proc
        #  function only for US-specific transformations?
        x /= factor

    # Paddings
    if paddings is not None:
        # TODO: would probably make more sense to remove paddings
        #  from this US-specific pre-processing function
        # x = _batched_pad(x, paddings=paddings)
        paddings = np.array(paddings)
        valid_pad_shape = 2, 2
        pad_shape = paddings.shape
        # TODO: this test is too restrictive in general (e.g. 3D)
        #  but ok for now as we only work on 2D images
        if pad_shape != valid_pad_shape:
            raise ValueError(
                "Incompatible 'paddings' shape. Current: {}. "
                "Expected {}".format(pad_shape, valid_pad_shape)
            )
        paddings = [[0, 0], *paddings.tolist()]
        pad_kwargs = {
            'paddings': tf.constant(paddings, dtype='int32'),
            'mode': 'CONSTANT',
            'constant_values': 0
        }
        x = tf.pad(x, **pad_kwargs)

    # Channel axis
    channel_axis = get_channel_axis(data_format=data_format)

    # Signal type
    if signal_type is not None:
        if signal_type == 'rf':
            x = tf.math.real(x)
        elif signal_type == 'iq':
            # Stack complex components in channels
            x = tf.stack((tf.math.real(x), tf.math.imag(x)), axis=channel_axis)
        elif signal_type == 'env':
            # Takes modulus of complex IQ signal
            x = tf.math.abs(x)
        elif signal_type == 'bm':
            # Takes modulus of complex IQ signal
            x = tf.math.abs(x)
            # Compress to dB
            x = compress_db(tensor=x)
        elif signal_type == 'raw':
            pass
        else:
            raise ValueError("Invalid signal type")

    # Expand dimension
    if signal_type != 'iq':
        x = tf.expand_dims(x, axis=channel_axis)

    return x


import logging
import json
import gzip
import ipaddress
import datetime

from c99api import EndpointClient
from typing import List, Dict, Optional
from os.path import exists

from pydantic import BaseModel

logger = logging.getLogger()

def enrich_object_c99(object, c99_key:str=""):
    c99 = EndpointClient
    c99.key = c99_key
    ip = object["IPAddress"]
    resp = c99.gethostname(ip)
    if resp["success"] and ip != resp["hostname"] and resp["hostname"] not in object["hostname_list"]:
        logging.info(f"gethostname: {resp['hostname']}")
        object["hostname_list"].append(resp["hostname"])
    resp = c99.ip2domains(ip)
    if resp["success"] and resp["count"] >= 1:
        logging.info(f"ip2domains: {resp['data']}")
        object["domain_list"].extend([hname for hname in resp["data"] if hname not in object["domain_list"]])

def merge_config(current_config: Dict[int, str] = {}, custom_config: Dict[int, str] = {}):
    for key, value in custom_config.items():
        if key in current_config.keys():
            if isinstance(value, (list,)):
                current_config[key] = list(set(current_config[key].extend(current_config[key])))
            elif isinstance(value, (dict,)):
                current_config[key] = merge_config(current_config[key], custom_config[key])
            else:
                current_config[key] = value
        else:
            current_config.update({key: value})
    return current_config

def load_config(default_config: str = "config.default.json", override_config: str = ""):
    config_builder = {}
    if exists(default_config):
        with open(default_config, "r", encoding="utf-8") as config_file:
            config_builder = json.load(config_file)
    else:
        raise ValueError("config file not found")
    if exists(override_config):
        with open(override_config, "r", encoding="utf-8") as config_file:
            try:
                configData = json.load(config_file)
                config_builder = merge_config(current_config=config_builder, custom_config=configData)
            except Exception as e:
                logger.error(f"Error adding override config\n{e}")
    return config_builder

def decode_shodan(obj:dict={}):
    try:
        parsed_object = {
            "domain_list": obj["domains"] if "domains" in obj else [],
            "hostname_list": [obj["_shodan"]["options"]["hostname"]] if "hostname" in obj["_shodan"]["options"] else [],
            "cloud_provider": None,
            "operating_system": obj["os"],
            "product": obj["product"] if "product" in obj else "",
            "IPAddress": ipaddress.ip_address(obj["ip_str"]),
            "timestamp": datetime.datetime.fromisoformat(obj["timestamp"]),
            "protocol": obj["transport"] if "transport" in obj else "",
            "internet_service_provider": obj["isp"],
            "version": obj["version"] if "version" in obj else "",
            "organisation": obj["org"],
            "country": obj["location"]["country_name"] if "country_name" in obj["location"] else "",
            "city": obj["location"]["city"] if "city" in obj["location"] else "",
            "port": obj["port"]
        }
        parsed_object["hostname_list"].extend([hname.strip() for hname in obj["hostnames"]])
    except Exception as e:
        logging.error(e)
        return {}
    try:
        if "ssl" in obj and "cert" in obj["ssl"]:
            cert = obj["ssl"]
            #parsed_object["ssl_fingerprint"] = cert["cert"]["fingerprint"]["sha256"]
            #parsed_object["ssl_serial"] = cert["cert"]["serial"]
            parsed_object["ssl_SAN"] = [cert["cert"]["subject"]["CN"]] if "CN" in cert["cert"]["subject"]["CN"] else []
            for alt in cert["cert"]["extensions"]:
                if alt["name"]=="subjectAltName" and alt["data"]:
                    i = 0
                    while i < len(alt["data"]):
                        if alt["data"][i] == "\\":
                            i += 4
                            continue
                        next_slash = alt["data"][i:].find("\\")
                        if next_slash >= 0:
                            parsed_object["ssl_SAN"].append(alt["data"][i:i+next_slash])
                            i += next_slash
                        else:
                            parsed_object["ssl_SAN"].append(alt["data"][i:])
                            i = len(alt["data"])
                        if parsed_object["ssl_SAN"][-1] == "0.":
                            parsed_object["ssl_SAN"].pop()
            parsed_object["ssl_SAN"] = list(set(parsed_object["ssl_SAN"]))
            parsed_object["ssl_issuer"] = cert["cert"]["issuer"]["O"] if "O" in cert["cert"]["issuer"] else cert["cert"]["issuer"]["CN"]
            #parsed_object["ssl_ja3"] = cert["ja3s"]
            #parsed_object["ssl_jarm"] = cert["jarm"]
            parsed_object["ssl_expiration"] = datetime.datetime.strptime(cert["cert"]["expires"], "%Y%m%d%H%M%SZ")
        else:
            #parsed_object["ssl_fingerprint"] = ""
            #parsed_object["ssl_serial"] = -1
            parsed_object["ssl_SAN"] = []
            parsed_object["ssl_issuer"] = ""
            #parsed_object["ssl_ja3"] = ""
            #parsed_object["ssl_jarm"] = ""
            parsed_object["ssl_expiration"] = datetime.datetime.fromordinal(1)
    except Exception as e:
        #parsed_object["ssl_fingerprint"] = ""
        #parsed_object["ssl_serial"] = -1
        parsed_object["ssl_SAN"] = []
        parsed_object["ssl_issuer"] = ""
        #parsed_object["ssl_ja3"] = ""
        #parsed_object["ssl_jarm"] = ""
        parsed_object["ssl_expiration"] = datetime.datetime.fromordinal(1)
        logging.error(e)
    return parsed_object

def load_shodan_files(filename:str="", config:Dict={}):
    if not exists(filename):
        logging.error(f"File not found: {filename}")
        raise FileNotFoundError
    logging.info(f"Loading file: {filename}")
    if filename.endswith(".json.gz"):
        with gzip.open(filename, "rb") as archive:
            lines = archive.readlines()
    else:
        with open(filename, "rb") as raw_file:
            lines = raw_file.readlines()
    data = []
    error_count = 0
    for line in lines:
        try:
            json_obj = json.loads(line)
            try:
                obj = decode_shodan(obj=json_obj)
                data.append(obj)
            except Exception as e:
                logger.warning(f"JSON data could not be parsed")
                logger.warning(e)
        except:
            error_count += 1
            continue
    if error_count > 0:
        logging.error(f"{filename} - Errors occurred during loading of data: {error_count}")
    return data

if __name__ == "__main__":
    configFileName = "../../config/config.default.json"
    loaded_config = load_config(default_config=configFileName)
    logConfig = loaded_config["logging"]
    logging.basicConfig(
        level=logConfig["log_level"],
        format=logConfig["log_fstr_std"],
        datefmt=logConfig["log_date_formt"]
    )
    pass


import csv
import matplotlib.pyplot as plt

Zb = [[], []]
# with open('F:/zhengwangwork/test csv/4.csv','rb')as f:
#     reader=csv.reader(f)
#     for row in reader:
#         print(row[0])
file = open('../gold.csv', 'r', encoding='UTF-8')  # 打开csv文件
reader = csv.reader(file)  # 读取csv文件
data = list(reader)  # 将csv数据转化为列表
length_h = len(data)  # 得到数据行数
lenght_l = len(data[0])  # 得到每行长度

x = list()
y = list()

for i in range(0, length_h):  # 从第一行开始读取
    x.append(data[i][0])  # 将第一列数据从第一行读取到最后一行付给列表x
    y.append(data[i][2])  # 将第三列数据从第一行读取到最后一行付给列表y
plt.plot(x, y)  # 绘制折线图
plt.show()  # 显示折线图


import random
import numpy as np
from scipy.stats import entropy

class CNF:
    def __init__(self, path=None, clauses=None):
        if path:
            with open(path, 'r') as cnf:
                formula = cnf.read()
                formula = formula.split('\n')
                start_index = 0
                while formula[start_index][0] != 'p':
                    start_index += 1
                self.n = int(formula[start_index].split()[2])
                self.variables = [i+1 for i in range(self.n)]
                self.m = int(formula[start_index].split()[3])
                self.clauses = [list(map(int, formula[start_index + 1 + i].split()[:-1])) for i in range(self.m)]
        else:
            variables = set()
            for clause in clauses:
                for var in clause:
                    variables.add(abs(var))
            self.n = len(variables)
            self.variables = list(variables)
            self.m = len(clauses)
            self.clauses = clauses

    def get_size(self):
        return [self.m, self.n, float(self.m) / self.n]

    def get_vc(self):
        nodes = {i: set() for i in self.variables}
        for j in range(self.m):
            for var in self.clauses[j]:
                nodes[abs(var)].add(j)
        nodes = [len(nodes.get(i)) for i in nodes]
        nodes_np = np.array(nodes)
        nodes_proba = np.unique(nodes_np, return_counts=True)[1]/float(len(nodes_np))
        nodes_entropy = entropy(list(nodes_proba))
        clause = []
        for j in range(self.m):
            cl = set()
            for i in range(len(self.clauses[j])):
                cl.add(abs(self.clauses[j][i]))
            clause.append(len(cl))
        clause_np = np.array(clause)
        clause_proba = np.unique(clause_np, return_counts=True)[1]/float(len(clause_np))
        clause_entropy = entropy(list(clause_proba))
        return [nodes_np.mean(), nodes_np.std()/nodes_np.mean(), nodes_np.min(), nodes_np.max(), nodes_entropy,
                clause_np.mean(), clause_np.std()/clause_np.mean(), clause_np.min(), clause_np.max(), clause_entropy]

    def get_v(self):
        variables = {i: set() for i in self.variables}
        for j in range(self.m):
            for var in self.clauses[j]:
                for var_o in self.clauses[j]:
                    if abs(var_o) != abs(var):
                        variables[abs(var)].add(abs(var_o))
        var_deg = [len(variables.get(i)) for i in variables]
        var_deg_np = np.array(var_deg)
        return [var_deg_np.mean(), var_deg_np.std()/var_deg_np.mean(), var_deg_np.min(), var_deg_np.max()]

    def get_balance(self):
        ratio_clause = []
        for clause in self.clauses:
            pos, neg = 0, 0
            for var in clause:
                if var > 0:
                    pos += 1
                else:
                    neg += 1
            ratio_clause.append(float(pos) / (pos + neg))
        ratio_clause_np = np.array(ratio_clause)
        ratio_clause_proba = np.unique(ratio_clause_np, return_counts=True)[1] / float(len(ratio_clause_np))
        ratio_clause_entropy = entropy(list(ratio_clause_proba))
        ration_var = {i: [0, 0] for i in self.variables}
        for j in range(self.m):
            for var in self.clauses[j]:
                if var > 0:
                    ration_var.get(abs(var))[0] += 1
                else:
                    ration_var.get(abs(var))[1] += 1
        ration_var = [float(ration_var.get(i)[0]) / (ration_var.get(i)[0] + ration_var.get(i)[1]) for i in ration_var]
        ration_var_np = np.array(ration_var)
        ration_var_proba = np.unique(ration_var_np, return_counts=True)[1] / float(len(ration_var_np))
        ration_var_entropy = entropy(list(ration_var_proba))
        binary, ternary = 0, 0
        for clause in self.clauses:
            if len(clause) == 2:
                binary += 1
            elif len(clause) == 3:
                ternary += 1
        return [ratio_clause_np.mean(), ratio_clause_np.std()/ratio_clause_np.mean(), ratio_clause_entropy,
                ration_var_np.mean(), ration_var_np.std()/ration_var_np.mean(), ration_var_np.min(),
                ration_var_np.max(), ration_var_entropy, float(binary)/self.m, float(ternary)/self.m]

    def get_horn(self):
        num_of_horns = 0
        horn_var = {i: 0 for i in self.variables}
        for clause in self.clauses:
            horn = True
            cnt = 0
            for var in clause:
                if var > 0:
                    cnt += 1
                if cnt > 1:
                    horn = False
                    break
            if horn:
                num_of_horns += 1
                for vr in clause:
                    horn_var[abs(vr)] += 1
        horn_var = [horn_var.get(i) for i in horn_var]
        horn_var_np = np.array(horn_var)
        horn_var_proba = np.unique(horn_var_np, return_counts=True)[1] / float(len(horn_var_np))
        horn_var_entropy = entropy(list(horn_var_proba))
        return [float(num_of_horns) / self.m, horn_var_np.mean(), horn_var_np.std()/horn_var_np.mean(),
                horn_var_np.min(), horn_var_np.max(), horn_var_entropy]

    def get_features(self):
        size = self.get_size()
        vc = self.get_vc()
        v = self.get_v()
        balance = self.get_balance()
        horn = self.get_horn()
        return size + vc + v + balance + horn

    def set_var(self, var=None):
        if not var:
            var = random.choice(self.variables + [-i for i in self.variables])
        new_clauses = [[i for i in clause if i != -var] for clause in self.clauses if var not in clause]
        return CNF(clauses=new_clauses)


import argparse
from math import sqrt, floor, ceil
from os import listdir
import io

from Player import Player
import pygame
import time
from View import View
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
import cv2

BUTTONPRESS_CSV = "buttonpress.csv"
GAME_CSV = "game.csv"
FACE_FILES = ["faces_0.txt", "faces_1.txt"]

class DummyClock:
    def __init__(self):
        pass

    def get_fps(self):
        return 60

    def get_rawtime(self):
        return 1

    def get_time(self):
        return 1

class GameReplay:
    def __init__(self, replay_dir):
        self.field = None
        self.players = []
        self.boxes = set()
        self.bombs = dict()
        self.explosions = set()
        self.inactive_traps = dict()
        self.active_traps = set()
        self.falling_boxes = dict()
        self.crushing_boxes = dict()
        self.power_ups = dict()
        self.active_taunts = dict()
        self.view = View(32 * 15, 32 * 16, "Replay Bombangerman")
        self.plots = []
        self.vlines = []
        self.replay_dir = "../replays/" + replay_dir
        self.face_replay_file_handles = [None,None]
        self.next_faces = [[0,np.zeros((48,48))],[0,np.zeros((48,48))]]

        plt.ion()
        plt.show()
        self.figure = plt.figure()
        self.figure.autofmt_xdate()
        plt.xticks([])
        plt.yticks([])

        self.last_timestamp = None
        self.first_timestamp = None
        replay_files = listdir(self.replay_dir)

        if GAME_CSV in replay_files:
            with open(self.replay_dir + "/" + GAME_CSV) as f:
                lines = f.readlines()
                self.first_timestamp = float(lines[3].split(";")[0])
                self.last_timestamp = float(lines[-1].split(";")[0])

        # plot data files
        data_files = [f for f in replay_files if f not in [GAME_CSV, BUTTONPRESS_CSV] + FACE_FILES]
        self.nr_plots = len(data_files)
        for file in data_files:
            print("FILENAME:", file)
            self.replay_data(file)

        # Faces Replay display setup
        # Yes, this has constantly open file handles. but this is read only, so we will prolly get away with it.
        # Here, we preload the first entry in those files
        for i,filename in enumerate(FACE_FILES):
            if filename in replay_files:
                f = open(self.replay_dir + "/" + filename)
                self.face_replay_file_handles[i] = f
                print("Opened Face Data File for Player", i)
                self.load_next_image_and_timestamp(i,f)
        # prepare 2 windows if 2 players have replay data here
        for i,h in enumerate(self.face_replay_file_handles):
            if h is not None:
                cv2.namedWindow("Player " + str(i))

        # buttonpress
        if BUTTONPRESS_CSV in replay_files:
            try:
                with open(self.replay_dir + "/" + BUTTONPRESS_CSV) as f:
                    content = f.readlines()

                bps = [float(x) for x in content]

                for b in bps:
                    for plot in self.plots:
                        plot.axvline(x=b, c="b")
            except FileNotFoundError:
                print(BUTTONPRESS_CSV + " not found")

        # game replay
        if GAME_CSV in replay_files:
            with open(self.replay_dir + "/" + GAME_CSV) as f:
                for i, line in enumerate(f):
                    if i == 0:
                        self.field = eval(line)
                    elif i == 1:
                        for x, y, ticks in eval(line):
                            self.inactive_traps[(x, y)] = [ticks, ticks]
                    elif i == 2:
                        player_data = eval(line)
                        self.update_player_data(player_data)
                    else:
                        break

    def replay_data(self, filename):
        content = pd.read_csv(self.replay_dir + "/" + filename, delimiter=";", names=["time", "player0", "player1"], dtype="float")
        if self.last_timestamp is not None:
            content = content[content["time"] <= self.last_timestamp]
        if self.first_timestamp is not None:
            content = content[content["time"] >= self.first_timestamp]
        content = content.fillna(method='ffill').dropna()
        content = content.sort_values('time', axis=0)

        plot_nr = len(self.plots) + 1

        nr_rows = ceil(sqrt(self.nr_plots))
        nr_columns = ceil(sqrt(self.nr_plots))

        plot = self.figure.add_subplot(nr_rows, nr_columns, plot_nr)
        plot.title.set_text(filename)
        self.plots.append(plot)
        content.plot(x="time", ax=plot)
        vline = plot.axvline(x=0, c="r")
        self.vlines.append(vline)

    def new_plot(self, x, ys, title):
        plot_nr = len(self.plots) + 1

        nr_rows = floor(sqrt(self.nr_plots))
        nr_columns = ceil(sqrt(self.nr_plots))

        plot = self.figure.add_subplot(nr_rows, nr_columns, plot_nr)
        plot.title.set_text(title)
        self.plots.append(plot)
        for y in ys:
            plot.plot(x, y)
        vline = plot.axvline(x=x[0], c="r")
        self.vlines.append(vline)

    def replay(self):
        clock = DummyClock()
        run = True
        frame = 3
        last_timestamp = None
        while run:

            events = pygame.event.get()
            for event in events:
                if event.type == pygame.QUIT:
                    run = False
                if event.type == pygame.K_ESCAPE:
                    run = False

            with open(self.replay_dir + "/" + GAME_CSV) as f:
                for i, line in enumerate(f):
                    if i == frame:
                        if line == "\n": # no idea why but needed in windows
                            continue
                        timestamp = eval(line.split(";")[0])

                        events = eval(line.split(";")[1])
                        self.handle_events(events)
                    elif i < frame:
                        continue
                    else:
                        break

            if frame % 30 == 0:
                for vline in self.vlines:
                    vline.set_xdata(timestamp)
                # TODO what's this line for?
                plt.pause(1e-10)

            if not last_timestamp is None:
                time.sleep((timestamp - last_timestamp)/5)

            last_timestamp = timestamp

            # intermission: draw player face when present
            for i,h in enumerate(self.face_replay_file_handles):
                if h is not None:
                    if timestamp >= self.next_faces[i][0]:
                        img = self.next_faces[i][1] / 255.0
                        img = cv2.resize(img, dsize=(48*3, 48*3), interpolation=cv2.INTER_NEAREST)
                        cv2.imshow("Player " + str(i), img)
                        self.load_next_image_and_timestamp(i,h)
            # intermission end

            self.update_counters()
            self.view.draw_game(self.field, self.boxes, self.inactive_traps, self.active_traps, self.power_ups,
                                self.bombs, self.explosions, self.falling_boxes, self.crushing_boxes,
                                self.players, self.active_taunts, 0, clock)
            self.view.update()
            frame += 1

        pygame.quit()

    def update_counters(self):
        """ Updates all tick counters for the client
        """
        for bomb in self.bombs:
            self.bombs[bomb] -= 1

        for x, y in self.inactive_traps:
            data = self.inactive_traps[(x, y)]
            data[1] -= 1  # [max_ticks, remaining_ticks]

        for falling_box in self.falling_boxes:
            self.falling_boxes[falling_box] -= 1

        for crushing_box in self.crushing_boxes:
            self.crushing_boxes[crushing_box] -= 1

        remove = []
        for id, ticks in self.active_taunts.items():
            if ticks <= 0:
                remove.append(id)
            else:
                self.active_taunts[id] -= 1
        for id in remove:
            del self.active_taunts[id]

    def update_player_data(self, player_data):
        while len(self.players) < len(player_data):
            self.players.append(None)
        for d in player_data:
            pid = d["id"]
            x = d["x"]
            y = d["y"]
            lifes = d["l"]
            bombs = d["b"]
            power = d["p"]
            player = self.players[pid] if len(self.players) > pid else None
            if player == None:
                self.players[pid] = Player(pid, x, y, lifes, power, bombs)
            else:
                player.x = x
                player.y = y
                player.lifes = lifes
                player.bombs = bombs
                player.power = power

    def handle_events(self, events):
        for type, data in events:
            if type == 0:
                # GENERIC
                pass
            elif type == 1:
                # PLAYER_INIT
                # unused
                pass
            elif type == 2:
                # PLAYER_MORTAL
                self.players[data["id"]].immortal = False
            elif type == 3:
                # PLAYER_DAMAGED
                self.players[data["id"]].lifes -= data["dmg"]
                self.players[data["id"]].immortal = True
            elif type == 4:
                # PLAYER_MOVED
                p = self.players[data["id"]]
                p.x = data["x"]
                p.y = data["y"]
                p.facing = data["f"]
            elif type == 12:
                # PLAYER_NOT_SLIMEY
                self.players[data["id"]].slimey = False
            elif type == 13:
                # PLAYER_SLIMED
                self.players[data["id"]].slimey = True
            elif type == 5:
                # SPAWN_BOX
                self.boxes.add((data["x"], data["y"]))
            elif type == 6:
                # SPAWN_BOMB
                self.bombs[(data["x"], data["y"])] = data["t"]  # x,y -> ticks
            elif type == 7:
                # SPAWN_EXPLOSION
                self.explosions.add((data["x"], data["y"]))
            elif type == 8:
                # UPDATE_TRAP
                # unused
                pass
            elif type == 9:
                # REMOVE_BOX
                self.boxes.discard((data["x"], data["y"]))
            elif type == 10:
                # REMOVE_BOMB
                self.bombs.pop((data["x"], data["y"]), None)
            elif type == 11:
                # REMOVE_EXPLOSION
                self.explosions.discard((data["x"], data["y"]))
            elif type == 15:
                # SPAWN_FALLING_BOX
                self.falling_boxes[(data["x"], data["y"])] = data["t"]  # x,y -> ticks
            elif type == 16:
                # REMOVE_FALLING_BOX
                self.falling_boxes.pop((data["x"], data["y"]), None)
            elif type == 17:
                # SPAWN_CRUSHING_BOX
                self.crushing_boxes[(data["x"], data["y"])] = data["t"]
            elif type == 18:
                # REMOVE_CRUSHING_BOX
                self.crushing_boxes.pop((data["x"], data["y"]), None)
            elif type == 19:
                # PLAYER_TAUNT
                if data["id"] not in self.active_taunts:
                    self.active_taunts[data["id"]] = data["t"]
            elif type == 20:
                # SPAWN_POWER_UP
                self.power_ups[(data["x"], data["y"])] = data["t"]  # type
            elif type == 21:
                # REMOVE_POWER_UP
                self.power_ups.pop((data["x"], data["y"]), None)
            elif type == 22:
                # ANGER_INFO
                self.players[0].set_anger(data["0"])
                self.players[1].set_anger(data["1"])
            elif type == 23:
                # ACTIVATE_TRAP
                self.inactive_traps.pop((data["x"], data["y"]), None)
                self.active_traps.add((data["x"], data["y"]))
            elif type == 24:
                # RESET_TRAP
                self.inactive_traps[(data["x"], data["y"])] = [data["t"], data["t"]]  # [max_ticks, remaining_ticks]
                self.active_traps.discard((data["x"], data["y"]))
            elif type == 25:
                # PLAYER_INVERT_KEYBOARD_ON
                self.players[data["id"]].inverted_keyboard = True
            elif type == 26:
                # PLAYER_INVERT_KEYBOARD_OFF
                self.players[data["id"]].inverted_keyboard = False
            elif type == 27:
                # PLAYER_CHANGE_BOMBS_COUNT
                self.players[data["id"]].bombs = data["b"]
            elif type == 28:
                # PLAYER_CHANGE_POWER_AMOUNT
                self.players[data["id"]].power = data["p"]
            elif type == 29:
                # PLAYER_AUTOWALK_ON
                self.players[data["id"]].autowalk = True
            elif type == 30:
                # PLAYER_AUTOWALK_OFF
                self.players[data["id"]].autowalk = False

    def load_next_image_and_timestamp(self, player_id, opened_handle):
        f = opened_handle
        timestamp = f.readline()
        if timestamp == "":
            self.face_replay_file_handles[player_id] = None
            f.close()
            return
        else:
            # print([timestamp])
            timestamp = float(timestamp)
            self.next_faces[player_id][0] = timestamp
            image_data = []
            for _ in range(48):
                line = f.readline()
                if line == "":
                    self.face_replay_file_handles[player_id] = None
                    f.close()
                    return
                image_data.append(line.strip())
            print(len(image_data))
            image_data = "\n".join(image_data)
            image_data = io.StringIO(initial_value=image_data.strip() + "\n")
            img = np.loadtxt(image_data)
            self.next_faces[player_id][1] = img
            line = f.readline().strip()
            if line != "":
                print(line)
                print("ERROR: Wanted to jump empty line but was not empty")

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("-r", "--replay_dir", help="The directory containing the replays for this run.", type=str)
    args = vars(parser.parse_args())

    gr = GameReplay(**args)
    gr.replay()


import logging
import json
import os
import requests

logger = logging.getLogger()
logger.setLevel(logging.INFO)

def lambda_handler(inputPayload, context):
    slack_link = os.environ['SLACK_URL']
    try:
        url = inputPayload['issue']['html_url']
    except Exception as e:
        logger.error(e)
        # return a 500 error code
        res = json.dumps({'statusCode': 500, 'body': f'Error: {e}'})
        return res
    reply = {'text': f"Issue Created: {url}"}
    res = requests.post(slack_link, json=reply)
    return res


# -*- coding: utf-8 -*-
"""
Mini project 1

Dennis Brown, COMP6636, 03 MAR 2021
"""

import numpy as np
import copy
import matplotlib.pyplot as plt

def libsvm_scale_import(filename):
    """
    Read data from a libsvm .scale file
    """
    datafile = open(filename, 'r')

    # First pass: get dimensions of data
    num_samples = 0
    max_feature_id = 0
    for line in datafile:
        num_samples += 1
        tokens = line.split()
        for feature in tokens[1:]:
            feature_id = int(feature.split(':')[0])
            max_feature_id = max(feature_id, max_feature_id)

    # Second pass: read data into array
    data = np.zeros((num_samples, max_feature_id + 1))
    curr_sample = 0
    datafile.seek(0)
    for line in datafile:
        tokens = line.split()
        data[curr_sample][0] = float(tokens[0])
        for feature in tokens[1:]:
            feature_id = int(feature.split(':')[0])
            feature_val = float(feature.split(':')[1])
            data[curr_sample][feature_id] = feature_val
        curr_sample += 1
    datafile.close()

    print('LOADED:', filename, ':', data.shape)

    return data

def get_neighbors(data, test_sample, num_neighbors):
    """
    Given training data, a test sample, and a number of
    neighbors, return the closest neighbors.
    """
    # Calculate all distances from the training samples
    # to this test sample. Collect index, distance into a list.
    indices_and_distances = list()
    for i in range(len(data)):
        dist = np.linalg.norm(test_sample[1:] - (data[i])[1:]) # leave out classification at pos 0
        indices_and_distances.append([i, dist])

    # Sort list by distance
    indices_and_distances.sort(key=lambda _: _[1])

    # Make a list of requested number of closest neighbors from sorted
    # list of indices+distances
    neighbors = list()
    for i in range(num_neighbors):
        neighbors.append(indices_and_distances[i][0])

    return neighbors

def classify_one_sample(data, test_sample, num_neighbors):
    """
    Given training data, a test sample, and a number of neighbors,
    predict which classification the test sample belongs to.
    """
    # Get closest neighbors
    neighbors = get_neighbors(data, test_sample, num_neighbors)

    # Create list of classifications of the neighbors
    classifications = list()
    for i in range(len(neighbors)):
        classifications.append(data[neighbors[i]][0]) # 0 = classification

    # Return the most common classification of the neighbors
    prediction = max(set(classifications), key = classifications.count)
    return prediction

def k_nearest_neighbors(data, test_samples, num_neighbors):
    """
    Given sample data (samples are rows, columns
    features, and samples have classifications in position 0),
    test data, and a number of neighbors, predict which classification
    each test sample belongs to.
    """
    classifications = list()
    for i in range(len(test_samples)):
        output = classify_one_sample(data, test_samples[i], num_neighbors)
        classifications.append(output)
        if ((i % 20) == 0):
            print('\rknn test sample', i, end='')
    print()

    return(classifications)

def check_knn_classifications(y, y_hat):
    """
    Given actual values y and classiciations y_hat,
    return the number of errors
    """
    errors = 0
    for i in range(len(y)):
        if (y[i] != y_hat[i]):
            errors += 1

    return errors

def train_perceptron(data, beta, step_limit):
    """
    Perceptron. Given a set of data (samples are rows, columns
    features, and samples have classifications in position 0),
    a step size (beta), and a step limit, train and return a
    weight vector that can be used to classify the given data.
    """

    # Initialize the weight vector including bias element
    w = np.zeros(len(data[0]))

    # Initialize y_hat
    y_hat = np.zeros(len(data))

    # Slice off y
    y = data[:,0]

    # Repeat the main loop until we have convergence or reach the
    # iteration limit
    steps = 0
    converged = False
    while(not(converged) and (steps < step_limit)):
        converged = True

        # For each sample in the data, calculate w's classification error
        # and update w.
        for i in range(len(data)):
            # Replace classification in sample[0] with a 1 to allow
            # for a biased weight vector
            biased_sample = np.copy(data[i])
            biased_sample[0] = 1

            # Get prediction and error, then update weight vector
            y_hat[i] = 1 if (np.matmul(w.T, biased_sample) > 0) else -1
            error = y[i] - y_hat[i]
            w += biased_sample * error * beta
            steps += 1

            # If error on this element is > a very small value, we have
            # not converged.
            if (abs(error) > 0.000001):
                converged = False

    print('Perceptron:' ,steps, 'steps; converged?', converged)

    return w

def multiclass_train_perceptron(data, beta, step_limit):
    """
    Perceptron. Given a set of data (samples are rows, columns
    features, and samples have classifications in position 0),
    a step size (beta), and a step limit, train and return a
    weight vector that can be used to classify the given data.

    This version works on data with multiple classes by one-vs-rest.
    """
    # Find unique classes
    classes = []
    for i in range(data.shape[0]):
        if (not(data[i][0] in classes)):
            classes.append(data[i][0])

    # For each classification, train perceptron on current class vs.
    # rest of the untrained classes.
    ws = []
    curr_data = copy.deepcopy(data)
    for curr_class in range(len(classes)):

        # Save original classification data
        orig_classes = copy.deepcopy(curr_data[:,0])

        # Reset classification data to 1 (for current class) or -1 for other
        for i in range(curr_data.shape[0]):
            if (curr_data[i][0] == classes[curr_class]):
                curr_data[i][0] = 1
            else:
                curr_data[i][0] = -1

        # Train and find weights
        ws.append(train_perceptron(curr_data, beta, step_limit))

        # Put original classifications back
        for i in range(curr_data.shape[0]):
            curr_data[i][0] = orig_classes[i]

    return ws

def test_perceptron(data, w):
    """
    Given test data and a weight vector w, return number of
    num_misclass when classifying the test data using the
    weights.
    """
    errors = 0

    # Initialize y_hat
    y_hat = np.zeros(len(data))

    # Slice off y
    y = data[:,0]

    # Determine how weights classify each test sample and count
    # num_misclass
    for i in range(len(data)):
        biased_sample = np.copy(data[i])
        biased_sample[0] = 1
        y_hat[i] = 1 if (np.matmul(w.T, biased_sample) > 0) else -1
        if (y[i] != y_hat[i]):
            errors += 1

    return errors

def multiclass_test_perceptron(data, ws):
    """
    Given test data and a weight vector w, return number of
    num_misclass when classifying the test data using the
    weights.

    This version works on data with multiple classes by One vs. All (OVA).
    """
    # Find unique classes
    classes = []
    for i in range(data.shape[0]):
        if (not(data[i][0] in classes)):
            classes.append(data[i][0])

    # For each classification, test perceptron on current class vs.
    # rest of the untested classes.
    errors = []
    curr_data = copy.deepcopy(data)
    for curr_class in range(len(classes)):

        # Save original classification data
        orig_classes = copy.deepcopy(curr_data[:,0])

        # Reset classification data to 1 (for current class) or -1 for other
        for i in range(curr_data.shape[0]):
            if (curr_data[i][0] == classes[curr_class]):
                curr_data[i][0] = 1
            else:
                curr_data[i][0] = -1

        # Train and find weights
        errors.append(test_perceptron(curr_data, ws[curr_class]))

        # Put original classifications back
        for i in range(curr_data.shape[0]):
            curr_data[i][0] = orig_classes[i]

    return errors

def iris_knn():
    """
    Run kNN on the iris dataset for the various numbers of neighbors.
    """
    print("----------\niris kNN")

    # Load data
    data = libsvm_scale_import('data/iris.scale')

    # Shuffle the data because we want to split it into train & test,
    # and it is pre-sorted (we would test against classes we didn't
    # see in training)
    np.random.seed(1) # ensure consistent shuffling
    np.random.shuffle(data)

    # Split up data into training and test data based on split value
    split = 50
    train_data = data[:split]
    test_data = data[split:]

    # Test multiple values of k
    test_ks = np.arange(1, split)
    error_rates = np.zeros(test_ks.shape[0])
    for i in range(len(test_ks)):
        # Classify the test data
        print('Classify with k =', test_ks[i])
        classifications = k_nearest_neighbors(train_data, test_data,
                                              test_ks[i])
        # Check accuracy
        errors = check_knn_classifications(test_data[:,0], classifications)
        error_rates[i] = errors / test_data.shape[0]
        print(errors, 'errors in', test_data.shape[0], 'samples')

    print('ks:', test_ks)
    print('error rates:', error_rates)
    plt.clf()
    plt.plot(test_ks, error_rates, marker='.')
    plt.title('Iris kNN: error rate vs. k')
    plt.xlabel('k')
    plt.ylabel('error rate')
    plt.xlim(left = 0)
    plt.ylim(bottom = 0)
    plt.grid(True)
    plt.savefig('iris_knn.png', dpi = 600)

def iris_perceptron():
    """
    Run Perceptron on the iris dataset in various ways.
    """
    print("----------\niris Perceptron")

    # Load data
    data = libsvm_scale_import('data/iris.scale')

    # Shuffle the data because we want to split it into train & test,
    # and it is pre-sorted (we would test against classes we didn't
    # see in training)
    np.random.seed(1) # ensure consistent shuffling
    np.random.shuffle(data)

    # Split up data into training and test data based on split value
    split = 50
    train_data = data[:split]
    test_data = data[split:]

    # Perform multi-class training and test and collect
    # a weight vector and number of errors for each class
    ws = multiclass_train_perceptron(train_data, 0.1, 100000)
    errors = multiclass_test_perceptron(test_data, ws)

    # Report errors
    print(errors, 'errors in', test_data.shape[0], 'samples')

    # Show sorted weights for every class
    for i in range(len(ws)):

        # Sort weights to find most important
        w = list(ws[i][1:])
        feature_ids = range(1, len(w) + 1)
        print('W:', w)
        labels = []
        for id in feature_ids:
            labels.append(str(int(id)))

        # Report top weights
        plt.clf()
        plt.bar(labels, w)
        plt.title('iris Perceptron: feature weights for class = ' + str(i+1))
        plt.xlabel('feature ID')
        plt.ylabel('weight')
        plt.grid(True)
        plt.savefig('iris_weights' + str(i+1) + '.png', dpi = 600)

def a4a_knn():
    """
    Run kNN on the a4a dataset for various numbers of neighbors.
    """
    print("----------\na4a kNN")

    # Load data
    train_data = libsvm_scale_import('data/a4a')
    test_data = libsvm_scale_import('data/a4a.t')

    # Training data has 1 fewer feature than test data, so add a column
    # of zeros to it so samples have same number of features in train and test
    zero_col = np.zeros((len(train_data), 1))
    train_data = np.hstack((train_data, zero_col))

    # Test multiple values of k
    # This takes over 3 hours to run on my fastest computer.
    test_ks = np.array([1, 3, 5, 11, 21, 31, 41, 51, 61, 71, 81, 91, 101, 201, 301, 401, 501, 601, 701, 801, 901, 1001])
    error_rates = np.zeros(len(test_ks))
    for i in range(len(test_ks)):
        print('Classify with k =', test_ks[i])
        # Classify the test data
        classifications = k_nearest_neighbors(train_data, test_data,
                                              test_ks[i])
        # Check accuracy
        errors = check_knn_classifications(test_data[:,0], classifications)
        error_rates[i] = errors / test_data.shape[0]
        print(errors, 'errors in', test_data.shape[0], 'samples')

    print('ks:', test_ks)
    print('error rates:', error_rates)
    plt.clf()
    plt.plot(test_ks, error_rates, marker='.')
    plt.title('a4a kNN: error rate vs. k')
    plt.xlabel('k')
    plt.ylabel('error rate')
    plt.xlim(left = 0)
    plt.ylim(bottom = 0)
    plt.grid(True)
    plt.savefig('a4a_knn.png', dpi = 600)

def a4a_perceptron():
    """
    Run Perceptron on the a4a dataset in various ways.
    """
    print("----------\na4a Perceptron")

    # Load data
    train_data = libsvm_scale_import('data/a4a')
    test_data = libsvm_scale_import('data/a4a.t')

    # Training data has 1 fewer feature than test data, so add a column
    # of zeros to it so samples have same number of features in train and test
    zero_col = np.zeros((len(train_data), 1))
    train_data = np.hstack((train_data, zero_col))

    # Test multiple values of beta
    test_betas = np.array([0.0001, 0.001, 0.01, 0.1, 1.0, 10.0, 100.0])
    error_rates = np.zeros(test_betas.shape[0])
    ws = []
    best_beta = -1
    best_error_rate = 999999
    for i in range(len(test_betas)):
        print('Classify with beta =', test_betas[i])

        # Train and find weights
        ws.append(train_perceptron(train_data, test_betas[i], 100000))

        # Check accuracy
        errors = test_perceptron(test_data, ws[i])
        error_rates[i] = errors / test_data.shape[0]
        if (error_rates[i] < best_error_rate):
            best_error_rate = error_rates[i]
            best_beta = i
        print(errors, 'errors in', test_data.shape[0], 'samples')

    # Report error rates
    print('betas:', test_betas)
    print('error rates:', error_rates)
    plt.clf()
    plt.plot(test_betas, error_rates, marker='.')
    plt.title('a4a Perceptron: error rate vs. step size for 100000 iterations')
    plt.xscale('log')
    plt.xlabel('step size')
    plt.ylabel('error rate')
    plt.ylim(bottom = 0)
    plt.grid(True)
    plt.savefig('a4a_perceptron.png', dpi = 600)

    # Sort weights to find most important
    w = list(ws[best_beta][1:])
    feature_ids = range(1, len(w) + 1)
    bar_data = list(zip(feature_ids, w))
    bar_data.sort(key = lambda _: abs(_[1]), reverse = True)
    bar_data = np.array(bar_data[:20])
    labels = []
    for id in bar_data[:,0]:
        labels.append(str(int(id)))

    # Report top weights
    plt.clf()
    plt.bar(labels, bar_data[:,1])
    plt.title('a4a Perceptron: 20 most important features')
    plt.xlabel('feature ID')
    plt.ylabel('weight')
    plt.grid(True)
    plt.savefig('a4a_weights.png', dpi = 600)

def main():
    iris_knn()
    iris_perceptron()
    a4a_knn()
    a4a_perceptron()

if __name__ == '__main__':
    main()


import sys
from itertools import combinations as comb
sys.stdin = open('input/20529.txt')
input = sys.stdin.readline

def d(A, B):
    return sum([A[i] != B[i] for i in range(4)])

T = int(input())
for tc in range(T):
    N = int(input())
    mbti = input().split()
    if len(mbti) > 32:
        print(0)
    else:
        ans = 12
        for c in set(list(comb(mbti, 3))):
            dist = d(c[0], c[1]) + d(c[0], c[2]) + d(c[1], c[2])
            ans = min(ans, dist)
        print(ans)

baekjoon/[S1] 20529. 가장 가까운 세 사람의 심리적 거리.py

[S1] 20529. 가장 가까운 세 사람의 심리적 거리.py


from decimal import Decimal
import ffmpeg
import math
import gc

def get_aspect_ratio(width, height):
    gcd = math.gcd(width, height)
    lhs = int(width / gcd)
    rhs = int(height / gcd)
    return f"{lhs}x{rhs}"

def get_raw_duration(video):
    duration_raw = None
    # check framerate at index 0 and 1, because its given like '25/1'
    # ToDo: add other sources for NUMBER_OF_FRAMES => check some files
    try:
        if 'NUMBER_OF_FRAMES-eng' in video['tags'] and 'avg_frame_rate' in video:
            duration_raw = int(video['tags']['NUMBER_OF_FRAMES-eng']) / \
                int((video['avg_frame_rate'][0] + video['avg_frame_rate'][1]))
    except:
        #raise TypeError('Some error happened during the calculation of the raw duration!')
        return duration_raw
    return duration_raw

def get_duration(video):
    duration = None
    try:
        if 'DURATION-eng' in video['tags']:
            # could also be DURATION => search for something with DURATION in its name; might be this one: [value for key, value in programs.items() if 'new york' in key.lower()]
            duration = video['tags']['DURATION-eng'].split('.')[0]
        elif 'DURATION-de' in video['tags']:
            duration = video['tags']['DURATION-de'].split('.')[0]
        elif 'DURATION' in video['tags']:
            duration = video['tags']['DURATION'].split('.')[0]
        else:
            raise TypeError('Cant find duration in tags!')
    except:
        #raise TypeError('Some error happened during the calculation of the duration!')
        return duration
    return duration

def try_get_width(video):
    width = None
    if 'width' in video:
        width = video['width']
    elif 'coded_width' in video:
        width = video['coded_width']
    return width

def try_get_height(video):
    height = None
    if 'height' in video:
        height = video['height']
    elif 'coded_height' in video:
        height = video['coded_height']
    return height

def get_data(path):
    # read the audio/video file from the command line arguments
    media_file = str(path)
    # uses ffprobe command to extract all possible metadata from the media file
    probe = ffmpeg.probe(media_file)
    bitrate = 0.00
    if 'format' in probe:
        bitrate = round(
            Decimal(probe['format'].get('bit_rate'))/(1024*1024), 2)
    streams = probe["streams"]
    video = streams[0]
    codec = video['codec_name']

    # for other codecs => needs to be included in the output file!
    other_codecs = []
    first_cd = True
    for cd in streams:
        if first_cd:
            first_cd = False
            continue
        # creates object with name, type, language, title
        codec_name = cd.get('codec_name', '')
        codec_type = cd.get('codec_type', '')
        codec_language = str
        codec_title = str
        if 'tags' in cd:
            codec_language = cd['tags'].get('language', '')
            codec_title = cd['tags'].get("title", '')
        other_codecs.append({"name": str(codec_name), "type": codec_type,
                            "language": codec_language, "title": codec_title})

    # ToDo: add FPS, and think of a good output for other codecs (e.g. ac3, eac3, aac) => so just comma seperated names
    # could also add audio language (comma seperated) and subtitle language

    duration = get_duration(video)
    duration_raw = get_raw_duration(video)

    height = try_get_height(video)
    width = try_get_width(video)

    aspect_ratio = '0x0' # might look for a better option => 16:9 - excel will convert this to datetime
    if width != None and height != None:
        aspect_ratio = get_aspect_ratio(width, height)

    # clear data
    del streams, video
    gc.collect()

    return {"codec": codec, "other_codecs": other_codecs, "bitrate": bitrate, "duration": duration, "aspect_ratio": aspect_ratio, "dimensions": {"width": width, "height": height}, "raw": {"duration_raw": duration_raw}}


import ROOT
import uproot
from hipe4ml.tree_handler import TreeHandler
import numpy as np

import argparse
import yaml

import sys
sys.path.append('utils')
import utils as utils

utils.set_style()
kBlueC = ROOT.TColor.GetColor('#1f78b4')
kOrangeC  = ROOT.TColor.GetColor('#ff7f00')
ROOT.gROOT.SetBatch()

## create signal extraction class
class SignalExtraction:
    def __init__(self, input_data_hdl, input_mc_hdl=None): ## could be either a pandas or a tree handler

        self.data_hdl = input_data_hdl
        self.mc_hdl = input_mc_hdl
        self.is_3lh = True
        self.bins = 40
        self.n_evts = 1e9
        self.is_matter = False
        self.signal_fit_func = 'dscb'
        self.bkg_fit_func = 'pol1'
        self.performance = False
        self.additional_pave_text = '' ## additional text to be added to the ALICE performance pave

        ## variables
        self.pdf = None
        self.roo_dataset = None

        ### frames to be saved to file
        self.mc_frame_fit = None
        self.data_frame_fit = None
        self.local_pvalue_graph = None

    def process_fit(self, extended_likelihood=True, rooworkspace_path=None):

        if self.is_3lh:
            self.inv_mass_string = '#it{M}_{^{3}He+#pi^{-}}' if self.is_matter else '#it{M}_{^{3}#bar{He}+#pi^{+}}'
            decay_string = '{}^{3}_{#Lambda}H #rightarrow ^{3}He+#pi^{-}' if self.is_matter else '{}^{3}_{#bar{#Lambda}}#bar{H} #rightarrow ^{3}#bar{He}+#pi^{+}'
            tree_var_name = 'fMassH3L'
        else:
            self.inv_mass_string = '#it{M}_{^{4}He+#pi^{-}}' if self.is_matter else '#it{M}_{^{4}#bar{He}+#pi^{+}}'
            decay_string = '{}^{4}_{#Lambda}H #rightarrow ^{4}He+#pi^{-}' if self.is_matter else '{}^{4}_{#bar{#Lambda}}#bar{H} #rightarrow ^{4}#bar{He}+#pi^{+}'
            tree_var_name = 'fMassH4L'

        # define signal and bkg variables
        if self.is_3lh:
            mass = ROOT.RooRealVar('m', self.inv_mass_string, 2.96, 3.04, 'GeV/c^{2}')
            mu = ROOT.RooRealVar('mu', 'hypernucl mass', 2.97, 2.992, 'GeV/c^{2}')
        else:
            mass = ROOT.RooRealVar('m', self.inv_mass_string, 3.89, 3.97, 'GeV/c^{2}')
            mu = ROOT.RooRealVar('mu', 'hypernucl mass', 3.9, 3.95, 'GeV/c^{2}')

        sigma = ROOT.RooRealVar('sigma', 'hypernucl width', 0.001, 0.004, 'GeV/c^{2}')
        a1 = ROOT.RooRealVar('a1', 'a1', 0, 5.)
        a2 = ROOT.RooRealVar('a2', 'a2', 0, 5.)
        n1 = ROOT.RooRealVar('n1', 'n1', 1, 5.)
        n2 = ROOT.RooRealVar('n2', 'n2', 1, 5.)
        c0 = ROOT.RooRealVar('c0', 'constant c0', -1., 1)
        c1 = ROOT.RooRealVar('c1', 'constant c1', -1., 1)

        if self.signal_fit_func == 'dscb':
            signal = ROOT.RooCrystalBall('cb', 'cb', mass, mu, sigma, a1, n1, a2, n2)
        elif self.signal_fit_func == 'gaus':
            signal = ROOT.RooGaussian('gaus', 'gaus', mass, mu, sigma)
        else:
            raise ValueError(f'Invalid signal fit function. Expected one of: dscb, gaus')

        # define background pdf
        if self.bkg_fit_func == 'pol1':
            background = ROOT.RooChebychev('bkg', 'pol1 bkg', mass, ROOT.RooArgList(c0))
        elif self.bkg_fit_func == 'pol2':
            background = ROOT.RooChebychev('bkg', 'pol2 bkg', mass, ROOT.RooArgList(c0, c1))
        else:
            raise ValueError(f'Invalid background fit function. Expected one of: pol1, pol2')

        if extended_likelihood:
            n_signal = ROOT.RooRealVar('n_signal', 'n_signal', 10., 1e6)
            n_background = ROOT.RooRealVar('n_background', 'n_background', 0., 1e6)
        else:
            f = ROOT.RooRealVar('f', 'fraction of signal', 0., 0.4)

        # fix DSCB parameters to MC
        if self.mc_hdl != None:
            mass_roo_mc = utils.ndarray2roo(np.array(self.mc_hdl['fMassH3L'].values, dtype=np.float64), mass, 'histo_mc')
            signal.fitTo(mass_roo_mc, ROOT.RooFit.Range(2.97, 3.01), ROOT.RooFit.PrintLevel(-1))
            a1.setConstant()
            a2.setConstant()
            n1.setConstant()
            n2.setConstant()
            sigma.setRange(sigma.getVal(), sigma.getVal()*1.5)
            self.mc_frame_fit = mass.frame(80)
            self.mc_frame_fit.SetName('mc_frame_fit')
            mass_roo_mc.plotOn(self.mc_frame_fit)
            signal.plotOn(self.mc_frame_fit)
            fit_param = ROOT.TPaveText(0.6, 0.6, 0.9, 0.9, 'NDC')
            fit_param.SetBorderSize(0)
            fit_param.SetFillStyle(0)
            fit_param.SetTextAlign(12)
            fit_param.AddText('#mu = ' + f'{mu.getVal()*1e3:.2f} #pm {mu.getError()*1e3:.2f}' + ' MeV/#it{c}^{2}')
            fit_param.AddText('#sigma = ' + f'{sigma.getVal()*1e3:.2f} #pm {sigma.getError()*1e3:.2f}' + ' MeV/#it{c}^{2}')
            self.mc_frame_fit.addObject(fit_param)

        # define the fit function and perform the actual fit
        if extended_likelihood:
            self.pdf = ROOT.RooAddPdf('total_pdf', 'signal + background', ROOT.RooArgList(signal, background), ROOT.RooArgList(n_signal, n_background))
        else:
            self.pdf = ROOT.RooAddPdf('total_pdf', 'signal + background', ROOT.RooArgList(signal, background), ROOT.RooArgList(f))

        mass_array = np.array(self.data_hdl[tree_var_name].values, dtype=np.float64)
        self.roo_dataset = utils.ndarray2roo(mass_array, mass)
        self.pdf.fitTo(self.roo_dataset, ROOT.RooFit.Extended(extended_likelihood), ROOT.RooFit.Save(True), ROOT.RooFit.PrintLevel(-1))
        ## get fit parameters
        fit_pars = self.pdf.getParameters(self.roo_dataset)
        sigma_val = fit_pars.find('sigma').getVal()
        sigma_val_error = fit_pars.find('sigma').getError()
        mu_val = fit_pars.find('mu').getVal()
        mu_val_error = fit_pars.find('mu').getError()

        if extended_likelihood:
            signal_counts = n_signal.getVal()
            signal_counts_error = n_signal.getError()
            background_counts = n_background.getVal()
            background_counts_error = n_background.getError()

        else:
            signal_counts = (1-f.getVal())*self.roo_dataset.sumEntries()
            signal_counts_error = (1-f.getVal()) * self.roo_dataset.sumEntries()*f.getError()/f.getVal()
            background_counts = f.getVal()*self.roo_dataset.sumEntries()
            background_counts_error = f.getVal() * self.roo_dataset.sumEntries()*f.getError()/f.getVal()

        self.data_frame_fit = mass.frame(self.n_bins)
        self.data_frame_fit.SetName('data_frame_fit')

        self.roo_dataset.plotOn(self.data_frame_fit, ROOT.RooFit.Name('data'), ROOT.RooFit.DrawOption('p'))
        self.pdf.plotOn(self.data_frame_fit, ROOT.RooFit.Components('bkg'), ROOT.RooFit.LineStyle(ROOT.kDashed), ROOT.RooFit.LineColor(kOrangeC))
        self.pdf.plotOn(self.data_frame_fit, ROOT.RooFit.LineColor(kBlueC), ROOT.RooFit.Name('fit_func'))

        self.data_frame_fit.GetYaxis().SetTitleSize(0.06)
        self.data_frame_fit.GetYaxis().SetTitleOffset(0.9)
        self.data_frame_fit.GetYaxis().SetMaxDigits(2)
        self.data_frame_fit.GetXaxis().SetTitleOffset(1.1)

        # signal within 3 sigma
        mass.setRange('signal', mu_val-3*sigma_val, mu_val+3*sigma_val)
        signal_int = signal.createIntegral(ROOT.RooArgSet(mass), ROOT.RooArgSet(mass), 'signal')
        signal_int_val_3s = signal_int.getVal()*signal_counts
        signal_int_val_3s_error = signal_int_val_3s*signal_counts_error/signal_counts
        # background within 3 sigma
        mass.setRange('bkg', mu_val-3*sigma_val, mu_val+3*sigma_val)
        bkg_int = background.createIntegral(ROOT.RooArgSet(mass), ROOT.RooArgSet(mass), 'bkg')
        bkg_int_val_3s = bkg_int.getVal()*background_counts
        bkg_int_val_3s_error = bkg_int_val_3s*background_counts_error/background_counts
        significance = signal_int_val_3s / np.sqrt(signal_int_val_3s + bkg_int_val_3s)
        significance_err = utils.significance_error(signal_int_val_3s, bkg_int_val_3s, signal_int_val_3s_error, bkg_int_val_3s_error)
        s_b_ratio_err = np.sqrt((signal_int_val_3s_error/signal_int_val_3s)**2 + (bkg_int_val_3s_error/bkg_int_val_3s)**2)*signal_int_val_3s/bkg_int_val_3s

        # add pave for stats
        pinfo_vals = ROOT.TPaveText(0.632, 0.5, 0.932, 0.85, 'NDC')
        pinfo_vals.SetBorderSize(0)
        pinfo_vals.SetFillStyle(0)
        pinfo_vals.SetTextAlign(11)
        pinfo_vals.SetTextFont(42)
        pinfo_vals.AddText(f'Signal (S): {signal_counts:.0f} #pm {signal_counts_error:.0f}')
        pinfo_vals.AddText(f'S/B (3 #sigma): {signal_int_val_3s/bkg_int_val_3s:.1f} #pm {s_b_ratio_err:.1f}')
        pinfo_vals.AddText('S/#sqrt{S+B} (3 #sigma): ' + f'{significance:.1f} #pm {significance_err:.1f}')
        pinfo_vals.AddText('#mu = ' + f'{mu_val*1e3:.2f} #pm {mu.getError()*1e3:.2f}' + ' MeV/#it{c}^{2}')
        pinfo_vals.AddText('#sigma = ' + f'{sigma_val*1e3:.2f} #pm {sigma.getError()*1e3:.2f}' + ' MeV/#it{c}^{2}')

        ## add pave for ALICE performance
        if self.performance:
            pinfo_alice = ROOT.TPaveText(0.6, 0.5, 0.93, 0.85, 'NDC')
        else:
            pinfo_alice = ROOT.TPaveText(0.14, 0.6, 0.42, 0.85, 'NDC')
        pinfo_alice.SetBorderSize(0)
        pinfo_alice.SetFillStyle(0)
        pinfo_alice.SetTextAlign(11)
        pinfo_alice.SetTextFont(42)
        pinfo_alice.AddText('ALICE Performance')
        pinfo_alice.AddText('Run 3, pp #sqrt{#it{s}} = 13.6 TeV')
        if not self.performance:
            pinfo_alice.AddText('N_{ev} = ' f'{self.n_evts:.0f} '  '#times 10^{9}')
        pinfo_alice.AddText(decay_string)
        if self.additional_pave_text != '':
            pinfo_alice.AddText(self.additional_pave_text)

        if not self.performance:
            self.data_frame_fit.addObject(pinfo_vals)
        self.data_frame_fit.addObject(pinfo_alice)

        fit_stats = {'signal': [signal_counts, signal_counts_error],
                     'significance': [significance, significance_err], 's_b_ratio': [signal_int_val_3s/bkg_int_val_3s, s_b_ratio_err]}

        if rooworkspace_path != None:
            w = ROOT.RooWorkspace('w')
            sb_model = ROOT.RooStats.ModelConfig('sb_model', w)
            sb_model.SetPdf(self.pdf)
            sb_model.SetParametersOfInterest(ROOT.RooArgSet(n_signal))
            sb_model.SetObservables(ROOT.RooArgSet(mass))
            getattr(w, 'import')(sb_model)
            getattr(w, 'import')(self.roo_dataset)
            w.writeToFile(rooworkspace_path + '/rooworkspace.root', True)

        return fit_stats

    def compute_significance_asymptotic_calc(self, rooworkspace_path, do_local_p0plot=False):
        print("-----------------------------------------------")
        print("Computing significance with asymptotic calculator")
        ## get saved workspace
        workspace_file = ROOT.TFile(rooworkspace_path + '/rooworkspace.root', 'READ')
        w = workspace_file.Get('w')
        roo_abs_data = w.data('data')
        sb_model = w.obj('sb_model')
        poi = sb_model.GetParametersOfInterest().first()
        sb_model.SetSnapshot(ROOT.RooArgSet(poi))
        ## create the b-only model
        b_model = sb_model.Clone()
        b_model.SetName('b_model')
        poi.setVal(0)
        b_model.SetSnapshot(poi)
        b_model.Print()
        # w.var('sigma').setConstant(True)
        w.var('mu').setConstant(True)

        asymp_calc = ROOT.RooStats.AsymptoticCalculator(roo_abs_data, sb_model, b_model)
        asymp_calc.SetPrintLevel(0)
        asymp_calc_result = asymp_calc.GetHypoTest()
        null_p_value = asymp_calc_result.NullPValue()
        null_p_value_err = asymp_calc_result.NullPValueError()
        significance = asymp_calc_result.Significance()
        significance_err = asymp_calc_result.SignificanceError()

        if do_local_p0plot:
            ### perform a scan in mass and compute the significance
            masses = []
            p0_values = []
            p0_values_expected = []
            mass_array = np.linspace(w.var('mu').getMin(), w.var('mu').getMax(), 100)
            for mass in mass_array:

                w.var('mu').setVal(mass)
                w.var('mu').setConstant(True)
                asymp_calc_scan = ROOT.RooStats.AsymptoticCalculator(roo_abs_data, sb_model, b_model)
                asymp_calc_scan.SetOneSidedDiscovery(True)
                asym_calc_result_scan = asymp_calc_scan.GetHypoTest()
                null_p_value_scan = asym_calc_result_scan.NullPValue()
                masses.append(mass)
                p0_values.append(null_p_value_scan)

                print(f"Mass: {mass} MeV/c^2, p0: {null_p_value_scan:.10f}")

            ## create a graph with the p0 values
            self.local_pvalue_graph = ROOT.TGraph(len(masses), np.array(masses), np.array(p0_values))
            self.local_pvalue_graph.SetName('p0_values')
            self.local_pvalue_graph.GetXaxis().SetTitle(self.inv_mass_string)
            self.local_pvalue_graph.GetYaxis().SetTitle('Local p-value')
            # log Y axis
            self.local_pvalue_graph.SetMarkerStyle(20)
            self.local_pvalue_graph.SetMarkerColor(kBlueC)
            self.local_pvalue_graph.SetMarkerSize(0)
            self.local_pvalue_graph.SetLineColor(kBlueC)
            self.local_pvalue_graph.SetLineWidth(2)

        print("****************************************************")
        print(f'p0: {null_p_value:.3E} +/- {null_p_value_err:.3E}')
        print(f'significance: {significance:.5f} +/- {significance_err:.5f}')
        print("****************************************************")

if __name__ == '__main__':

    # set parameters
    parser = argparse.ArgumentParser(
        description='Configure the parameters of the script.')
    parser.add_argument('--config-file', dest='config_file', default='',
                        help='path to the YAML file with configuration.')
    parser.add_argument('--nbins', dest='n_bins', default=30,
                        help='number of bins in the final plot.')
    parser.add_argument('--performance', action='store_true',
                        help="True for performance plot", default=False)
    args = parser.parse_args()

    config_file = open(args.config_file, 'r')
    config = yaml.full_load(config_file)

    input_parquet_data = config['input_parquet_data']
    input_analysis_results = config['input_analysis_results']

    input_parquet_mc = config['input_parquet_mc']

    output_dir = config['output_dir']
    output_file = config['output_file']

    is_4lh = config['is_4lh']
    matter_type = config['matter_type']
    n_bins = config['n_bins']

    performance = args.performance
    data_hdl = TreeHandler(input_parquet_data)
    mc_hdl =  TreeHandler(input_parquet_mc)

    an_vtx_z = uproot.open(input_analysis_results)['hyper-reco-task']['hZvtx']
    n_evts = an_vtx_z.values().sum() / 1e9
    n_evts = round(n_evts, 0)

    signal_extraction = SignalExtraction(data_hdl, mc_hdl)
    signal_extraction.n_bins = n_bins
    signal_extraction.n_evts = n_evts
    signal_extraction.matter_type = matter_type
    signal_extraction.performance = performance
    signal_extraction.is_3lh = not is_4lh
    signal_extraction.bkg_fit_func = 'pol1'

    signal_extraction.process_fit(extended_likelihood=True, rooworkspace_path="../results")
    signal_extraction.compute_significance_asymptotic_calc(rooworkspace_path="../results", do_local_p0plot=True)

    # create output file and save frames
    out_file = ROOT.TFile(f'{output_dir}/{output_file}', 'recreate')
    out_file.cd()
    signal_extraction.data_frame_fit.Write()
    signal_extraction.mc_frame_fit.Write()
    signal_extraction.local_pvalue_graph.Write()
    out_file.Close()

    if is_4lh:
        state_label = '4lh'
    else:
        state_label = '3lh'

    cSignalExtraction = ROOT.TCanvas('cSignalExtraction', 'cSignalExtraction', 800, 600)
    signal_extraction.data_frame_fit.SetTitle('')
    signal_extraction.data_frame_fit.Draw()
    cSignalExtraction.SaveAs(f'{output_dir}/cSignalExtraction_{matter_type}_{state_label}.pdf')


import requests
import datetime as dt
from twilio.rest import Client

account_sid = 'Twilio_api_sid'
auth_token = 'twilio_auth_token'

STOCK = "TSLA"
COMPANY_NAME = "Tesla Inc"

stock_api_key = 'alpha_vantage_api_key'
news_api_key = 'news_api_key'

STOCK_ENDPOINT = "https://www.alphavantage.co/query"
NEWS_ENDPOINT = "https://newsapi.org/v2/everything"

today = dt.datetime.now().date()
lag_1day = str(today - dt.timedelta(days=1))
lag_2day = str(today - dt.timedelta(days=2))
up_down = ''
percent_change = 0

stock_parameters = {
    'function': 'TIME_SERIES_DAILY_ADJUSTED',
    'symbol': STOCK,
    'outputsize': 'compact',
    'apikey': stock_api_key,
    'pageSize': 3,
    'page': 1,
}

news_parameters = {
    'q': COMPANY_NAME,
    'from': lag_2day,
    'to': lag_1day,
    'sortBy': 'publishedAt',
    'apiKey': news_api_key,
}

# Make api request to stock api
stock_response = requests.get(STOCK_ENDPOINT, params=stock_parameters)
stock_response.raise_for_status()
stock_data = stock_response.json()

# Get closing price
try:
    lag_1day_close = float(stock_data['Time Series (Daily)'][lag_1day]['4. close'])
except KeyError:
    lag_1day_close = None
try:
    lag_2day_close = float(stock_data['Time Series (Daily)'][lag_2day]['4. close'])
except KeyError:
    lag_2day_close = None

# Find percent change, and set up_down symbol
if lag_1day_close is not None and lag_2day_close is not None:
    difference = lag_1day_close - lag_2day_close
    percent_change = round((difference / lag_1day_close) * 100)
    if difference < 0:
        up_down = 'ðŸ”»'
    else:
        up_down = 'ðŸ”º'

# Make api request to get news articles
news_response = requests.get(NEWS_ENDPOINT, params=news_parameters)
news_response.raise_for_status()
news_data = news_response.json()
top_news = news_data['articles'][:3]
news_title_list = [top_news[_]['title'] for _ in range(len(top_news))]
news_description_list = [top_news[_]['description'] for _ in range(len(top_news))]

# Send text messages
if percent_change >= 5 or percent_change <= -5:
    for i in range(len(news_title_list)):
        client = Client(account_sid, auth_token)
        message = client.messages \
            .create(
            body=f'{STOCK}: {up_down}{percent_change}%\nHeadline: {news_title_list[i]}\nBrief: {news_description_list[i]}',
            from_='+19257226085',
            to='+15551234567'
        )


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

from skopt import gp_maximize
import numpy as np
from skopt.plots import plot_convergence
np.random.seed(123)

#%matplotlib inline
import matplotlib.pyplot as plt

noise_level = 0.1

def f(x, noise_level=noise_level):
    return np.sin(5 * x[0]) * (1 - np.tanh(x[0] ** 2)) + np.random.randn() * noise_level
# Plot f(x) + contours
x = np.linspace(-2, 2, 400).reshape(-1, 1)
fx = [f(x_i, noise_level=0.0) for x_i in x]
plt.plot(x, fx, "r--", label="True (unknown)")
plt.fill(np.concatenate([x, x[::-1]]),
         np.concatenate(([fx_i - 1.9600 * noise_level for fx_i in fx],
                         [fx_i + 1.9600 * noise_level for fx_i in fx[::-1]])),
         alpha=.2, fc="r", ec="None")
#plt.legend()
#plt.grid()
#plt.show()

res = gp_maximize(f,                  # the function to minimize
                  [(-2.0, 2.0)],      # the bounds on each dimension of x
                  acq_func="EI",      # the acquisition function
                  n_calls=15,         # the number of evaluations of f
                  n_random_starts=5,  # the number of random initialization points
                  noise=0.1**2,       # the noise level (optional)
                  random_state=123)   # the random seed

#print(res)
#plot_convergence(res);
plt.rcParams["figure.figsize"] = (6, 4)

# Plot f(x) + contours
x = np.linspace(-2, 2, 400).reshape(-1, 1)
x_gp = res.space.transform(x.tolist())

fx = [f(x_i, noise_level=0.0) for x_i in x]
plt.plot(x, fx, "r--", label="True (unknown)")
plt.fill(np.concatenate([x, x[::-1]]),
         np.concatenate(([fx_i - 1.9600 * noise_level for fx_i in fx],
                         [fx_i + 1.9600 * noise_level for fx_i in fx[::-1]])),
         alpha=.2, fc="r", ec="None")

# Plot GP(x) + contours
gp = res.models[-1]
y_pred, sigma = gp.predict(x_gp, return_std=True)

plt.plot(x, y_pred, "g--", label=r"$\mu_{GP}(x)$")
plt.fill(np.concatenate([x, x[::-1]]),
         np.concatenate([y_pred - 1.9600 * sigma,
                         (y_pred + 1.9600 * sigma)[::-1]]),
         alpha=.2, fc="g", ec="None")

# Plot sampled points
plt.plot(res.x_iters,
         res.func_vals,
         "r.", markersize=15, label="Observations")

plt.title(r"$x^* = %.4f, f(x^*) = %.4f$" % (res.x[0], res.fun))
plt.legend(loc="best", prop={'size': 8}, numpoints=1)
plt.grid()

plt.show()


# from django.contrib import messages
from json import loads, dumps
from .models import Link
from django.db.models import Sum
from django.db import OperationalError
from tenacity import (retry, stop_after_attempt, wait_fixed,
                      retry_if_exception_type)

import random
import string
import datetime
from django.shortcuts import render, redirect, get_object_or_404
from django.http import (HttpResponse, HttpResponseServerError, Http404,
                         HttpResponseBadRequest)

# For Google Web Crawler to work and website to show up on Google
def robots_txt(request):

    lines = [
        "User-Agent: *",
        "Disallow: /admin/"
        # "Disallow: /*"
    ]
    return HttpResponse("\n".join(lines), content_type="text/plain")

# Returning home page
def index(request):
    stats = get_stats()
    return render(request, 'shortner/index.html', context=stats)

# returns stats for rendering in index.html

def return_last_value(retry_state):
    print(f'\n\n attempt number {retry_state.attempt_number} \n \
        function for which retry was called: {retry_state.fn} \n\n')

@retry(retry=retry_if_exception_type(OperationalError),
       stop=stop_after_attempt(3),
       wait=wait_fixed(0.75),
       retry_error_callback=return_last_value)
def get_stats():

    # generating date information
    d1 = datetime.datetime(2020, 8, 30)
    d2 = datetime.datetime.now()
    time_difference = d2-d1
    months = round(time_difference.days / 30)

    stats = {
        'total_links': Link.objects.all().count(),
        'total_clicks':
            Link.objects.aggregate(total_clicks=Sum('clicks'))['total_clicks'],
        'active_months': months
    }

    return stats

def check(request, shortlink):
    if linkExists(shortlink):
        return HttpResponse(dumps({'link': shortlink, 'available': False}))
    else:
        return HttpResponse(dumps({'link': shortlink, 'available': True}))
    # not strictly required but might be useful for debugging
    print('nothing got returned')

def create(request):
    # assump1: post body exists
    # assump2: post body has 'longlink' defined

    if request.method != 'POST':
        return redirect('/')

    reqBody = loads(request.body)
    longlink = reqBody['longlink']
    shortlink = ''  # temporary empty value

    try:
        shortlink = reqBody['shortlink']

        if shortlink == '':
            # ik it's wrong...sorry.
            raise KeyError('Empty shortlink')
        if linkExists(shortlink):
            res = HttpResponseBadRequest()
            res.reason_phrase = 'Shortlink already taken'
            res.status_code = 400
            return res
    except KeyError:
        shortlink = getShortRandomLink(5)

    obj = Link(shortlink=shortlink, longlink=longlink)
    obj.save()
    return HttpResponse(dumps(obj.getDict()))

@retry(retry=retry_if_exception_type(OperationalError),
       stop=stop_after_attempt(3),
       wait=wait_fixed(0.75),
       retry_error_callback=return_last_value)
def rediretor(request, shortlink):
    shortlinkObj = get_object_or_404(Link, pk=shortlink)

    # uncomment below lines when adding feature
    shortlinkObj.clicks += 1
    shortlinkObj.save()

    return redirect(shortlinkObj.longlink)

def custom_404(request, exception):
    return render(request, 'shortner/404.html', status=404)

def linkExists(shortlink):
    try:
        Link.objects.get(pk=shortlink)
        return True
    except Link.DoesNotExist:
        return False

# ------- helper functions ---------

def getShortRandomLink(length):
    temp = get_random_string(length)
    if linkExists(temp):
        # recursion!
        getShortRandomLink(length)
    return temp

def get_random_string(length):
    letters = string.ascii_lowercase
    result_str = ''.join(random.choice(letters) for i in range(length))
    return result_str

# function to tell user how many clicks their link have gotten
# usable as api/clicky/<shortlink>
def clicks(request, shortlink):
    # print(f"shortlink of cliks is {shortlink}\n")
    if linkExists(shortlink):
        link = Link.objects.get(pk=shortlink)
        return HttpResponse(link.clicks)

    else:
        return HttpResponse('0')