import os, sys
import pandas as pd
import numpy as np
import json, copy
from .split_datasets import read_data,ALL_KEYS,ONE_KEYS,extend_multi_concepts,save_dcur
from .split_datasets import train_test_split,KFold_split,calStatistics,get_max_concepts,id_mapping,write_config

def generate_sequences(df, effective_keys, min_seq_len=3, maxlen = 200, pad_val = -1):
    save_keys = list(effective_keys) + ["selectmasks"]
    dres = {"selectmasks": []}
    dropnum = 0
    for i, row in df.iterrows():
        dcur = save_dcur(row, effective_keys)

        rest, lenrs = len(dcur["responses"]), len(dcur["responses"])
        j = 0
        while lenrs >= j + maxlen:
            rest = rest - (maxlen)
            for key in effective_keys:
                dres.setdefault(key, [])
                if key not in ONE_KEYS:
                    dres[key].append(",".join(dcur[key][j: j + maxlen]))#[str(k) for k in dcur[key][j: j + maxlen]]))
                else:
                    dres[key].append(dcur[key])
            dres["selectmasks"].append(",".join(["1"] * maxlen))

            j += maxlen
        if rest < min_seq_len:# delete sequence len less than min_seq_len
            dropnum += rest
            continue

        pad_dim = maxlen - rest
        for key in effective_keys:
            dres.setdefault(key, [])
            if key not in ONE_KEYS:
                paded_info = np.concatenate([dcur[key][j:], np.array([pad_val] * pad_dim)])
                dres[key].append(",".join([str(k) for k in paded_info]))
            else:
                dres[key].append(dcur[key])
        dres["selectmasks"].append(",".join(["1"] * rest + [str(pad_val)] * pad_dim))

    # after preprocess data, report
    dfinal = dict()
    for key in ALL_KEYS:
        if key in save_keys:
            dfinal[key] = dres[key]
    finaldf = pd.DataFrame(dfinal)
    print(f"dropnum: {dropnum}")
    return finaldf

def generate_window_sequences(df, effective_keys, maxlen=200, pad_val=-1):
    save_keys = list(effective_keys) + ["selectmasks"]
    dres = {"selectmasks": []}
    for i, row in df.iterrows():
        dcur = save_dcur(row, effective_keys)
        lenrs = len(dcur["responses"])
        if lenrs > maxlen:
            for key in effective_keys:
                dres.setdefault(key, [])
                if key not in ONE_KEYS:
                    dres[key].append(",".join(dcur[key][0: maxlen]))#[str(k) for k in dcur[key][0: maxlen]]))
                else:
                    dres[key].append(dcur[key])
            dres["selectmasks"].append(",".join(["1"] * maxlen))
            for j in range(maxlen+1, lenrs+1):
                for key in effective_keys:
                    dres.setdefault(key, [])
                    if key not in ONE_KEYS:
                        dres[key].append(",".join([str(k) for k in dcur[key][j-maxlen: j]]))
                    else:
                        dres[key].append(dcur[key])
                dres["selectmasks"].append(",".join([str(pad_val)] * (maxlen - 1) + ["1"]))
        else:
            for key in effective_keys:
                dres.setdefault(key, [])
                if key not in ONE_KEYS:
                    pad_dim = maxlen - lenrs
                    paded_info = np.concatenate([dcur[key][0:], np.array([pad_val] * pad_dim)])
                    dres[key].append(",".join([str(k) for k in paded_info]))
                else:
                    dres[key].append(dcur[key])
            dres["selectmasks"].append(",".join(["1"] * lenrs + [str(pad_val)] * pad_dim))

    dfinal = dict()
    for key in ALL_KEYS:
        if key in save_keys:
            # print(f"key: {key}, len: {len(dres[key])}")
            dfinal[key] = dres[key]
    finaldf = pd.DataFrame(dfinal)
    return finaldf

def save_id2idx(dkeyid2idx, save_path):
    with open(save_path, "w+") as fout:
        fout.write(json.dumps(dkeyid2idx))

def id_mapping_que(df):
    id_keys = ["questions", "concepts","uid"]
    dres = dict()
    dkeyid2idx = dict()
    print(f"df.columns: {df.columns}")
    for key in df.columns:
        if key not in id_keys:
            dres[key] = df[key]
    for i, row in df.iterrows():
        for key in id_keys:
            if key not in df.columns:
                continue
            dkeyid2idx.setdefault(key, dict())
            dres.setdefault(key, [])
            curids = []
            for id in row[key].split(","):
                sub_ids = id.split('_')
                sub_curids = []
                for sub_id in sub_ids:
                    if sub_id not in dkeyid2idx[key]:
                        dkeyid2idx[key][sub_id] = len(dkeyid2idx[key])
                    sub_curids.append(str(dkeyid2idx[key][sub_id]))
                curids.append("_".join(sub_curids))
            dres[key].append(",".join(curids))
    finaldf = pd.DataFrame(dres)
    return finaldf, dkeyid2idx

def main(dname, fname, dataset_name, configf, min_seq_len = 3, maxlen = 200, kfold = 5):
    """split main function

    Args:
        dname (str): data folder path
        fname (str): the data file used to split, needs 6 columns, format is: (NA indicates the dataset has no corresponding info)
            uid,seqlen: 50121,4
            quetion ids: NA
            concept ids: 7014,7014,7014,7014
            responses: 0,1,1,1
            timestamps: NA
            cost times: NA
        dataset_name (str): dataset name
        configf (str): the dataconfig file path
        min_seq_len (int, optional): the min seqlen, sequences less than this value will be filtered out. Defaults to 3.
        maxlen (int, optional): the max seqlen. Defaults to 200.
        kfold (int, optional): the folds num needs to split. Defaults to 5.

    """
    stares = []

    total_df, effective_keys = read_data(fname)
    #cal max_concepts
    if 'concepts' in effective_keys:
        max_concepts = get_max_concepts(total_df)
    else:
        max_concepts = -1

    oris, _, qs, cs, seqnum = calStatistics(total_df, stares, "original")
    print("="*20)
    print(f"original total interactions: {oris}, qs: {qs}, cs: {cs}, seqnum: {seqnum}")

     # just for id map
    total_df, dkeyid2idx = id_mapping_que(total_df)
    dkeyid2idx["max_concepts"] = max_concepts

    save_id2idx(dkeyid2idx, os.path.join(dname, "keyid2idx.json"))
    effective_keys.add("fold")

    df_save_keys = []
    for key in ALL_KEYS:
        if key in effective_keys:
            df_save_keys.append(key)

    # train test split
    train_df, test_df = train_test_split(total_df, 0.2)
    splitdf = KFold_split(train_df, kfold)
    splitdf[df_save_keys].to_csv(os.path.join(dname, "train_valid_quelevel.csv"), index=None)
    ins, ss, qs, cs, seqnum = calStatistics(splitdf, stares, "original train+valid question level")
    print(f"train+valid original interactions num: {ins}, select num: {ss}, qs: {qs}, cs: {cs}, seqnum: {seqnum}")

    # generate sequences
    split_seqs = generate_sequences(splitdf, effective_keys, min_seq_len, maxlen)
    ins, ss, qs, cs, seqnum = calStatistics(split_seqs, stares, "train+valid sequences question level")
    print(f"train+valid sequences interactions num: {ins}, select num: {ss}, qs: {qs}, cs: {cs}, seqnum: {seqnum}")
    split_seqs.to_csv(os.path.join(dname, "train_valid_sequences_quelevel.csv"), index=None)
    # print(f"split seqs dtypes: {split_seqs.dtypes}")

    # for test dataset
    # add default fold -1 to test!
    test_df["fold"] = [-1] * test_df.shape[0]
    test_seqs = generate_sequences(test_df, list(effective_keys), min_seq_len, maxlen)
    ins, ss, qs, cs, seqnum = calStatistics(test_df, stares, "test original question level")
    print(f"original test interactions num: {ins}, select num: {ss}, qs: {qs}, cs: {cs}, seqnum: {seqnum}")
    ins, ss, qs, cs, seqnum = calStatistics(test_seqs, stares, "test sequences question level")
    print(f"test sequences interactions num: {ins}, select num: {ss}, qs: {qs}, cs: {cs}, seqnum: {seqnum}")
    print("="*20)

    test_window_seqs = generate_window_sequences(test_df, list(effective_keys), maxlen)

    test_df = test_df[df_save_keys]
    test_df.to_csv(os.path.join(dname, "test_quelevel.csv"), index=None)
    test_seqs.to_csv(os.path.join(dname, "test_sequences_quelevel.csv"), index=None)
    test_window_seqs.to_csv(os.path.join(dname, "test_window_sequences_quelevel.csv"), index=None)

    ins, ss, qs, cs, seqnum = calStatistics(test_window_seqs, stares, "test window question level")
    print(f"test window interactions num: {ins}, select num: {ss}, qs: {qs}, cs: {cs}, seqnum: {seqnum}")

    other_config = {
        "train_valid_original_file_quelevel": "train_valid_quelevel.csv",
        "train_valid_file_quelevel": "train_valid_sequences_quelevel.csv",
        "test_file_quelevel": "test_sequences_quelevel.csv",
        "test_window_file_quelevel": "test_window_sequences_quelevel.csv",
        "test_original_file_quelevel": "test_quelevel.csv"
    }

    write_config(dataset_name=dataset_name, dkeyid2idx=dkeyid2idx, effective_keys=effective_keys,
                configf=configf, dpath = dname, k=kfold,min_seq_len = min_seq_len, maxlen=maxlen,other_config=other_config)

    print("="*20)
    print("\n".join(stares))


import unittest

# navigate to the Features folder
import sys
sys.path.append(sys.path[0]+'/../Features')

# UNIVERSAL IMPORT
from universal_imports import *

# import other features
import helper

# Import testing specific libraries
from pandas.testing import assert_frame_equal

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

class test_get_study_sets(unittest.TestCase):

    def setUp(self):

        # test user
        self.user = "sample_user_1"

        # the expected sets to be returned from test user
        self.expected_sets = ["Food", "French Elementary", "Simple English", "Weather"]

        # sort the list in the default order of sort() function
        self.expected_sets.sort()

    def test_get_study_sets(self):

        # get the output from the function we're testing
        found_sets = helper.get_study_sets(self.user)

        # sort the list in the default order of sort() function
        found_sets.sort()

        # check if the produced dataframe and the processed dataframe are the same
        self.assertEqual(found_sets, self.expected_sets)

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

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


from dash import Dash, CeleryManager
from flask_caching import Cache
import dash_bootstrap_components as dbc
from app_main import layout
from app_main.callbacks import register_callbacks
from celery import Celery

REDIS_URL = "redis://localhost:6379"
celery_app = Celery(__name__, broker=REDIS_URL, backend=REDIS_URL)
background_callback_manager = CeleryManager(celery_app)

DEBUG = False
if DEBUG:
    requests_pathname_prefix = None
else:
    requests_pathname_prefix = '/mangem/'

app = Dash(__name__, external_stylesheets=[dbc.icons.FONT_AWESOME], background_callback_manager=background_callback_manager,
           requests_pathname_prefix=requests_pathname_prefix)

app.title = 'MANGEM'
application = app.server

CACHE_CONFIG = {
    'CACHE_TYPE': 'FileSystemCache',
    'CACHE_DIR': 'cache_dir',
    'CACHE_THRESHOLD': 500,  # maximum number of items the cache will store
    'CACHE_DEFAULT_TIMEOUT': 172800  # Default timeout, seconds (3600 s/hr * 48 hr = 172800)
}
cache = Cache()
cache.init_app(application, config=CACHE_CONFIG)

app.layout = layout.get_layout()
register_callbacks(app, cache, background_callback_manager)

if __name__ == '__main__':
    app.run_server(debug=True)


from typing import Dict
from lesson_20.decorator.hobby import Hobby

@Hobby("dancing")
class Human:
    def __init__(self, name: str, age: int) -> None:
        self.__name = name
        self.__age = age

    def __modify_key(self, key: str) -> str:
        return key.replace(f"_{self.__class__.__name__}__", "")

    def json(self) -> Dict[str, str]:
        data = {}

        for key, value in self.__dict__.items():
            data[self.__modify_key(key)] = value

        return data

if __name__ == "__main__":
    john = Human("John", 32)
    print(john.hobby)


from collections import deque

class Solution:
    def getFood(self, grid: [[str]]) -> int:
        """ O(n*m) time, O(n*m) space, worst case go through all points """
        # this is a typical BFS question, each time we try all the possible dirs by adding it to queue
        # iterate the element in curr level, add all next levels into the queue
        # the first time we met a food will be the shorted length by the nature of BFS
        # mark the visited sell so we don't need to go there again
        queue = deque()
        rows = len(grid)
        cols = len(grid[0])
        for row in range(rows):
            for col in range(cols):
                if grid[row][col] == '*':
                    queue.append([row, col])
                    break

        dirs = [(0, 1), (0, -1), (1, 0), (-1, 0)]
        # bfs to check all the neighboring cells, use "V" to indicate this cell has been visited
        curr_step = 0
        while queue:
            curr_step += 1
            for _ in range(len(queue)):
                row, col = queue.popleft()
                for d in dirs:
                    next_row = row + d[0]
                    next_col = col + d[1]
                    if 0 <= next_row < rows and 0 <= next_col < cols:
                        if grid[next_row][next_col] == '#':
                            return curr_step
                        if grid[next_row][next_col] == 'O':
                            grid[next_row][next_col] = 'V'
                            queue.append([next_row, next_col])
        return -1

s = Solution()
print(s.getFood(grid=[["X", "X", "X", "X", "X", "X"], ["X", "*", "O", "O", "O", "X"], ["X", "O", "O", "#", "O", "X"],
                      ["X", "X", "X", "X", "X", "X"]]))


'''
Requirements
1. Create a multiprocessing program that reads in files with defined tasks to perform.
2. The program should use a process pool per task type and use apply_async calls with callback functions.
3. The callback functions will store the results in global lists based on the task to perform.
4. Once all 4034 tasks are done, the code should print out each list and a breakdown of
   the number of each task performed.

Questions:
1. How many processes did you specify for each pool:
   >Finding primes: 4
   >Finding words in a file: 8
   >Changing text to uppercase: 1
   >Finding the sum of numbers: 6
   >Web request to get names of Star Wars people: 12

   >How do you determine these numbers: I tested each task indavidulay and found the minimum pool size for it to
   >run optimally. I then did it all together and adjusted it a bit from there.

2. Specify whether each of the tasks is IO Bound or CPU Bound?
   >Finding primes: CPU Bound
   >Finding words in a file: I/O Bound
   >Changing text to uppercase: CPU Bound
   >Finding the sum of numbers: CPU Bound
   >Web request to get names of Star Wars people: I/O Bound

3. What was your overall time, with:
   >one process in each of your five pools:  63 seconds
   >with the number of processes you show in question one:  19 seconds
'''

import glob
import json
import math
import multiprocessing as mp
import os
import time
from datetime import datetime, timedelta

import numpy as np
import requests

TYPE_PRIME = 'prime'
TYPE_WORD = 'word'
TYPE_UPPER = 'upper'
TYPE_SUM = 'sum'
TYPE_NAME = 'name'

# Global lists to collect the task results
result_primes = []
result_words = []
result_upper = []
result_sums = []
result_names = []

def is_prime(n: int):
    """Primality test using 6k+-1 optimization.
    From: https://en.wikipedia.org/wiki/Primality_test
    """
    if n <= 3:
        return n > 1
    if n % 2 == 0 or n % 3 == 0:
        return False
    i = 5
    while i ** 2 <= n:
        if n % i == 0 or n % (i + 2) == 0:
            return False
        i += 6
    return True

def prime_cb(results):
    global result_primes
    result_primes.append(results)

def task_prime(value):
    """
    Use the is_prime() above
    Add the following to the global list:
        {value} is prime
            - or -
        {value} is not prime
    """
    return f"{value} is prime" if is_prime(value) else f"{value} is not prime"

def word_cb(results):
    global result_words
    result_words.append(results)

def task_word(word):
    """
    search in file 'words.txt'
    Add the following to the global list:
        {word} Found
            - or -
        {word} not found *****
    """
    with open('words.txt') as f:
            found = word in map(lambda x:x[:-1],f.readlines())
            return f"{word} Found" if found else f"{word} not found *****"

def upper_cb(results):
    global result_upper
    result_upper.append(results)

def task_upper(text):
    """
    Add the following to the global list:
        {text} ==>  uppercase version of {text}
    """
    return f"{text} ==> {text.upper()}"

def sum_cb(results):
    global result_sums
    result_sums.append(results)

def task_sum(start_value, end_value):
    """
    Add the following to the global list:
        sum of {start_value:,} to {end_value:,} = {total:,}
    """
    total = sum(range(start_value, end_value + 1))
    return f"sum of {start_value:,} to {end_value:,} = {total:,}"

def name_cb(results):
    global result_names
    result_names.append(results)

def task_name(url):
    """
    use requests module
    Add the following to the global list:
        {url} has name <name>
            - or -
        {url} had an error receiving the information
    """
    response = requests.get(url)
    if response.status_code == 200:
        data = response.json()
        if "name" in data:
            name = data['name']
            return f"{url} has name {name}"
    return f"{url} had an error receiving the information"

def load_json_file(filename):
    if os.path.exists(filename):
        with open(filename) as json_file:
            data = json.load(json_file)
        return data
    else:
        return {}

def main():
    begin_time = time.time()
    # Call all global lists so we can append to them
    global result_primes, result_words, result_upper, result_sums, result_names

    # Create all the process pools for each task type
    prime_pool = mp.Pool(4)
    word_pool = mp.Pool(8)
    upper_pool = mp.Pool(1)
    sum_pool = mp.Pool(6)
    name_pool = mp.Pool(12)

    # List to take the value from the tasks
    prime_params = []
    word_params = []
    upper_params = []
    sum_params = []
    name_params  = []

    count = 0
    print("Starting to retrieve tasks from file")
    # Open each task file and add the task to a list
    task_files = glob.glob("tasks/*.task")
    for filename in task_files:
        task = load_json_file(filename)
        count += 1
        task_type = task['task']
        if task_type == TYPE_PRIME:
            prime_params.append((task['value'],))
        elif task_type == TYPE_WORD:
            word_params.append((task['word'],))
        elif task_type == TYPE_UPPER:
            upper_params.append((task['text'],))
        elif task_type == TYPE_SUM:
            sum_params.append((task['start'], task['end']))
        elif task_type == TYPE_NAME:
            name_params.append((task['url'],))
        else:
            print(f'Error: unknown task type {task_type}')

    # start each task asynchronously using a call back function to
    # add results to a global list
    print("Running tasks")
    for args in prime_params:
        prime_pool.apply_async(task_prime, args, callback = prime_cb)

    for args in word_params:
        word_pool.apply_async(task_word, args, callback= word_cb)

    for args in upper_params:
        upper_pool.apply_async(task_upper, args, callback= upper_cb)

    for args in sum_params:
        sum_pool.apply_async(task_sum, args, callback= sum_cb)

    for args in name_params:
        name_pool.apply_async(task_name, args, callback= name_cb)

    # Close and end each pool
    prime_pool.close()
    word_pool.close()
    upper_pool.close()
    sum_pool.close()
    name_pool.close()

    prime_pool.join()
    word_pool.join()
    upper_pool.join()
    sum_pool.join()
    name_pool.join()

    # Display all the results
    def print_list(lst):
        for item in lst:
            print(item)
        print(' ')

    print('-' * 80)
    print(f'Primes: {len(result_primes)}')
    print_list(result_primes)

    print('-' * 80)
    print(f'Words: {len(result_words)}')
    print_list(result_words)

    print('-' * 80)
    print(f'Uppercase: {len(result_upper)}')
    print_list(result_upper)

    print('-' * 80)
    print(f'Sums: {len(result_sums)}')
    print_list(result_sums)

    print('-' * 80)
    print(f'Names: {len(result_names)}')
    print_list(result_names)

    print(f'Number of Primes tasks: {len(result_primes)}')
    print(f'Number of Words tasks: {len(result_words)}')
    print(f'Number of Uppercase tasks: {len(result_upper)}')
    print(f'Number of Sums tasks: {len(result_sums)}')
    print(f'Number of Names tasks: {len(result_names)}')
    print(f'Finished processes {count} tasks = {time.time() - begin_time}')

if __name__ == '__main__':
    main()


from typing import Any, Optional, Sequence, Tuple, Union

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops.layers.torch import Rearrange
from monai.networks.blocks.convolutions import Convolution
from monai.networks.blocks.segresnet_block import ResBlock
from monai.networks.blocks.upsample import UpSample
from monai.networks.layers.factories import Act
from monai.networks.layers.utils import get_norm_layer
from monai.utils import InterpolateMode, UpsampleMode

from . import layers

conv1x1 = layers.ddpm_conv1x1
conv3x3 = layers.ddpm_conv3x3
NIN = layers.NIN
default_init = layers.default_init
AttentionBlock = layers.AttentionBlock

class GaussianFourierProjection(nn.Module):
    """Gaussian Fourier embeddings for noise levels."""

    def __init__(self, embedding_size=256, scale=1.0, learnable=False):
        super().__init__()
        self.W = nn.Parameter(
            torch.randn(embedding_size) * scale, requires_grad=learnable
        )

    def forward(self, x):
        x_proj = x[:, None] * self.W[None, :] * 2 * np.pi
        # print("x:", x.shape, "x_proj:", x_proj.shape)
        return torch.cat([torch.sin(x_proj), torch.cos(x_proj)], dim=-1)

########## Code for 3D brain reconstruction models ##############

def get_conv_layer(
    spatial_dims: int,
    in_channels: int,
    out_channels: int,
    kernel_size: int = 3,
    stride: int = 1,
    bias: bool = True,
    dilation: int = 1,
    init_scale=1.0,
):
    conv = Convolution(
        spatial_dims,
        in_channels,
        out_channels,
        strides=stride,
        kernel_size=kernel_size,
        bias=bias,
        conv_only=True,
        dilation=dilation,
    )

    conv.conv.weight.data = default_init(init_scale)(conv.conv.weight.data.shape)
    if bias:
        nn.init.zeros_(conv.conv.bias)

    return conv

def get_upsample_layer(
    spatial_dims: int,
    in_channels: int,
    upsample_mode: Union[UpsampleMode, str] = "nontrainable",
    interp_mode=InterpolateMode.LINEAR,
    scale_factor: int = 2,
):
    return UpSample(
        spatial_dims=spatial_dims,
        in_channels=in_channels,
        out_channels=in_channels,
        scale_factor=scale_factor,
        mode=upsample_mode,
        interp_mode=interp_mode,
        align_corners=None if interp_mode == InterpolateMode.NEAREST else False,
    )

def make_dense_layer(in_sz, out_sz):
    dense = nn.Linear(in_sz, out_sz)
    dense.weight.data = default_init()(dense.weight.shape)
    nn.init.zeros_(dense.bias)
    return dense

class ResBlockpp(nn.Module):
    """
    Modified ResBlock form Monai implementation
    [LINK]
    ResBlock employs skip connection and two convolution blocks and is used
    in SegResNet based on `3D MRI brain tumor segmentation using autoencoder regularization
    <https://arxiv.org/pdf/1810.11654.pdf>`_.
    """

    def __init__(
        self,
        spatial_dims: int,
        in_channels: int,
        norm: Union[Tuple, str],
        act: str = "swish",
        kernel_size: int = 3,
        dilation: int = 1,
        init_scale=0.0,
    ) -> None:
        """
        Args:
            spatial_dims: number of spatial dimensions, could be 1, 2 or 3.
            in_channels: number of input channels.
            norm: feature normalization type and arguments.
            kernel_size: convolution kernel size, the value should be an odd number. Defaults to 3.
            dilation: dilation size of kernel for larger receptive fields
            init_scale: variance scale used to initialize kernel params via ddpm initializer
        """

        super().__init__()

        if kernel_size % 2 != 1:
            raise AssertionError("kernel_size should be an odd number.")

        self.norm1 = get_norm_layer(
            name=norm, spatial_dims=spatial_dims, channels=in_channels
        )
        self.norm2 = get_norm_layer(
            name=norm, spatial_dims=spatial_dims, channels=in_channels
        )
        # self.act = Act["mish"](inplace=True)
        self.act = Act[act]()

        # Following convention of the BigGAN blocks above
        # first conv uses default init scale, second uses config
        self.conv1 = get_conv_layer(
            spatial_dims,
            in_channels=in_channels,
            out_channels=in_channels,
            dilation=dilation,
            kernel_size=kernel_size,
        )
        self.conv2 = get_conv_layer(
            spatial_dims,
            in_channels=in_channels,
            out_channels=in_channels,
            dilation=dilation,
            init_scale=init_scale,
            kernel_size=kernel_size,
        )

    def forward(self, x):
        identity = x

        x = self.norm1(x)
        x = self.act(x)
        x = self.conv1(x)

        x = self.norm2(x)
        x = self.act(x)
        x = self.conv2(x)

        x += identity
        # x /= torch.sqrt(torch.tensor(2.0, requires_grad=False))

        return x

class SegResBlockpp(nn.Module):
    """
    ResBlock employs skip connection and two convolution blocks and is used
    in SegResNet based on `3D MRI brain tumor segmentation using autoencoder regularization
    <https://arxiv.org/pdf/1810.11654.pdf>`_.
    """

    def __init__(
        self,
        in_channels: int,
        norm: Union[Tuple, str],
        spatial_dims: int = 3,
        act: str = "swish",
        kernel_size: int = 3,
        temb_dim: int = None,
        pre_conv: Any = None,
        attention_heads: int = 0,
        dilation: int = 1,
        resblock_pp: bool = False,
        jit: bool = False,
    ) -> None:
        super().__init__()

        self.pre_conv = pre_conv
        self.n_channels = in_channels

        if not resblock_pp:
            self.resblock = ResBlock(
                spatial_dims, in_channels, norm, kernel_size=kernel_size
            )
        else:
            self.resblock = ResBlockpp(
                spatial_dims,
                in_channels,
                norm,
                dilation=dilation,
                kernel_size=kernel_size,
                act=act,
            )

        if jit:
            # print("Jitting resblock")
            self.resblock = torch.jit.script(self.resblock)
        self.attention = attention_heads

        if temb_dim is not None:
            self.dense = make_dense_layer(temb_dim, in_channels * 2)

        self.act = nn.SiLU()

        if attention_heads > 0:
            self.attn = AttentionBlock(channels=in_channels, num_heads=attention_heads)

    def forward(self, x, temb=None):
        if self.pre_conv is not None:
            x = self.pre_conv(x)

        x = self.resblock(x)

        # If time embedding provided
        # Conditioning is acheived via time embedding
        if temb is not None:
            cond_info = self.dense(self.act(temb))[:, :, None, None, None]
            gamma, beta = torch.split(
                cond_info, (self.n_channels, self.n_channels), dim=1
            )
            x = x * gamma + beta

        if self.attention:
            x = self.attn(x)

        return x

class ResnetBlockBigGANpp(nn.Module):
    """
    BigGAN block adapted from song. Notably, the conditioning is done b/w convs
    norm - act - conv --> time_cond --> norm - act - drop - conv --> skip
    """

    def __init__(
        self,
        in_channels,
        kernel_size: int = 3,
        spatial_dims: int = 3,
        act: str = "swish",
        temb_dim: int = None,
        dropout: float = 0.0,
        init_scale: float = 0.0,
        pre_conv: Any = None,
    ):
        super().__init__()

        self.n_channels = in_channels
        self.pre_conv = pre_conv

        # print("IN_CHANNELS:", in_channels)
        self.norm_0 = get_norm_layer(
            name=("GROUP", {"num_groups": min(in_channels // 4, 32), "eps": 1e-6}),
            spatial_dims=spatial_dims,
            channels=in_channels,
        )

        self.conv_0 = get_conv_layer(
            spatial_dims,
            in_channels=in_channels,
            out_channels=in_channels,
            kernel_size=kernel_size,
        )

        if temb_dim is not None:
            self.dense = make_dense_layer(temb_dim, in_channels * 2)

        self.norm_1 = get_norm_layer(
            name=("GROUP", {"num_groups": min(in_channels // 4, 32), "eps": 1e-6}),
            spatial_dims=spatial_dims,
            channels=in_channels,
        )
        self.conv_1 = get_conv_layer(
            spatial_dims,
            in_channels=in_channels,
            out_channels=in_channels,
            init_scale=init_scale,
            kernel_size=kernel_size,
        )
        self.dropout = nn.Dropout(dropout)

        self.act = Act[act]()

    def forward(self, x, temb):
        if self.pre_conv is not None:
            # This is where the downsample happens
            x = self.pre_conv(x)

        h = self.act(self.norm_0(x))

        h = self.conv_0(h)
        # FiLM-like conditioning for each feature map via time embedding
        cond_info = self.dense(self.act(temb))[:, :, None, None, None]
        gamma, beta = torch.split(cond_info, (self.n_channels, self.n_channels), dim=1)
        h = h * gamma + beta

        h = self.act(self.norm_1(h))
        h = self.dropout(h)
        h = self.conv_1(h)

        x = x + h

        return x

class AttentionBlock3d(nn.Module):
    """Channel-wise 3D self-attention block."""

    def __init__(
        self, channels, num_head_channels=None, skip_scale=False, init_scale=0.1
    ):
        super().__init__()
        torch.random.manual_seed(42)
        self.norm = nn.GroupNorm(num_groups=8, num_channels=channels, eps=1e-6)
        self.qkv = nn.Conv3d(channels, channels * 3, kernel_size=1)
        self.proj = nn.Conv3d(channels, channels, kernel_size=1)
        self.spatial_flatten = Rearrange(pattern="b c h w d -> b c (h w d)")
        self.scale = int(channels) ** (-0.5)
        self.skip_scale = np.sqrt(2.0) ** -1 if skip_scale else 1.0

        # self.num_heads = channels // num_head_channels if num_head_channels is not None else 1
        # self.scale = 1 / np.sqrt(channels / self.num_heads)

        # Initialize weights
        self.qkv.weight.data = default_init(init_scale)(self.qkv.weight.data.shape)
        self.proj.weight.data = default_init(init_scale)(self.proj.weight.data.shape)
        nn.init.zeros_(self.qkv.bias)
        nn.init.zeros_(self.proj.bias)

    def forward(self, x):
        B, C, H, W, D = x.shape
        q, k, v = self.qkv(self.norm(x)).chunk(3, dim=1)

        q = self.spatial_flatten(q)
        k = self.spatial_flatten(k)
        v = self.spatial_flatten(v)

        with torch.cuda.amp.autocast(enabled=True, dtype=torch.float32):
            w = torch.einsum("b c q, b c k -> b q k", q, k) * self.scale
            w = F.softmax(w, dim=-1)
            # print("INSIDE ATTENTION BLOCK:", w.dtype)
        h = torch.einsum("b q k , b c k -> b c q", w, v)
        # print("OUTSIDE ATTENTION BLOCK:", h.dtype)
        h = torch.reshape(h, (B, C, H, W, D))
        h = self.proj(h)
        x = x + h
        x = x * self.skip_scale

        return x


""" Test case for service functions of firestorm app of project ten. """
import logging

from django.test import TestCase
from django.test.client import RequestFactory

from advertiser.factories.advertiser_factory import ADVERTISER_FACTORY
from common.test_utils import EnhancedTestCase
from firestorm.factories.ad_rep_factory import AD_REP_FACTORY
from firestorm.models import AdRepAdvertiser
from firestorm.service import build_adv_url_with_ad_rep, get_consumer_bonus_pool

LOG = logging.getLogger('ten.%s' % __name__)
LOG.setLevel(logging.INFO)

class TestGetConsumerBonusPool(TestCase):
    """ Test case for get_consumer_bonus_pool. """

    def test_get_consumer_bonus_pool(self):
        """ Assert the consumer bonus pool is calculated correctly. """
        ad_rep = AD_REP_FACTORY.create_ad_rep()
        pool = get_consumer_bonus_pool()
        self.assertEqual(pool, 0)
        AD_REP_FACTORY.qualify_ad_rep(ad_rep)
        self.assertEqual(ad_rep.verified_consumers().count(), 10)
        ad_rep.consumer_points = '100'
        ad_rep.save()
        pool = get_consumer_bonus_pool()
        self.assertEqual(pool, 100)

class TestBuildURLForAdRepLoading(EnhancedTestCase):
    """ Test cases for service function build_adv_url_with_ad_rep that builds
    a url that calls join-me view and then redirects to final destination URL
    (after loading ad rep in session). If there is no ad rep, for advertiser
    (and no url supplied), then just pass back the destination URL.
    """

    def prep_test(self, ad_rep_advertiser=False):
        """ Create a request to test with. """
        self.advertiser = ADVERTISER_FACTORY.create_advertiser()
        self.ad_rep = AD_REP_FACTORY.create_ad_rep()
        factory = RequestFactory()
        self.request = factory.get('/hudson-valley/', follow=True)
        self.request.session = self.session
        self.request.session['ad_rep_id'] = self.ad_rep.id
        self.request.session['referring_ad_rep_dict'] = self.ad_rep
        if ad_rep_advertiser:
            AdRepAdvertiser.objects.create_update_rep(
                self.request, self.advertiser, self.ad_rep)

    def test_no_ad_rep_no_ad_rep_url(self):
        """ Assert when advertiser does not have an ad rep and no ad rep url
        supplied as a parameter, the destination URL is passed back.
        """
        self.prep_test()
        test_url = '/omega/alpha/pluralis/'
        formatted_url = build_adv_url_with_ad_rep(self.advertiser, test_url)
        self.assertEqual(formatted_url, test_url)

    def test_ad_rep_no_ad_rep_url(self):
        """ Assert when advertise  has an ad rep the link is formatted and
        returned correctly.
        """
        self.prep_test(True)
        test_url = '/omega/alpha/pluralis/'
        formatted_url = build_adv_url_with_ad_rep(self.advertiser, test_url)
        self.assertEqual(formatted_url,
            '/join-me/omega/alpha/pluralis/%s/' % self.ad_rep.url)

    def test_ad_rep_url(self):
        """ Assert whatever is passed at the ad rep url is used in the finally
        return (minimizes db hits).
        """
        self.prep_test(False)
        test_url = '/omega/alpha/pluralis/'
        formatted_url = build_adv_url_with_ad_rep(
            self.advertiser, test_url, self.ad_rep.url)
        self.assertEqual(formatted_url,
            '/join-me/omega/alpha/pluralis/%s/' % self.ad_rep.url)


from flask import Blueprint, request
from big_data.db import get_db, close_db
import json
from big_data.spark import get_spark, get_client
import urllib.parse

from pyspark.ml.feature import StringIndexer
from pyspark.ml.feature import RFormula
from pyspark.ml.feature import StandardScaler

from pyspark.ml.classification import LogisticRegression
from pyspark.ml.classification import DecisionTreeClassifier
from pyspark.ml.classification import RandomForestClassifier
from pyspark.ml.classification import GBTClassifier
from pyspark.ml.classification import NaiveBayes
from pyspark.ml.classification import LinearSVC
from pyspark.ml.regression import LinearRegression
from pyspark.ml.regression import RandomForestRegressor
from pyspark.ml.regression import GBTRegressor
from pyspark.ml.evaluation import BinaryClassificationEvaluator
from pyspark.ml.evaluation import MulticlassClassificationEvaluator
from pyspark.ml.evaluation import RegressionEvaluator

import logging
LOG_FORMAT = "/********** %(asctime)s - %(levelname)s - %(message)s **********/"
logging.basicConfig(level=logging.DEBUG, format=LOG_FORMAT)

bp = Blueprint('add-model', __name__, url_prefix='/add-model')

# 根据 dataset_id 获取 dataframe
def getDF(dataset_id):
    db = get_db()
    c = db.cursor()

    # 通过 dataset_id 获取 project_id 和 filename
    t = (dataset_id,)
    c.execute("SELECT project_id, filename FROM dataset WHERE id=?", t)
    res = c.fetchone()
    project_id = tuple(res)[0]
    filename = tuple(res)[1]

    # 通过 project_id 获取 username
    t = (project_id,)
    c.execute("SELECT username FROM project WHERE id=?", t)
    res = c.fetchone()
    username = tuple(res)[0]

    # 拼接出 hdfs_path，获得数据集 df
    hdfs_path = "/user/" + username + "/" + str(project_id) + "/" + filename + ".parquet"
    df = get_spark().read.load(hdfs_path)

    return df

# 数据预处理
def dataPreProcess(df, label, features, problem_type):
    # 把 df 中除 label 列的数据类型是字符串的各列先 StringIndexer 一下
    for i, dtype in enumerate(df.dtypes):
        if dtype[1] == "string" and i != int(label):
            df = df.withColumnRenamed(dtype[0], dtype[0] + "_old")
            indexer = StringIndexer(inputCol = dtype[0] + "_old", outputCol = dtype[0])
            df = indexer.fit(df).transform(df)

    # 如果是分类问题，用 StringIndexer 对 label 列进行处理
    if problem_type != 3:
        label_name = df.dtypes[int(label)][0]
        df = df.withColumnRenamed(label_name, label_name + "_old")
        indexer = StringIndexer(inputCol = label_name + "_old", outputCol = label_name)
        df = indexer.fit(df).transform(df)

    # 装配 label 和 features vector
    formula_str = df.columns[-1] + "~"
    for feature in features:
        formula_str = formula_str + df.columns[int(feature)] + "+"
    formula_str = formula_str[:-1]
    formula = RFormula(
        formula = formula_str,
        featuresCol = "features",
        labelCol = "label"
    )
    df = formula.fit(df).transform(df)

    # 特征归一化处理
    df = df.withColumnRenamed("features", "features_old")
    standardscaler = StandardScaler().setInputCol("features_old").setOutputCol("features")
    df = standardscaler.fit(df).transform(df)

    return df

# 根据参数配置分类器
def getClassifier(type, params):
    if type == 1: # Logistic Regression
        lr_aggregation_depth = int(params["lr_aggregation_depth"])
        lr_elastic_net_param = float(params["lr_elastic_net_param"])
        lr_fit_intercept = True if params["lr_fit_intercept"] == "true" else False
        lr_max_iter = int(params["lr_max_iter"])
        lr_reg_param = float(params["lr_reg_param"])
        cf = LogisticRegression(
            aggregationDepth = lr_aggregation_depth,
            elasticNetParam = lr_elastic_net_param,
            fitIntercept = lr_fit_intercept,
            maxIter = lr_max_iter,
            regParam = lr_reg_param
        )
    elif type == 2: # Decision Tree Classifier
        dt_max_depth = int(params["dt_max_depth"])
        dt_impurity = params["dt_impurity"]
        cf = DecisionTreeClassifier(
            maxDepth = dt_max_depth,
            impurity = dt_impurity
        )
    elif type == 3: # Random Forest Classifier
        rf_impurity = params["rf_impurity"]
        rf_max_depth = int(params["rf_max_depth"])
        rf_num_trees = int(params["rf_num_trees"])
        cf = RandomForestClassifier(
            impurity = rf_impurity,
            maxDepth = rf_max_depth,
            numTrees = rf_num_trees
        )
    elif type == 4: # Gradient-boosted Tree Classifier
        gbt_max_depth = int(params["gbt_max_depth"])
        gbt_max_iter = int(params["gbt_max_iter"])
        gbt_step_size = float(params["gbt_step_size"])
        gbt_subsampling_rate = float(params["gbt_subsampling_rate"])
        cf = GBTClassifier(
            maxDepth = gbt_max_depth,
            maxIter = gbt_max_iter,
            stepSize = gbt_step_size,
            subsamplingRate = gbt_subsampling_rate
        )
    elif type == 5: # Naive Bayes
        nb_smoothing = float(params["nb_smoothing"])
        nb_model_type = params["nb_model_type"]
        cf = NaiveBayes(
            smoothing = nb_smoothing,
            modelType = nb_model_type
        )
    elif type == 6: # Linear Support Vector Machine
        lsvc_max_iter = int(params["lsvc_max_iter"])
        lsvc_reg_param = float(params["lsvc_reg_param"])
        lsvc_fit_intercept = True if params["lsvc_fit_intercept"] == "true" else False
        lsvc_standardization = True if params["lsvc_standardization"] == "true" else False
        lsvc_aggregation_depth = int(params["lsvc_aggregation_depth"])
        cf = LinearSVC(
            maxIter = lsvc_max_iter,
            regParam = lsvc_reg_param,
            fitIntercept = lsvc_fit_intercept,
            standardization = lsvc_standardization,
            aggregationDepth = lsvc_aggregation_depth
        )
    elif type == 7: # Linear Regression
        l_rgs_max_iter = int(params["l_rgs_max_iter"])
        l_rgs_reg_param = float(params["l_rgs_reg_param"])
        cf = LinearRegression(
            maxIter = l_rgs_max_iter,
            regParam = l_rgs_reg_param,
        )
    elif type == 8: # Random Forest Regression
        rf_rgs_max_depth = int(params["rf_rgs_max_depth"])
        rf_rgs_num_trees = int(params["rf_rgs_num_trees"])
        cf = RandomForestRegressor(
            maxDepth = rf_rgs_max_depth,
            numTrees = rf_rgs_num_trees
        )
    elif type == 9: # Gradient-boosted Tree Regression
        gbt_rgs_max_depth = int(params["gbt_rgs_max_depth"])
        gbt_rgs_max_iter = int(params["gbt_rgs_max_iter"])
        cf = GBTRegressor(
            maxDepth = gbt_rgs_max_depth,
            maxIter = gbt_rgs_max_iter
        )

    return cf

# 获取当前项目下的数据集
@bp.route('get-dataset/', methods = ['POST'])
def getDataset():
    # 变量初始化 & 接收请求数据
    status = 0
    dataset_num = 0
    datasets = []
    received_data = request.get_json(silent=True)

    # logging
    logging.info("received_data: " + json.dumps(received_data))

    # 应该接收到有:
    # project_id
    project_id = received_data['id']

    # 获取数据集数量
    db = get_db()
    c = db.cursor()
    t = (project_id,)
    c.execute("SELECT COUNT(*) from dataset WHERE project_id=?", t)
    res = c.fetchone()
    dataset_num = tuple(res)[0]

    # 获取数据集
    if dataset_num == 0:
        datasets = []
    else:
        tmp = {
            "id": 0,
            "name": ""
        }
        for dataset in c.execute("SELECT id, name FROM dataset WHERE project_id=?", t):
            tmp["id"] = tuple(dataset)[0]
            tmp["name"] = tuple(dataset)[1]
            datasets.append(tmp.copy())

    status = 1

    # 接口约定：
    # status: 状态码
    # dataset_num: 数据集数量
    # datasets: 数据集数组
    #     [{"id": "...", "name": "..."}]
    return {
        "status": status,
        "dataset_num": dataset_num,
        "datasets": datasets
    }

# 获取数据集的列
@bp.route('get-columns/', methods = ['POST'])
def getColumns():
    # 变量初始化 & 接收请求数据
    status = 0
    columns = []
    received_data = request.get_json(silent=True)

    # logging
    logging.info("received_data: " + json.dumps(received_data))

    # 应该接收到有:
    # dataset_id
    dataset_id = received_data['dataset_id']

    # 获取 dataframe
    df = getDF(dataset_id)

    # 获取 columns
    tmp = {
        "index": 0,
        "name": "",
    }
    for i, column in enumerate(df.columns):
        tmp["index"] = i
        tmp["name"] = column
        columns.append(tmp.copy())

    status = 1

    # 接口约定：
    # status: 状态码
    # columns: 列
    return {
        "status": status,
        "columns": columns
    }

# 训练模型
@bp.route('train/', methods = ['POST'])
def train():
    # 变量初始化 & 接收请求数据
    status = 0
    results = {}
    received_data = request.get_json(silent=True)

    # logging
    logging.info("received_data: " + json.dumps(received_data))

    # 应该接收到有:
    # dataset_id, label, features, problem_type, classifier, classifier_params
    dataset_id = received_data['dataset_id']
    label = received_data['label']
    features = received_data['features']
    problem_type = received_data['problem_type']
    classifier = received_data['classifier']
    classifier_params = received_data['classifier_params']

    # 获取 dataframe
    df = getDF(dataset_id)

    # 数据预处理
    df = dataPreProcess(df, label, features, problem_type)

    # 划分训练集和测试集
    (trainingData, testData) = df.randomSplit([0.7, 0.3])

    # 训练
    cf = getClassifier(classifier, classifier_params)
    model = cf.fit(trainingData)

    # 把训练好的模型分别作用在训练集和测试集上
    predict_train = model.transform(trainingData)
    predict_test = model.transform(testData)

    # 模型评估
    if problem_type == 1: # 二分类问题
        evaluator = BinaryClassificationEvaluator( # 计算 area under ROC
            rawPredictionCol = "rawPrediction",
            labelCol = "label"
        )
        results["train_roc"] = evaluator.evaluate(predict_train)
        results["test_roc"] = evaluator.evaluate(predict_test)

        evaluator = BinaryClassificationEvaluator( # 计算 area under PR
            rawPredictionCol = "rawPrediction",
            labelCol = "label",
            metricName = "areaUnderPR"
        )
        results["train_pr"] = evaluator.evaluate(predict_train)
        results["test_pr"] = evaluator.evaluate(predict_test)

    elif problem_type == 2: # 多分类问题
        evaluator = MulticlassClassificationEvaluator( # 计算 acc
            labelCol = "label",
            predictionCol = "prediction",
            metricName = 'accuracy'
        )
        results["train_acc"] = evaluator.evaluate(predict_train)
        results["test_acc"] = evaluator.evaluate(predict_test)

        evaluator = MulticlassClassificationEvaluator( # 计算 f1 score
            labelCol = "label",
            predictionCol = "prediction",
            metricName = 'f1'
        )
        results["train_f1"] = evaluator.evaluate(predict_train)
        results["test_f1"] = evaluator.evaluate(predict_test)
    elif problem_type == 3: # 回归问题
        evaluator = RegressionEvaluator( # 计算 rmse
            labelCol = "label",
            predictionCol = "prediction",
            metricName = 'rmse'
        )
        results["train_rmse"] = evaluator.evaluate(predict_train)
        results["test_rmse"] = evaluator.evaluate(predict_test)

        evaluator = RegressionEvaluator( # 计算 mae
            labelCol = "label",
            predictionCol = "prediction",
            metricName = 'mae'
        )
        results["train_mae"] = evaluator.evaluate(predict_train)
        results["test_mae"] = evaluator.evaluate(predict_test)

    status = 1

    # 接口约定：
    # status: 状态码
    # results: 评估结果
    return {
        "status": status,
        "results": results
    }

# 保存模型
@bp.route('save/', methods = ['POST'])
def save():
    # 变量初始化 & 接收请求数据
    status = 0
    received_data = request.get_json(silent=True)

    # logging
    logging.info("received_data: " + json.dumps(received_data))

    # 应该接收到有:
    # dataset_id, label, features, problem_type, classifier
    # classifier_params, model_name, model_description
    dataset_id = received_data['dataset_id']
    label = received_data['label']
    features = received_data['features']
    problem_type = received_data['problem_type']
    classifier = received_data['classifier']
    classifier_params = received_data['classifier_params']
    model_name = received_data['model_name']
    model_description = received_data['model_description']

    # 获取 dataframe
    df = getDF(dataset_id)

    # 数据预处理
    df = dataPreProcess(df, label, features, problem_type)

    # 训练
    cf = getClassifier(classifier, classifier_params)
    model = cf.fit(df)

    db = get_db()
    c = db.cursor()

    #####################
    # 在 hdfs 中保存模型 #
    #####################

    # 通过 dataset_id 获取 project_id
    t = (dataset_id,)
    c.execute("SELECT project_id FROM dataset WHERE id=?", t)
    res = c.fetchone()
    project_id = tuple(res)[0]

    # 通过 project_id 获取 username
    t = (project_id,)
    c.execute("SELECT username FROM project WHERE id=?", t)
    res = c.fetchone()
    username = tuple(res)[0]

    # 拼接出 model_path 并保存模型
    model_path = "/user/" + username + "/" + str(project_id) + "/" + urllib.parse.quote(model_name)
    model.save(model_path)

    ####################
    # 在数据库中进行记录 #
    ####################

    t = (project_id, model_name, model_description, classifier, dataset_id, int(label))
    c.execute("INSERT INTO model (project_id, name, description, type, dataset_id, label_index) VALUES (?, ?, ?, ?, ?, ?)", t)
    db.commit()
    close_db()

    status = 1

    # 接口约定：
    # status: 状态码
    return {
        "status": status
    }


from sigma.backends.splunk import SplunkBackend
from sigma.rule import SigmaRule
from sigma.collection import SigmaCollection
from sigmaiq.backends.sigmaiq_abstract_backend import AbstractGenericSigmAIQBackendClass
from copy import copy, deepcopy
import re
from importlib_resources import files

SAVEDSEARCHES_TEMPLATE = files("sigmaiq.backends.splunk").joinpath("savedsearches_template.txt").read_text()

class SigmAIQSplunkBackend(AbstractGenericSigmAIQBackendClass, SplunkBackend):
    """SigmAIQ backend interface for the pySigma Splunk Backend library to translate a SigmaRule object
    to a Splunk search query"""

    custom_formats = {"stanza": "Enterprise Security savedsearches.conf stanza"}
    associated_pipelines = ["splunk_windows", "splunk_wineventlog", "splunk_windows_sysmon_acc", "splunk_cim_dm"]
    default_pipeline = "splunk_windows"

    def handle_output_format(self, sigma_rule, output):
        """Converts Splunk search to savedsearches.conf stanza if output_format given"""
        if self.custom_output_format == "stanza":
            output = self._convert_to_stanza(sigma_rule, output)
        return output

    def _convert_to_stanza(self, sigma_rule, output: list):
        """Converts a list of SigmaRule Splunk query outputs to savedsearches.conf stanzas to use in
        Splunk Enterprise Security Correlation Searches
        """
        # Convert SigmaRule to SigmaCollection just for ease of use
        sigma_rule_temp = copy(sigma_rule)
        if isinstance(sigma_rule_temp, SigmaRule):
            collection = SigmaCollection([sigma_rule_temp])
        else:
            collection = sigma_rule_temp

        stanzas = []
        for i, rule in enumerate(collection):
            title = rule.title
            search = output[i]
            tags = self._extract_mitre_tags(rule.tags)
            replacements = {
                "%TITLE_PLACEHOLDER%": title,
                "%MITRE_ATTACK_PLACEHOLDER%": f"{tags}",
                "%NOTABLE_TITLE_PLACEHOLDER%": title,
                "%DRILLDOWN_TITLE_PLACEHOLDER%": title,
                "%DRILLDOWN_SEARCH_PLACEHOLDER%": search,
                "%SEARCH_PLACEHOLDER%": search,
            }

            stanza = deepcopy(SAVEDSEARCHES_TEMPLATE)
            for k, v in replacements.items():
                stanza = stanza.replace(k, v)
            stanzas.append(stanza)
        return stanzas

    @staticmethod
    def _extract_mitre_tags(tags):
        """Takes tags from a SigmaRule and outputs them into a format that a Splunk Correlation Search stanza will
        recognize for the action.correlationsearch.annotations field:
        i.e.    action.correlationsearch.annotations = {"mitre_attack":["T1055.011","T1560.001","T1205.002","T1583"]}
        """
        mitre_tags = {"mitre_attack": []}
        if not tags:
            return mitre_tags
        for tag in tags:
            matches = re.match(r"^(attack\.)(t\d{4}(\.\d{3})?)", str(tag))
            if matches:
                mitre_tags["mitre_attack"].append(matches[2].upper())
        return mitre_tags


from __future__ import annotations

from dataclasses import asdict, dataclass
from typing import TYPE_CHECKING

from sec_parser.semantic_elements.abstract_semantic_element import (
    AbstractSemanticElement,
)

if TYPE_CHECKING:
    from sec_parser.parsing_engine.html_parsers.html_tag import HtmlTag

class HighlightedTextElement(AbstractSemanticElement):
    """
    The HighlightedElement class used for detecting title elements.
    The process begins with the detection of highlighted elements,
    which are then further classified into title elements.
    """

    def __init__(
        self,
        html_tag: HtmlTag,
        inner_elements: list[AbstractSemanticElement],
        *,
        style: TextStyle | None = None,
    ) -> None:
        super().__init__(html_tag, inner_elements)
        if style is None:
            msg = "styles must be specified for HighlightedElement"
            raise ValueError(msg)
        self.style = style

    @classmethod
    def convert_from(
        cls,
        source: AbstractSemanticElement,
        *,
        style: TextStyle | None = None,
    ) -> HighlightedTextElement:
        return cls(
            source.html_tag,
            source.inner_elements,
            style=style,
        )

@dataclass(frozen=True)
class TextStyle:
    PERCENTAGE_THRESHOLD = 80
    BOLD_THRESHOLD = 600

    bold_with_font_weight: bool
    italic: bool
    # underline?
    # all-caps?

    def __bool__(self) -> bool:
        return any(asdict(self).values())

    @classmethod
    def from_style_string(
        cls,
        style_string: dict[tuple[str, str], float],
    ) -> TextStyle:
        filtered_styles = {
            (k, v): p for (k, v), p in style_string.items()
            if p >= cls.PERCENTAGE_THRESHOLD
        }

        bold_with_font_weight = any(
            cls._is_bold_with_font_weight(k, v)
            for (k, v) in filtered_styles
        )

        italic = any(
            k == "font-style" and v == "italic"
            for (k, v) in filtered_styles
        )

        return cls(
            bold_with_font_weight=bold_with_font_weight,
            italic=italic,
        )

    @classmethod
    def _is_bold_with_font_weight(cls, key: str, value: str) -> bool:
        if key != "font-weight":
            return False
        if value == "bold":
            return True
        try:
            return int(value) >= cls.BOLD_THRESHOLD
        except ValueError:
            return False


""" Apply air quality modelling EPA functions and return statistics for NO2 specifically

Returns:
    [type]: [description]
"""

import pandas as pd

from .File_Utilties import _read_large_dataset, prepend_header_dataframe
import typing

def _compute_statistics(
    dataframe: pd.DataFrame, exceedances: float, percentile: float, window: float
) -> pd.DataFrame:
    outdf = pd.DataFrame(
        {
            "Max of Sensor": dataframe.max(),
            "Average of Sensor": dataframe.mean(),
            "Number of Exceedances": dataframe[dataframe > exceedances].count(),
            "Average Max Column": dataframe.max().mean(),
            "Max Value": dataframe.values.max(),
            f"{percentile * 100}th Percentile of Sensor": dataframe.quantile(
                percentile
            ),
            f"Max of {percentile * 100}th Percentile": dataframe.quantile(
                percentile
            ).max(),
            f"Average {window} Hour Rolling Average of Sensor": dataframe.rolling(
                window=window
            )
            .mean()
            .mean(),
            f"Max {window} Hour Rolling Average of Sensor": dataframe.rolling(
                window=window
            )
            .mean()
            .max(),
            # f"Max {window} Hour Rolling Average": dataframe.rolling(window=window)
            # .mean()
            # .max()
            # .max(),
        }
    )
    outdf.index.name = "Sensor ID"
    outdf.insert(0, "Sensor ID", outdf.index)
    return outdf

def process(
    header_length: int,
    initial: float,
    exceedance: float,
    background_name: str,
    input_data: str,
    ozone_column_name: str,
    fill_invalid_value: float,
    calc_without_background: bool,
    background_column_name: str,
    ozone_scale: float,
    percentile: float,
    window: int,
    top_header_length: int
) -> typing.Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame, pd.DataFrame]:
    """ Apply EPA modelling functions for NO2 and generate statistics

    Args:
        header_length (int): Number of columns before data starts (eg 3 for year/month/day)
        initial (float): Initial percentage to work with (eg 0.1 = 10%)
        exceedance (float): How many exceedances to compute (eg, compare how many are greater than 246)
        background_name (str): File path to background NO2 data set
        input_data (str): Source data file path
        ozone_column_name (str): Column name of ozone data in background data set
        fill_invalid_value (float): Value to fill any missing values
        calc_without_background (bool): Whether to calculate the same statistics for the data set without the background summed
        background_column_name (str): Column name of background data in background data set (eg, Background NO2)
        ozone_scale (float): Number to scale the ozone values by (default 46/48)
        top_header_length (int): Number of rows to ignore before data begins

    Returns:
        typing.Tuple[pd.DataFrame, pd.DataFrame,pd.DataFrame]: Temporary computed data set, output statistics on computed data, output background statistics
    """

    print(f"Processing NO2 statistics for {input_data}")

    # Read Data
    background = _read_large_dataset(background_name)
    background.dropna(how='all', inplace=True)

    data = _read_large_dataset(input_data)

    data_header = data.iloc[:, :header_length]

    data = data.iloc[:, header_length:]

    data = data[top_header_length - 1:]

    # Multiply background data by provided scalar

    data["background_NO2"] = background[background_column_name]

    data["background_O3"] = background[ozone_column_name]

    print("25%")

    data = data.fillna(fill_invalid_value)

    def olmfunction(row: pd.Series) -> pd.Series:
        """ Compute OLM which is the recorded value multiplied by the initial amount, and add the minimum of (1-iniial)*value or the ozone)

        Args:
            row (pd.Series): Row to be computed where last column is the ozone data scaled

        Returns:
            pd.Series: Computed row
        """
        # Take the ozone value
        ozone_value = float(row.iloc[-1]) * ozone_scale
        background_value = float(row.iloc[-2])
        new_row = []

        # Skip header rows and ignore last 2 values
        for value in row[:-2]:
            value = (float(value) * initial) + min(((1 - initial) * float(value)), ozone_value) + background_value
            new_row.append(value)
        return pd.Series(new_row)

    def backfunction(row: pd.Series) -> pd.Series:
        """ Background function for computing sum of recorded value and background

        Args:
            row (pd.Series): Row to be computed where last column is the background data

        Returns:
            pd.Series: Computed row
        """
        # Take the background value
        background_value = float(row.iloc[-1])
        new_row = []
        for value in row[:-1]:
            value = value - background_value
            new_row.append(value)

        return pd.Series(new_row)

    olm_data = data.apply(olmfunction, axis=1)
    olm_data.columns = data.columns[:-2]
    olm_data_with_background = olm_data.copy()
    olm_data["background"] = list(background[background_column_name])
    olm_data_without_background = olm_data.apply(backfunction, axis=1)
    olm_data_without_background.columns = data.columns[:-2]

    print("50%")
    # Compute statistics
    outdf = _compute_statistics(
        olm_data_with_background, exceedance, percentile, window
    )
    print("75%")
    no_bg_outdf = pd.DataFrame
    if calc_without_background:
        # data["background"] = list(background[background_column_name])
        # bg_data = data.apply(backfunction, axis=1)
        no_bg_outdf = _compute_statistics(
            olm_data_without_background, exceedance, percentile, window
        )

    olm_data_with_background = prepend_header_dataframe(
        olm_data_with_background, data_header
    )
    olm_data_without_background = prepend_header_dataframe(
        olm_data_without_background, data_header
    )

    print("100%")

    return olm_data_with_background, olm_data_without_background, outdf, no_bg_outdf


import sqlite3 as sql

banco = 'banco.bd'

# Função
def abrir_conexao(banco):
    conexao = sql.connect(banco)
    cursor = conexao.cursor()
    return conexao, cursor

def fechar_conexao(conexao):
    conexao.commit()
    conexao.close()

# Comandos SQL
contagem = "SELECT COUNT(*) FROM produtos;"
select_todos = "SELECT * FROM produtos;"
select_nome = "SELECT * FROM produtos WHERE nome = ?;"
insert = "INSERT INTO produtos VALUES (:nome,:preco,:qtd);"
delete = "DELETE FROM produtos WHERE nome = ?;"
update = '''
UPDATE produtos SET
    qtd = :qtd
WHERE nome like :nome
'''

def insert_produto(produtos):
    conexao, cursor = abrir_conexao(banco)
    cursor.execute(insert, produtos)
    fechar_conexao(conexao)

def func_leitura():
    conexao, cursor = abrir_conexao(banco)
    cursor.execute(select_todos)
    resultado = cursor.fetchall()
    fechar_conexao(conexao)
    return resultado

def func_delete(nome):
    conexao, cursor = abrir_conexao(banco)
    resultado = cursor.execute(delete, [nome]).rowcount
    fechar_conexao(conexao)

def func_read(nome):
    conexao, cursor = abrir_conexao(banco)
    resultado = cursor.execute(select_nome, [nome]).fetchall()
    fechar_conexao(conexao)
    return resultado

def func_update(nome, produtos):
    conexao, cursor = abrir_conexao(banco)
    resp = {"nome": nome, "qtd": produtos["qtd"]}
    cursor.execute(update, resp)
    fechar_conexao(conexao)
    return resp


'''
Estimate the SNR of seismic signals.

Autor: Peidong Shi
Contact: speedshi@hotmail.com  or  peidong.shi@sed.ethz.ch
Create time: 20221110
'''

from obspy import UTCDateTime
import numpy as np

def estimate_snr(trace, stime, noise_window, signal_window, method='maxamp'):
    # trace is one or three component seismic data recroded at a certain station; no gaps in traces;
    # 'stime' is the picked signal arrival-time (phase arrival-time) in datetime format, such as datetime.datetime(2022,2,3,15,21,30,345332).
    # 'noise_window' specify the time window in second relative to 'stime' for evaluating noise, such as [-4, -2].
    # 'signal_window' specify the time window in second relative to 'stime' for evaluating signal, such as [-0.4, 0.6].
    # 'method' specify how to calculate snr; 'maxamp': maximum amplitude (on an arbitrary channel) ratio;
    #                                        'maxeng': maximum energy ratio (need more than 2 channel);
    #                                         'std': standard deviation ratio;

    if trace.count() == 0:
        # no input seismic data
        return None
    else:
        assert(trace.count()<=3)

        # check traces are coming from the same station
        for itrace in trace:
            assert(itrace.stats.station == trace[0].stats.station)

        stime = UTCDateTime(stime)

        assert(len(noise_window)==2 and (noise_window[0]<=noise_window[1]))
        assert(len(signal_window)==2 and (signal_window[0]<=signal_window[1]))

        noise_start = stime + noise_window[0]
        noise_end = stime + noise_window[1]
        noises = (trace.copy()).trim(starttime=noise_start, endtime=noise_end, pad=True, fill_value=0)  # noise segments

        signal_start = stime + signal_window[0]
        signal_end = stime + signal_window[1]
        signals = (trace.copy()).trim(starttime=signal_start, endtime=signal_end, pad=True, fill_value=0)  # signal segments

        if method.lower() == 'maxeng':
            if noises.count() == 1:
                maxeng_noise = max(np.sqrt(noises[0].data*noises[0].data))
            elif noises.count() == 2:
                maxeng_noise = max(np.sqrt(noises[0].data*noises[0].data + noises[1].data*noises[1].data))
            elif noises.count() == 3:
                maxeng_noise = max(np.sqrt(noises[0].data*noises[0].data + noises[1].data*noises[1].data + noises[2].data*noises[2].data))
            else:
                raise ValueError

            if signals.count() == 1:
                maxeng_signal = max(np.sqrt(signals[0].data*signals[0].data))
            elif signals.count() == 2:
                maxeng_signal = max(np.sqrt(signals[0].data*signals[0].data + signals[1].data*signals[1].data))
            elif signals.count() == 3:
                maxeng_signal = max(np.sqrt(signals[0].data*signals[0].data + signals[1].data*signals[1].data + signals[2].data*signals[2].data))
            else:
                raise ValueError

            snr = maxeng_signal / maxeng_noise
        elif method.lower() == 'maxamp':
            maxamp_noise = np.amax(np.absolute(noises.max()))
            maxamp_signal = np.amax(np.absolute(signals.max()))
            snr = maxamp_signal / maxamp_noise
        elif method.lower() == 'std':
            std_noise = np.amax(noises.std())
            std_signal = np.amax(signals.std())
            snr = std_signal / std_noise
        else:
            raise ValueError("Unrecognized input for method: {method}.")

        del noises, signals
    return snr


from typing import AnyStr

from lxml.etree import strip_elements
from lxml.html import HTMLParser
from lxml.html import fromstring

USELESS_TAGS = frozenset(
    {
        "head",
        "script",
        "style",
        "iframe",
        "noscript",
    },
)

def get_root(
    html: AnyStr,
    *,
    useless_tags: list[str] | None = None,
    remove_comments: bool = True,
):
    """Get root of DOM Tree without useless data"""

    parser = HTMLParser(remove_comments=remove_comments)
    h = fromstring(html, parser=parser)
    if useless_tags is None:
        useless_tags = list(USELESS_TAGS)
    strip_elements(h, *useless_tags, with_tail=False)
    return h

def strip_tags(text: str) -> str:
    """Remove HTML Tags from input text"""

    return str(fromstring(text).text_content())


from matplotlib import pyplot as plt
import numpy as np
import itertools
import random
import dataset
import networkx as nx
import os
#import main

def time_lasted(t):
    if t <= 60:
        return f"{round(t, 3)}s"
    return f"{round(t / 60, 3)}min"

def make_random_graph(n_nodes):
    # Create nodes
    nodes = [str(i) for i in range(n_nodes)]

    # Create edges
    edges = list(itertools.combinations(range(n_nodes), 2))
    n_edges = len(edges)
    edge_weights = [random.randint(1,10) for _ in edges] # use random ints for edge weights for now, to be changed
    edge_list = [(edges[i][0], edges[i][1], edge_weights[i]) for i in range(n_edges)]

    # Add edges to the graph (also adds nodes)
    edge_labels = {}
    for i in range(len(edges)):
        key = edges[i]
        edge_labels[key] = edge_weights[i]

    graph = {
        "n_node": n_nodes,
        "n_edge": n_edges,
        "nodes": nodes,
        "edges": edge_list,
        "edge_labels": edge_labels
        }
    return graph

def updated_graph(nodes, edges):
    """
    Return a new graph based on given nodes and find corresponding edges.
    """
    n_nodes = len(nodes)
    new_edges = list(itertools.combinations(nodes, 2))
    n_edges = len(new_edges)

    edge_list = []
    edge_labels = {}
    for e in new_edges:
        edge_labels[e] = edges[e]
        edge_list.append((e[0], e[1], edges[e]))

    graph = {
        "n_node": n_nodes,
        "n_edge": n_edges,
        "nodes": nodes,
        "edges": edge_list,
        "edge_labels": edge_labels
        }
    return graph

def display_result(results):
    """
    Reform results into a list of dicts to be written into files.
    Input:
        results: list of course groups
    Return:
        list of dicts
    """
    final_result = []
    for i in range(len(results)):
        group = {f'Group {i+1}': results[i]}
        final_result.append(group)
    return final_result

def plot_log(logs, out_img, out_file):
    """
    Plot graph size vs run time.
    """
    ax = plt.subplots(figsize=(6, 6))[1]
    n_node = np.array([abs(int(i[0])) for i in logs], dtype=np.int8)
    run_time = np.array([i[1] for i in logs], dtype=np.float32)

    csv = dataset._CSV()
    result = [{"Number of Nodes": n_node, "Time (s)": run_time}]
    csv.write(result, out_file)

    ax.scatter(n_node, run_time)
    ax.set_title("Run Time vs Graph Size")
    ax.set_ylabel("Time (s)", fontsize=12)
    ax.set_xlabel("Number of Nodes", fontsize=12)
    plt.tight_layout()
    plt.show()
    plt.savefig(out_img, bbox_inches='tight')

def objective_value(result, nodes, edge_weights):
    """
    Calculate the final objective value optimized.

    Input:
        result: final result groups
        nodes: the nodes from the original graph
        edge_weights: a dictionary with the initial edge weights
    Return:
        Optimized final value and original total weights.
    """
    all_weights = sum([edge_weights[e] for e in edge_weights])
    final_value = all_weights
    for group in result:
        new_group = [nodes.index(i) for i in group]
        edges = list(itertools.combinations(new_group, 2))
        for e in edges:
            try:
                final_value -= edge_weights[e]
            except KeyError:
                final_value -= edge_weights[(e[1], e[0])]
    return all_weights, final_value

def plot_graph(graph, name, output_dir, color_list='0.1', g=None):
    edges = graph["edges"]
    #edge_labels = graph["edge_labels"]
    nodes = graph["nodes"]

    if g == None:
        g = nx.Graph()
        node_list = list(range(len(nodes)))
        g.add_nodes_from(node_list)
        g.add_weighted_edges_from(edges)

    plt.figure(figsize=(5,5))
    pos = nx.spring_layout(g)
    nx.draw_networkx_nodes(g, pos, node_size=15, node_color=color_list)
    nx.draw_networkx_edges(g, pos, alpha=0.05, edge_color='0.3')
    filename = f"{name}.png"
    out_name = os.path.join(output_dir, filename)
    print(f">>> The plot {filename} is saved to {out_name}")
    plt.savefig(out_name, bbox_inches='tight')

def plot_final_graph(groups, graph, output_dir, g=None):
    nodes = graph["nodes"]
    color_list = list(nodes)
    color_order = ["r", "c"]
    new_edges = []
    i = 0
    for group in groups:
        new_group = []
        #r, g, b = random.randint(1, 255), random.randint(1, 255), random.randint(1,255)
        for n in group:
            j = nodes.index(n)
            #color_list[j] = (r/255, g/255, b/255)
            color_list[j] = color_order[i]
            new_group.append(j)
        e = list(itertools.combinations(new_group, 2))
        new_edges += e
        i += 1

    edges = graph["edge_labels"]
    edge_list = []
    edge_labels = {}
    for e in new_edges:
        edge_labels[e] = edges[e]
        edge_list.append((e[0], e[1], edges[e]))

    new_graph = {
        "n_node": len(nodes),
        "n_edge": len(new_edges),
        "nodes": nodes,
        "edges": edge_list,
        "edge_labels": edge_labels
        }

    plot_graph(new_graph, "Output Graph", output_dir, color_list=color_list, g=g)
    return new_graph

def calc_mass(Pt1, Pt2, eta1, eta2, phi1, phi2):
    return np.sqrt(2*Pt1*Pt2*(np.cosh(eta1-eta2)-np.cos(phi1-phi2)))

def get_jet_info(event, idx):
    return [event.JetPt[idx], event.JetEta[idx], event.JetPhi[idx]]

def get_track_info(event, idx):
    return [event.TrackPt[idx], event.TrackEta[idx], event.TrackPhi[idx]]

def get_tower_info(event, idx):
    return [event.JetTowerEt[idx], event.JetTowerEta[idx], event.JetTowerPhi[idx]]

def cos_similarity(a, b):
    """
    Return the negative value of inner product.
    """
    return -np.dot(np.linalg.norm(a), np.linalg.norm(b))

def euclidean_distance(a, b):
    """
    Return root mean square error.
    """
    a, b = np.array(a), np.array(b)
    return np.linalg.norm(a - b)

def abe(a, b):
    """
    Return the absolute mean error.
    """
    a, b = np.array(a), np.array(b)
    return np.mean(np.abs(a - b))

def rename_file(file_path, overwrite=False):
    file_name = file_path[:-4]
    file_type = file_path[-4:]
    count = 0
    while os.path.exists(file_path) and not overwrite:
        file_path = f"{file_name}_copy{count}{file_type}"
        count += 1
    return file_path


from . import db
from .data_mapper import map_penalty_records, map_to_game_summary
from .models import UserEntity, GameEntity, GameStatus, PenaltyEntity, PenaltyRecordEntity, ParticipantStatus, ParticipantEntity, TotalFineEntity
from flask import Blueprint, render_template, request, flash, redirect, url_for, abort
from flask_login import login_required, current_user
import json

from .logger_config import setup_logger

DEBUG = True

class PlayerListError(ValueError):
    pass

logger = setup_logger()

views = Blueprint('views', __name__) # setup a blueprint for flask, can be named w.e. best practice to name same as file

# Display homepage endpoint
@views.route('/') # decorator
@login_required
def home():
    return render_template("home.html", user=current_user)

# Display user_management site
@views.route('/user_management', methods=['GET'])
@login_required
def user_management():
    participants = ParticipantEntity.query.filter_by(user_id=current_user.id, status = ParticipantStatus.active)
    return render_template("user_management.html", participants=participants, user=current_user)

# Display Game creation and loading site
@views.route('/create_or_load_game/', methods=['GET'])
@login_required
def create_or_load_game():
    participants = ParticipantEntity.query.filter_by(user_id=current_user.id, status = ParticipantStatus.active)
    all_saved_games = GameEntity.query.filter_by(user_id = current_user.id, status = GameStatus.active).all()
    len_all_saved_games = len(all_saved_games)
    logger.debug("Length of all saved_games for user {} is {}".format(current_user.email, len_all_saved_games))
    return render_template("create_or_load_game.html", game_participants=participants, user=current_user, all_saved_games = all_saved_games, len_all_saved_games=len_all_saved_games)

# Display certain game
@views.route('/view_game/<game_id>', methods=['GET'])
@login_required
def view_game(game_id):
    try:
        game = GameEntity.query.filter_by(id = game_id,user_id = current_user.id).one()
    except:
        flash("No game was found", category="error")
        return render_template("create_or_load_game.html",user=current_user)
    # game_participants is a json object storing all the player_ids as a string
    all_penalties = PenaltyEntity.query.filter_by(user_id=current_user.id).all()
    game_participants = game.game_participants
    if game_participants is None or game_participants == "":
        flash("invalid game", category="error")
        return render_template("create_or_load_game.html", user=current_user)
    logger.info("Game Status for user {} is {}".format(current_user.email,game.status))
    all_participant_ids = []
    if game.status == GameStatus.active:
        participants_data = []
        # Iterate through all participant ids
        # Get all their PenaltyRecordEntities from database
        # Get their TotalFine from database
        for player in json.loads(game_participants):
            all_participant_ids.append(int(player))
            participant_record = PenaltyRecordEntity.query.filter_by(game_id=game.id, participant_id=player).first()
            participant_fine_sum = TotalFineEntity.query.filter_by(game_id=game.id, participant_id=player).first()
            # Map the Database Entities to useful data_models, a single object that stores everything needed for the player
            # (name, quantity for each penalty, the total_finesum and more)
            # All stored in a list and given to the template
            participants_data.extend(map_penalty_records(participant_record, game_id,all_penalties, participant_fine_sum))
    if game.status == GameStatus.save:
        participants_data = []
        for player in json.loads(game_participants):
            all_participant_ids.append(player)
            participant_records = PenaltyRecordEntity.query.filter_by(game_id=game.id, participant_id=player).all()
            participant_fine_sum = TotalFineEntity.query.filter_by(game_id=game.id, participant_id=player).first()
            participants_data.append(map_to_game_summary(participant_records, participant_fine_sum))
        logger.debug("All participant_ids going into the game: ", all_participant_ids)
    return render_template("view_game.html",user=current_user, game=game, all_penalties=all_penalties, player_record_list = participants_data, all_participant_ids=all_participant_ids)

# Display finished game summary
@views.route('/game_summary/<game_id>', methods=['POST','GET'])
@login_required
def game_summary(game_id):
    if request.method != 'POST' and request.method !='GET':
        abort(405)
    # Create a game summary object and map everything so far obtained to it in the finish_game endpoint
    all_penalties = PenaltyEntity.query.filter_by(user_id = current_user.id).all()
    target_game = GameEntity.query.filter_by(id = game_id, user_id = current_user.id).first()
    if target_game and target_game.status == GameStatus.finish:
        game_participants = target_game.game_participants
        game_summary_objects = []
        for player in json.loads(game_participants):
            participant_records = PenaltyRecordEntity.query.filter_by(game_id=target_game.id, participant_id=player).all()
            participant_fine_sum = TotalFineEntity.query.filter_by(game_id=target_game.id, participant_id=player).first()
            game_summary_objects.append(map_to_game_summary(participant_records,participant_fine_sum))
        tier_list = retrieve_tier_list(game_summary_objects)
        penalty_sum = 1 + len(all_penalties)
        return render_template("game_summary.html", user= current_user, game = target_game, all_penalties = all_penalties, game_summary_objects = game_summary_objects, tier_list = tier_list, penalty_sum = penalty_sum)
    else:
        flash("This game has not been finished yet", category="error")

def retrieve_tier_list(PenaltyRecordList):
    max_quantities = {}
    for record in PenaltyRecordList:
        for penalty in record.penalties:
            if penalty.penalty_name not in max_quantities or max_quantities[penalty.penalty_name][1] < penalty.penalty_quantity:
                max_quantities[penalty.penalty_name] = (record.participant_name, penalty.penalty_quantity)
    return max_quantities

@views.route('/penalties', methods=['GET'])
@login_required
def penalties():
    penalties = PenaltyEntity.query.filter_by(user_id=current_user.id)
    return render_template("penalties.html", user=current_user, all_penalties = penalties)

@views.route('/view_archive', methods=['GET'])
@login_required
def view_archive():
    all_finished_games = GameEntity.query.filter_by(user_id = current_user.id, status = 'finish').all()
    return render_template("view_archive.html", user=current_user, all_finished_games = all_finished_games)


from dataclasses import dataclass, replace, field
from pprint import pprint
import abc
import typing
import torch

@dataclass(frozen=True)
class SequenceVariableValue:
    type = "sequence"
    value: torch.Tensor

@dataclass(frozen=True)
class MatrixVariableValue:
    type = "matrix"
    value: torch.Tensor

Variable = SequenceVariableValue | MatrixVariableValue

@dataclass(frozen=True)
class ProgramState:
    tokens: torch.Tensor
    sequence_variables: dict[str, SequenceVariableValue] = field(default_factory=dict)
    matrix_variables: dict[str, MatrixVariableValue] = field(default_factory=dict)

    def assign_sequence(
        self, variable_name: str, value: SequenceVariableValue
    ) -> "ProgramState":
        assert (
            variable_name not in self.sequence_variables
            and variable_name not in self.matrix_variables
        )
        return replace(
            self, sequence_variables={**self.sequence_variables, variable_name: value}
        )

    def assign_matrix(
        self, variable_name: str, value: MatrixVariableValue
    ) -> "ProgramState":
        assert (
            variable_name not in self.sequence_variables
            and variable_name not in self.matrix_variables
        )
        return replace(
            self, matrix_variables={**self.matrix_variables, variable_name: value}
        )

T = typing.TypeVar("T", bound=(SequenceVariableValue | MatrixVariableValue))

class Expression(abc.ABC, typing.Generic[T]):
    @abc.abstractmethod
    def evaluate(self, state: ProgramState) -> T:
        pass

SequenceExpression = Expression[SequenceVariableValue]
MatrixExpression = Expression[MatrixVariableValue]

Predicate = typing.Callable[[torch.Tensor, torch.Tensor], torch.Tensor]

def eq_predicate(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    return a == b

def lt_predicate(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    return a < b

def gt_predicate(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    return a > b

def leq_predicate(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    return a <= b

def geq_predicate(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    return a >= b

@dataclass(frozen=True)
class SelectExpression(MatrixExpression):
    k: SequenceExpression
    q: SequenceExpression
    p: Predicate

    def evaluate(self, state: ProgramState) -> MatrixVariableValue:
        return MatrixVariableValue(
            value=self.p(
                self.k.evaluate(state).value.unsqueeze(0),
                self.q.evaluate(state).value.unsqueeze(1),
            )
        )

@dataclass(frozen=True)
class SelectorWidthExpression(SequenceExpression):
    m: MatrixExpression

    def evaluate(self, state: ProgramState) -> SequenceVariableValue:
        return SequenceVariableValue(self.m.evaluate(state).value.sum(dim=1))

@dataclass(frozen=True)
class NotSequenceExpression(SequenceExpression):
    s: SequenceExpression

    def evaluate(self, state: ProgramState) -> SequenceVariableValue:
        return SequenceVariableValue(~self.s.evaluate(state).value)

@dataclass(frozen=True)
class IndicatorExpression(SequenceExpression):
    s: SequenceExpression

    def evaluate(self, state: ProgramState) -> SequenceVariableValue:
        return SequenceVariableValue((self.s.evaluate(state).value > 0).float())

@dataclass
class NotMatrixExpression(MatrixExpression):
    m: MatrixExpression

    def evaluate(self, state: ProgramState) -> MatrixVariableValue:
        return MatrixVariableValue(value=~self.m.evaluate(state).value)

@dataclass
class AndMatrixExpression(MatrixExpression):
    a: MatrixExpression
    b: MatrixExpression

    def evaluate(self, state: ProgramState) -> MatrixVariableValue:
        return MatrixVariableValue(
            value=self.a.evaluate(state).value & self.b.evaluate(state).value
        )

@dataclass(frozen=True)
class OrSequenceExpression(SequenceExpression):
    a: SequenceExpression
    b: SequenceExpression

    def evaluate(self, state: ProgramState) -> SequenceVariableValue:
        return SequenceVariableValue(
            value=self.a.evaluate(state).value | self.b.evaluate(state).value
        )

@dataclass(frozen=True)
class AndSequenceExpression(SequenceExpression):
    a: SequenceExpression
    b: SequenceExpression

    def evaluate(self, state: ProgramState) -> SequenceVariableValue:
        return SequenceVariableValue(
            value=self.a.evaluate(state).value & self.b.evaluate(state).value
        )

class TrueSequenceExpression(SequenceExpression):
    def evaluate(self, state: ProgramState) -> SequenceVariableValue:
        return SequenceVariableValue(
            value=torch.ones(state.tokens.shape[0], dtype=torch.bool)
        )

class TrueMatrixExpression(MatrixExpression):
    def evaluate(self, state: ProgramState) -> MatrixVariableValue:
        return MatrixVariableValue(
            value=torch.ones(
                state.tokens.shape[0], state.tokens.shape[0], dtype=torch.bool
            )
        )

class TokensExpression(SequenceExpression):
    def evaluate(self, state: ProgramState) -> SequenceVariableValue:
        return SequenceVariableValue(value=state.tokens)

class IndicesExpression(SequenceExpression):
    def evaluate(self, state: ProgramState) -> SequenceVariableValue:
        return SequenceVariableValue(value=torch.arange(state.tokens.shape[0]))

@dataclass
class SequenceReferenceExpression(SequenceExpression):
    variable_name: str

    def evaluate(self, state: ProgramState) -> SequenceVariableValue:
        return state.sequence_variables[self.variable_name]

@dataclass
class MatrixReferenceExpression(MatrixExpression):
    variable_name: str

    def evaluate(self, state: ProgramState) -> MatrixVariableValue:
        return state.matrix_variables[self.variable_name]

Token = float

@dataclass(frozen=True)
class LiteralSequenceExpression(SequenceExpression):
    literal_value: Token

    def evaluate(self, state: ProgramState) -> SequenceVariableValue:
        return SequenceVariableValue(
            value=torch.full((state.tokens.shape[0],), self.literal_value)
        )

class LastColumnExpression(SequenceExpression):
    def evaluate(self, state: ProgramState) -> SequenceVariableValue:
        return SequenceVariableValue(
            value=torch.full((state.tokens.shape[0],), state.tokens.shape[0] - 1)
        )

@dataclass(frozen=True)
class LiteralMatrixExpression(MatrixExpression):
    literal_value: bool

    def evaluate(self, state: ProgramState) -> MatrixVariableValue:
        return MatrixVariableValue(
            value=torch.full(
                (state.tokens.shape[0], state.tokens.shape[0]), self.literal_value
            )
        )

@dataclass(frozen=True)
class AggregateExpression(SequenceExpression):
    m: MatrixExpression
    s: SequenceExpression

    def evaluate(self, state: ProgramState) -> SequenceVariableValue:
        s_value = self.s.evaluate(state).value
        m_value = self.m.evaluate(state).value
        mult = s_value * m_value
        agg = mult.sum(dim=1)
        avg = agg / m_value.sum(dim=1)
        return SequenceVariableValue(value=avg)

@dataclass(frozen=True)
class CompareSequencesExpression(SequenceExpression):
    a: SequenceExpression
    b: SequenceExpression
    p: Predicate

    def evaluate(self, state: ProgramState) -> SequenceVariableValue:
        return SequenceVariableValue(
            value=self.p(
                self.a.evaluate(state).value,
                self.b.evaluate(state).value,
            )
        )

@dataclass(frozen=True)
class CompareMatricesExpression(MatrixExpression):
    a: MatrixExpression
    b: MatrixExpression
    p: Predicate

    def evaluate(self, state: ProgramState) -> MatrixVariableValue:
        return MatrixVariableValue(
            value=self.p(
                self.a.evaluate(state).value,
                self.b.evaluate(state).value,
            )
        )

@dataclass(frozen=True)
class SubtractSequencesExpression(SequenceExpression):
    a: SequenceExpression
    b: SequenceExpression

    def evaluate(self, state: ProgramState) -> SequenceVariableValue:
        return SequenceVariableValue(
            value=self.a.evaluate(state).value - self.b.evaluate(state).value
        )

@dataclass(frozen=True)
class Statement:
    assignment_name: str
    expression: SequenceExpression | MatrixExpression

@dataclass(frozen=True)
class Program:
    statements: tuple[Statement, ...]

    def invoke(self, tokens: torch.Tensor):
        state = ProgramState(tokens=tokens)
        for statement in self.statements:
            result = statement.expression.evaluate(state)
            if isinstance(result, SequenceVariableValue):
                state = state.assign_sequence(statement.assignment_name, result)
            elif isinstance(result, MatrixVariableValue):
                state = state.assign_matrix(statement.assignment_name, result)
            else:
                raise NotImplementedError()
        print(result)

double_histogram_program = Program(
    statements=(
        Statement(
            "same_tok",
            SelectExpression(TokensExpression(), TokensExpression(), eq_predicate),
        ),
        Statement(
            "hist", SelectorWidthExpression(MatrixReferenceExpression("same_tok"))
        ),
        Statement(
            "is_prev",
            SelectExpression(IndicesExpression(), IndicesExpression(), lt_predicate),
        ),
        Statement(
            "has_prev",
            CompareSequencesExpression(
                SelectorWidthExpression(
                    AndMatrixExpression(
                        MatrixReferenceExpression("is_prev"),
                        MatrixReferenceExpression("same_tok"),
                    ),
                ),
                LiteralSequenceExpression(0),
                gt_predicate,
            ),
        ),
        Statement(
            "first", NotSequenceExpression(SequenceReferenceExpression("has_prev"))
        ),
        Statement(
            "same_count",
            SelectExpression(
                SequenceReferenceExpression("hist"),
                SequenceReferenceExpression("hist"),
                eq_predicate,
            ),
        ),
        Statement(
            "first_mat",
            SelectExpression(
                SequenceReferenceExpression("first"),
                TrueSequenceExpression(),
                eq_predicate,
            ),
        ),
        Statement(
            "same_count_reprs",
            AndMatrixExpression(
                MatrixReferenceExpression("same_count"),
                MatrixReferenceExpression("first_mat"),
            ),
        ),
        Statement(
            "hist2",
            SelectorWidthExpression(MatrixReferenceExpression("same_count_reprs")),
        ),
    )
)

double_histogram_program.invoke(torch.tensor([1, 1, 1, 2, 2, 3, 3, 4, 5, 6]))

def frac_prevs(sop: SequenceExpression, val: Token) -> SequenceExpression:
    prevs = SelectExpression(IndicesExpression(), IndicesExpression(), leq_predicate)
    return AggregateExpression(
        prevs,
        IndicatorExpression(
            CompareSequencesExpression(
                sop, LiteralSequenceExpression(val), eq_predicate
            )
        ),
    )

def pair_balance(open_token: Token, close_token: Token) -> SequenceExpression:
    opens = frac_prevs(TokensExpression(), open_token)
    closes = frac_prevs(TokensExpression(), close_token)
    return SubtractSequencesExpression(opens, closes)

OPEN_A_TOKEN: Token = 1.0
CLOSE_A_TOKEN: Token = 2.0
OPEN_B_TOKEN: Token = 3.0
CLOSE_B_TOKEN: Token = 4.0

shuffle_dyck_2_program = Program(
    statements=(
        Statement("bal1", pair_balance(OPEN_A_TOKEN, CLOSE_A_TOKEN)),
        Statement("bal2", pair_balance(OPEN_B_TOKEN, CLOSE_B_TOKEN)),
        Statement(
            "negative",
            OrSequenceExpression(
                CompareSequencesExpression(
                    SequenceReferenceExpression("bal1"),
                    LiteralSequenceExpression(0),
                    lt_predicate,
                ),
                CompareSequencesExpression(
                    SequenceReferenceExpression("bal2"),
                    LiteralSequenceExpression(0),
                    lt_predicate,
                ),
            ),
        ),
        Statement(
            "had_neg",
            CompareSequencesExpression(
                AggregateExpression(
                    TrueMatrixExpression(),
                    IndicatorExpression(SequenceReferenceExpression("negative")),
                ),
                LiteralSequenceExpression(0),
                gt_predicate,
            ),
        ),
        Statement(
            "select_last",
            SelectExpression(IndicesExpression(), LastColumnExpression(), eq_predicate),
        ),
        Statement(
            "end_0",
            AggregateExpression(
                MatrixReferenceExpression("select_last"),
                AndSequenceExpression(
                    CompareSequencesExpression(
                        SequenceReferenceExpression("bal1"),
                        LiteralSequenceExpression(0),
                        eq_predicate,
                    ),
                    CompareSequencesExpression(
                        SequenceReferenceExpression("bal2"),
                        LiteralSequenceExpression(0),
                        eq_predicate,
                    ),
                ),
            ),
        ),
        Statement(
            "shuffle_dyck_2",
            AndSequenceExpression(
                CompareSequencesExpression(
                    SequenceReferenceExpression("end_0"),
                    LiteralSequenceExpression(1),
                    eq_predicate,
                ),
                NotSequenceExpression(SequenceReferenceExpression("had_neg")),
            ),
        ),
    )
)

shuffle_dyck_2_program.invoke(
    torch.tensor([OPEN_A_TOKEN, OPEN_B_TOKEN, CLOSE_A_TOKEN, CLOSE_B_TOKEN])
)

shuffle_dyck_2_program.invoke(
    torch.tensor(
        [OPEN_A_TOKEN, OPEN_B_TOKEN, CLOSE_A_TOKEN, CLOSE_B_TOKEN, CLOSE_B_TOKEN]
    )
)


import argparse
import logging
import dataset
from model import ListNet

logging.basicConfig(level=logging.INFO)
"""
    entrance
"""

def run(args):
    logging.info("Loading data...")
    id, x_train, x_test, y_train, y_test = dataset.load_random_data()
    logging.info("loading {} id".format(len(id)))
    logging.info("loading {} items data for training, dims {}".format(x_train.shape[0], x_train.shape[1]))
    logging.info("loading {} items data as training label, dims {}".format(y_train.shape[0], y_train.shape[1]))
    logging.info("loading {} items data for testing, dims {}".format(x_test.shape[0], x_test.shape[1]))
    logging.info("loading {} items data as testing result, dims {}".format(y_test.shape[0], y_test.shape[1]))
    listnet = ListNet(h1_input=x_train.shape[1],
                      h1_output=10,
                      h2_input=10,
                      h2_output=1)
    listnet.train(x_train, y_train, 100, 0.001, optmz='adam')
    listnet.predict(x_test)
    perms=listnet.permutation(id)
    print(perms)

if __name__ == '__main__':
    p = argparse.ArgumentParser()
    p.add_argument('--train',
                   type=str,
                   help='input train json file',
                   default='random')
    p.add_argument('--env',
                   required=False,
                   help='input dev/prod to alter environment',
                   default='dev')
    p.print_help()
    args = p.parse_args()
    print(args)
    run(args)


from django.urls import path

from . import views

urlpatterns = [
    path('settings/personal/',
         views.user_settings_personal,
         name='user-settings-personal'),
    path('settings/security/',
         views.user_settings_security,
         name='user-settings-security'),
    path('settings/notifications/',
         views.user_settings_notifications,
         name='user-settings-notifications'),
    path('settings/change-email/',
         views.change_email,
         name='change-email'),
    path('settings/change-email/confirm/',
         views.change_email_request_confirmation,
         name='change-email-request-confirmation'),
    path('settings/change-email/confirm/<uid>/<token>/<new_email>/',
         views.change_email_confirm,
         name='change-email-confirm'),
]


from argparse import ArgumentParser
import  os

parser = ArgumentParser()

model_path = os.path.join('/models/hubert/', 'hubert.pt')

parser.add_argument('--train_dataset', default='/train/dataset')
parser.add_argument('--process', default='/train/Literature/prepared')
parser.add_argument('--train', default='/train/Literature/ready')

parser.add_argument('--hubert-model', default='./models/hubert/hubert.pt', help='The hubert model to use for preparing the data and later creation of semantic tokens.')
parser.add_argument('--train-save-epochs', default=1, type=int, help='The amount of epochs to train before saving')

args = parser.parse_args()