# -*- coding: utf-8 -*-
"""BE Proj.ipynb

 Suspicious Link Detection

"""

import warnings
warnings.filterwarnings('ignore')

import numpy as np
import pandas as pd
import re
import nltk

############## Model 1 new########################
import joblib
import numpy as np

import re
#Use of IP or not in domain
def having_ip_address(url):
    match = re.search(
        '(([01]?\\d\\d?|2[0-4]\\d|25[0-5])\\.([01]?\\d\\d?|2[0-4]\\d|25[0-5])\\.([01]?\\d\\d?|2[0-4]\\d|25[0-5])\\.'
        '([01]?\\d\\d?|2[0-4]\\d|25[0-5])\\/)|'  # IPv4
        '((0x[0-9a-fA-F]{1,2})\\.(0x[0-9a-fA-F]{1,2})\\.(0x[0-9a-fA-F]{1,2})\\.(0x[0-9a-fA-F]{1,2})\\/)' # IPv4 in hexadecimal
        '(?:[a-fA-F0-9]{1,4}:){7}[a-fA-F0-9]{1,4}', url)  # Ipv6
    if match:
        # print match.group()
        return 1
    else:
        # print 'No matching pattern found'
        return 0

from urllib.parse import urlparse

def abnormal_url(url):
    hostname = urlparse(url).hostname
    hostname = str(hostname)
    match = re.search(hostname, url)
    if match:
        # print match.group()
        return 1
    else:
        # print 'No matching pattern found'
        return 0

def count_dot(url):
    count_dot = url.count('.')
    return count_dot

def count_www(url):
    url.count('www')
    return url.count('www')

def count_atrate(url):

    return url.count('@')

def no_of_dir(url):
    urldir = urlparse(url).path
    return urldir.count('/')

def no_of_embed(url):
    urldir = urlparse(url).path
    return urldir.count('//')

def shortening_service(url):
    match = re.search('bit\.ly|goo\.gl|shorte\.st|go2l\.ink|x\.co|ow\.ly|t\.co|tinyurl|tr\.im|is\.gd|cli\.gs|'
                      'yfrog\.com|migre\.me|ff\.im|tiny\.cc|url4\.eu|twit\.ac|su\.pr|twurl\.nl|snipurl\.com|'
                      'short\.to|BudURL\.com|ping\.fm|post\.ly|Just\.as|bkite\.com|snipr\.com|fic\.kr|loopt\.us|'
                      'doiop\.com|short\.ie|kl\.am|wp\.me|rubyurl\.com|om\.ly|to\.ly|bit\.do|t\.co|lnkd\.in|'
                      'db\.tt|qr\.ae|adf\.ly|goo\.gl|bitly\.com|cur\.lv|tinyurl\.com|ow\.ly|bit\.ly|ity\.im|'
                      'q\.gs|is\.gd|po\.st|bc\.vc|twitthis\.com|u\.to|j\.mp|buzurl\.com|cutt\.us|u\.bb|yourls\.org|'
                      'x\.co|prettylinkpro\.com|scrnch\.me|filoops\.info|vzturl\.com|qr\.net|1url\.com|tweez\.me|v\.gd|'
                      'tr\.im|link\.zip\.net',
                      url)
    if match:
        return 1
    else:
        return 0

def count_https(url):
    return url.count('https')

def count_http(url):
    return url.count('http')

def count_per(url):
    return url.count('%')

def count_ques(url):
    return url.count('?')

def count_hyphen(url):
    return url.count('-')

def count_equal(url):
    return url.count('=')

def url_length(url):
    return len(str(url))

#Hostname Length

def hostname_length(url):
    return len(urlparse(url).netloc)

def suspicious_words(url):
    match = re.search('PayPal|login|signin|bank|account|update|free|lucky|service|bonus|ebayisapi|webscr',
                      url)
    if match:
        return 1
    else:
        return 0

def digit_count(url):
    digits = 0
    for i in url:
        if i.isnumeric():
            digits = digits + 1
    return digits

def letter_count(url):
    letters = 0
    for i in url:
        if i.isalpha():
            letters = letters + 1
    return letters

#Importing dependencies
from urllib.parse import urlparse
from tld import get_tld
import os.path

#First Directory Length
def fd_length(url):
    urlpath= urlparse(url).path
    try:
        return len(urlpath.split('/')[1])
    except:
        return 0

def tld_length(tld):
    try:
        return len(tld)
    except:
        return -1

def main(url):

    status = []

    status.append(having_ip_address(url))
    status.append(abnormal_url(url))
    status.append(count_dot(url))
    status.append(count_www(url))
    status.append(count_atrate(url))
    status.append(no_of_dir(url))
    status.append(no_of_embed(url))

    status.append(shortening_service(url))
    status.append(count_https(url))
    status.append(count_http(url))

    status.append(count_per(url))
    status.append(count_ques(url))
    status.append(count_hyphen(url))
    status.append(count_equal(url))

    status.append(url_length(url))
    status.append(hostname_length(url))
    status.append(suspicious_words(url))
    status.append(digit_count(url))
    status.append(letter_count(url))
    status.append(fd_length(url))
    tld = get_tld(url,fail_silently=True)

    status.append(tld_length(tld))

    print(status)
    return status

def get_prediction_from_url(test_url):
    try:

        features_test = main(test_url)

        # Due to updates to scikit-learn, we now need a 2D array as a parameter to the predict function.
        features_test = np.array(features_test).reshape((1, -1))
        print(features_test)
        # features_test=[0, 1, 2, 0, 0, 3, 0, 0, 0, 1, 0, 0, 1, 0, 73, 34, 0, 3, 61, 10, 3]
        LGB_C = joblib.load('Models/rf_model1.h5')

        pred = LGB_C.predict(features_test)

        print(pred)

        if int(pred[0]) == 0:

            res="SAFE"
            return res
        elif int(pred[0]) == 1.0:

            res="DEFACEMENT"
            return res
        elif int(pred[0]) == 2.0:
            res="PHISHING"
            return res

        elif int(pred[0]) == 3.0:

            res="MALWARE"
        return 1
    except Exception as e:
        print("error occurs",e)

##############End  Model 1 new########################

############## Model 2 old########################

def getTokens(input):
    tokensBySlash = str(input.encode('utf-8')).split('/')
    allTokens=[]
    for i in tokensBySlash:
        tokens = str(i).split('-')
        tokensByDot = []
        for j in range(0,len(tokens)):
            tempTokens = str(tokens[j]).split('.')
            tokentsByDot = tokensByDot + tempTokens
        allTokens = allTokens + tokens + tokensByDot
    allTokens = list(set(allTokens))
    if 'com' in allTokens:
        allTokens.remove('com')
    return allTokens

############## End Model 2########################

import requests
from bs4 import BeautifulSoup

from flask import *
import json
app = Flask(__name__)
@app.route('/')
def home():
    return render_template("index.html")

@app.route('/news')
def news():
    return render_template("news2.html")

@app.route('/preventions')
def prevention():
    return render_template("prevention.html")

import joblib
import pickle
from sklearn.feature_extraction.text import TfidfVectorizer
from scipy.sparse import csr_matrix

@app.route('/getanswer', methods =["GET", "POST"])
def getrecm():
    if request.method == "POST":
        predictx = request.form.get("mname")
        if not predictx.startswith("http://") and not predictx.startswith("https://"):
            predictx = "http://" + predictx

        url=predictx
        print(predictx)

        try:
            X_predict = [predictx]
            # Load the serialized vectorizer object from the file
            with open('Models/vectorizer.pkl', 'rb') as f:
                vectorizer = pickle.load(f)

            # Define some new URLs to transform
            new_urls = X_predict

            # Use the loaded vectorizer to transform the new URLs into TF-IDF format
            new_urls_tfidf = vectorizer.transform(new_urls)

            # Print the transformed URLs
            X_predict=new_urls_tfidf.toarray()
            print(new_urls_tfidf.toarray())
            print("****")
            print(X_predict)

            # load the saved model
            clfsm = joblib.load('Models/rf_model2.h5')
            y_Predict = clfsm.predict(X_predict)
            #print(y_Predict)

            ## for new model
            y_Predict1= get_prediction_from_url(url)
            print(url)
            try:
                page = requests.get(url)
                soup = BeautifulSoup(page.content, "html.parser")

                final_url = page.url
                # serving_ip = page.raw._connection.sock.getpeername()[0]
                serving_ip=123
                status_code = page.status_code
                body_length = len(page.content)
                headers = json.dumps(dict(page.headers))
                meta_tags = soup.find_all("meta")
                arr=[y_Predict[0],url,status_code,body_length, headers,  meta_tags,y_Predict1]

            except:

                arr=[y_Predict[0],url,'NA','NA', 'NA',  'NA',y_Predict1]

            # url_details=f"Final URL: {final_url}
Serving IP Address: {serving_ip}
Status Code: {status_code}
Body Length: {body_length} B
Headers:
{headers}
Meta Tags:
{meta_tags}"

            # return y_Predict[0]+"
"+url_details
            return render_template("result2.html",url_details=arr)

        except Exception as e:
            print(e)
            return ("error")

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


""" Bout utility methods

Methods for extracting bouts from DataFrames and annotating other DataFrames
with this bout information.

A bout is a time range within a larger set of data that shares a particular feature.
"""

import pandas as pd

def extract_bouts(
    df, valid, range_column="t", valid_column="valid", keep_invalid=True, by=[]
):
    """
    Extract from a Pandas DataFrame a list of bouts, where each bout is indicated by a minimum and maximum
    timestamp range and determined by valid ranges.

    Parameters
    ----------
    df : pandas.Dataframe
        The data to extract the bouts from
    valid : pandas.Series
        A series of bool values indicating what rows are considered valid and invalid
    range_column : str
        Optional string indicating the column in df for the timestamp
    valid_column : str
        Optional string indicating the name in the output DataFrame indicating the validness of the bout
    keep_invalid : bool
        Optional value indicating whether to keep invalid bouts in the returned DataFrame
    by : list
        Optional list of columns to group the data by, before extracting the bouts, meaning that bout
        boundaries are both determined by the valid column and the group boundaries.

    Returns
    -------
    pandas.DataFrame
        Index:
            RangeIndex
        Columns:
            Name: count, dtype: int64
                Number of rows in df belonging to the bout
            Name: t_start, dtype: datetime64[ns]
                Starting timestamp of the bout (t_ prefix depends on range_column)
            Name: t_end, dtype: datetime64[ns]
                End timestamp of the bout (t_ prefix depends on range_column)
            Name: valid, dtype: bool
                Whether the bout is valid according to given criterium
    """
    df["_filter"] = valid.loc[df.index]

    dfBool = valid != valid.shift()
    dfCumsum = dfBool.cumsum()

    by_list = [dfCumsum]
    for b in by:
        by_list.append(b)

    groups = df.groupby(by=by_list, sort=False)

    bouts = groups.agg({range_column: ["count", "min", "max"], "_filter": ["first"]})
    bouts.columns = [
        "count",
        range_column + "_start",
        range_column + "_end",
        valid_column,
    ]

    df.drop(columns=["_filter"], inplace=True)

    if not keep_invalid:
        bouts = bouts[bouts[valid_column]]

    bouts.reset_index(drop=True, inplace=True)

    return bouts

def with_padded_bout_window(bouts, window=[0, 0], range_column="t"):
    """
    Pad the values in a Pandas DataFrame created with `extract_bouts` with a time window.

    Parameters
    ----------
    bouts : pandas.Dataframe
        The DataFrame containing the bouts
    window : list
        The number of seconds to add to the starting and end time of the bout.
    range_column : str
        Optional string indicating the column in original for the timestamp. This results in a prefix
        in the bouts DataFrame, timestamp column 't' leads to bout columns 't_min' and 't_max'.

    Returns
    -------
    pandas.DataFrame
        A copy of the bouts DataFrame with padded min and max timestamp values
    """
    bouts = bouts.copy()
    bouts[range_column + "_start"] = bouts[range_column + "_start"] + pd.to_timedelta(
        window[0], unit="s"
    )
    bouts[range_column + "_end"] = bouts[range_column + "_end"] + pd.to_timedelta(
        window[1], unit="s"
    )

    return bouts

def add_bouts_as_column(
    df,
    bouts,
    new_column="bout",
    range_column="t",
    valid_column="valid",
    value="column",
    reset_value=pd.Series(dtype="float64"),
):
    """
    Applies the time ranges in a bouts DataFrame created with `extract_bouts` to the rows in another DataFrame, by
    adding bout data to a new column.

    Parameters
    ----------
    df : pandas.DataFrame
        The DataFrame containing the data that has to be annotated by bout information
    bouts : pandas.Dataframe
        The DataFrame containing the bouts
    new_column : str
        The optional column name to add with bout information
    range_column : str
        Optional string indicating the column in original for the timestamp. This results in a prefix
        in the bouts DataFrame, timestamp column 't' leads to bout columns 't_start' and 't_end'.
    valid_column : str
        Optional string indicating the name in the output DataFrame indicating the validness of the bout
    value : object
        Optional value to insert for a valid bout. If 'index' it takes the bout index as bout identifier,
        'column' fills in the valid column else sets the constant value given.
    reset_value : object
        Optional default value set to the new bouts column if it does not yet exist

    Returns
    -------
    pandas.DataFrame
        A reference to the updated df DataFrame. The original DataFrame is updated in place.
    """
    if new_column not in df:
        df[new_column] = reset_value

    for idx, bout in bouts.iterrows():
        df.loc[
            (df[range_column] >= bout[range_column + "_start"])
            & (df[range_column] < bout[range_column + "_end"]),
            new_column,
        ] = (
            bout[valid_column]
            if value == "column"
            else idx
            if value == "index"
            else value
        )

    return df

def interpolate_bouts_as_column(
    df,
    df_values,
    bouts,
    new_column="bout",
    range_column="t",
    valid_column="valid",
    value_column="position",
    reset_value=pd.Series(dtype="float64"),
):
    """
    Applies the time ranges in a bouts DataFrame created with `extract_bouts` to the rows in another DataFrame, by
    adding bout data to a new column.

    Parameters
    ----------
    df : pandas.DataFrame
        The DataFrame containing the data that has to be annotated by bout information
    bouts : pandas.Dataframe
        The DataFrame containing the bouts
    df_values : pandas.DataFrame
        The DataFrame containing the values to interpolate into df
    value_column : str
        The column in df_values with the values to interpolate into df
    new_column : str
        The optional column name to add with bout information
    range_column : str
        Optional string indicating the column in original for the timestamp. This results in a prefix
        in the bouts DataFrame, timestamp column 't' leads to bout columns 't_start' and 't_end'.
    valid_column : str
        Optional string indicating the name in the output DataFrame indicating the validness of the bout
    reset_value : object
        Optional default value set to the new bouts column if it does not yet exist

    Returns
    -------
    pandas.DataFrame
        A reference to the updated df DataFrame. The original DataFrame is updated in place.
    """

    if new_column not in df:
        df[new_column] = reset_value

    for idx, bout in bouts.iterrows():
        bout_data = df[
            (df[range_column] >= bout[range_column + "_start"])
            & (df[range_column] < bout[range_column + "_end"])
        ]

        t_data = pd.to_numeric(bout_data[range_column])
        t_bout, val_bout = (
            pd.to_numeric(df_values[range_column]),
            df_values[value_column],
        )

        if len(t_bout) > 1:
            from scipy import interpolate

            f = interpolate.interp1d(t_bout, val_bout, fill_value="extrapolate")
            val_data = [f(t) for t in t_data]
            df[new_column].iloc[bout_data.index] = val_data
        else:
            val_data = [
                val_bout[0] + (t - t_bout) / 1e9 for t in t_data
            ]  # manual extrapolation from a single point assuming ns precision in timestamps
            df[new_column].iloc[bout_data.index] = val_data

    # changes the df inplace but return dataframe for chaining purposes
    return df

def select_range(df, window, range_column="t", include_end=True):
    """
    Select a range of data for multiple Pandas DataFrames at once.
    It was designed for using it with the output of the bout methods.

    Parameters
    ----------
    df : pandas.DataFrame/list
        The (list of) DataFrame to use
    window : list
        A 2-list containing the minimum and maximum value of the range to select
    range_column : str
        Optional column in the DataFrame to use for the range (usually timestamp)
    include_end : bool
        Optional whether to include the max value in the range

    Returns
    -------
    list(pandas.DataFrame)
        A list of DataFrames selections
    """

    if not isinstance(df, list):
        df = [df]

    return [
        dff[(dff[range_column] >= window[0]) & (dff[range_column] <= window[1])]
        if include_end
        else dff[(dff[range_column] >= window[0]) & (dff[range_column] < window[1])]
        for dff in df
    ]


import json
import string
import random
import time

from django.http import JsonResponse
from django.views.decorators.csrf import csrf_exempt

from ubskin_web_django.order import models as order_models
from ubskin_web_django.member import models as member_models
from ubskin_web_django.item import models as item_models
from ubskin_web_django.common import lib_data
from ubskin_web_django.common import decorators
from ubskin_web_django.common import common

def get_recv(request):
    if request.method == 'GET':
        return_value = {
            'status': 'error',
            'message': '',
        }
        current_page = request.GET.get('page', 1)
        value = request.GET.get('search_value')
        recv_list = order_models.Recv.get_recv_list(current_page, value)
        return_value['status'] = 'success'
        return_value['data'] = recv_list
        return JsonResponse(return_value)

@csrf_exempt
@decorators.wx_api_authenticated
def create_stock_batch_api(request):
    return_value = {
        'status': 'error',
        'message': '',
    }
    if request.method == 'POST':
        stock_batch_id = ''.join(
            random.choice(string.ascii_lowercase + string.digits) \
            for i in range(8)
        )
        while True:
            stock_batch = order_models.StockBatch.get_stock_dict_by_stock_batch_id(stock_batch_id)
            if stock_batch:
                stock_batch_id = ''.join(
                random.choice(string.ascii_lowercase + string.digits) \
                for i in range(8)
                )
            else:
                break
        data = json.loads(request.body)
        recv_code = data.get('shop_id')
        item_codes_dict = data.get('item_codes_dict')
        openid = request.COOKIES.get('openid')
        member = member_models.Member.get_member_by_wx_openid(openid)
        if not member and not member.is_staff:
            return_value['message'] = '无权限'
            return JsonResponse(return_value)
        order_models.create_model_data(
            order_models.StockBatch,
            {
                "stock_batch_id": stock_batch_id,
                "recv_code": recv_code,
                "create_user": member.member_id,
                "create_time": int(time.time()),
            }
        )
        for key, item in item_codes_dict.items():
            if not item:
                continue
            for i in item:
                if not (len(i) == 9 and i.startswith('U')):
                    return_value['message'] = '商品二维码格式错误'
                    return JsonResponse(return_value)
            stock_batch_count = order_models.create_model_data(
                order_models.StockBatchCount,
                {
                    'item_barcode': key,
                    'stock_batch_id': stock_batch_id,
                    'item_count': len(item),
                }
            )
            for i in item:
                qr_code_obj = order_models.ItemQRCode.get_qr_code_obj_by_qr_code(i)
                qr_code_obj.stock_batch_count_id = stock_batch_count.stock_batch_count_id
                qr_code_obj.create_user = member.member_id
            else:
                qr_code_obj.save()

        return_value['status'] = 'success'
        return JsonResponse(return_value)

def item_code(request, qr_code):
    return_value = {
        'status': 'error',
        'message': '',
    }
    # qr_code = request.GET.get('qr_code')
    if not (len(qr_code) == 9 and qr_code.startswith('U')):
        return_value['message'] = '商品码格式错误'
        return JsonResponse(return_value)

    qr_code_obj = order_models.ItemQRCode.get_qr_code_obj_by_qr_code(qr_code)
    if not qr_code_obj:
        return_value['message'] = '没有记录'
        return JsonResponse(return_value)
    qr_code_obj.search_count += 1
    qr_code_obj.save()
    batch_qr_code_id = qr_code_obj.batch_qr_code_id
    batch_qr_code_obj = order_models.get_model_obj_by_pk(
        order_models.BatchQrCode,
        batch_qr_code_id
    )
    if not qr_code_obj.stock_batch_count_id and batch_qr_code_obj.recv_code:
        recv_addr = order_models.Recv.get_recv_addr_by_recv_code(batch_qr_code_obj.recv_code)
        date = batch_qr_code_obj.create_time
        date = lib_data.parse_timestamps(date)
        member_obj = member_models.Member.get_member_by_id(batch_qr_code_obj.create_member)
        if member_obj:
            is_admin = member_obj.is_admin
            action = "出库扫码" if is_admin else '店铺扫码'
        else:
            action = '店铺扫码'
        return_value['data'] = [
            {"action": action, "to": recv_addr, "date": date},
        ]
        return_value["status"] = "success"
        return JsonResponse(return_value)
    elif not qr_code_obj.stock_batch_count_id and not batch_qr_code_obj.recv_code:
        return_value['messge'] = "该二维码未绑定"
        return JsonResponse(return_value)
    stock_batch_count_id = qr_code_obj.stock_batch_count_id
    obj = order_models.get_model_obj_by_pk(
        order_models.StockBatchCount,
        stock_batch_count_id,
    )
    item_obj = item_models.Items.get_item_obj_by_barcode(obj.item_barcode)
    item_name = '未查询到相关商品信息'
    brand_name = '未查询到相关品牌信息'
    item_photo = common.build_photo_url(None, pic_version='title', cdn=True)
    if item_obj:
        item_name = item_obj.item_name
        brand_id = item_obj.brand_id
        item_photo = common.build_photo_url(item_obj.photo_id, pic_version='title', cdn=True)
        if brand_id:
            brand_obj = item_models.Brands.get_brand_by_id(brand_id)
            brand_name = brand_obj.cn_name if brand_obj else brand_name
    stock_batch_dict = order_models.StockBatch.get_stock_dict_by_stock_batch_id(obj.stock_batch_id)
    recv_code = stock_batch_dict.get('recv_code')
    recv_addr = order_models.Recv.get_recv_addr_by_recv_code(recv_code)
    date = stock_batch_dict.get('create_time')
    date = lib_data.parse_timestamps(date)
    member_obj = member_models.Member.get_member_by_id(stock_batch_dict.get('create_user'))
    if member_obj:
        is_admin = member_obj.is_admin
        action = "出库扫码" if is_admin else '店铺扫码'
    else:
        action = '店铺扫码'
    return_value['data'] = [
        {
            "action": action, "to": recv_addr, "date": date,
            'item_name': item_name, 'brand_name': brand_name, 'image': item_photo,
        },
    ]
    return_value["status"] = "success"
    return JsonResponse(return_value)

def check_has_item_qr_code(request):
    return_value = {
        'status': 'error',
        'message': '',
    }
    if request.method == 'GET':
        item_qr_code = request.GET.get('item_qr_code')
        if not (len(item_qr_code) == 9 and item_qr_code.startswith('U')):
            return_value['message'] = '当前二维码无效'
            return JsonResponse(return_value)
        has = order_models.ItemQRCode.check_has_item_qr_code(item_qr_code)
        if not has:
            return_value['message'] = '当前二维码不在数据库中'
            return JsonResponse(return_value)
        else:
            if has.stock_batch_count_id:
                return_value['message'] = '当前二维码已被绑定'
                return JsonResponse(return_value)
            else:
                return_value['status'] = 'success'
                return JsonResponse(return_value)


from collections import defaultdict

from flask import Flask, request
from flask_restful import Resource, Api
from sqlalchemy import extract

from models import db_session
from models import ModeloColaborador, ModeloPonto

app = Flask(__name__)
api = Api(app)

class Colaborador(Resource):
    def get(self, id):
        colaboradores = ModeloColaborador.query.filter_by(
            id=id).first()
        try:
            response = {
                'id': colaboradores.id,
                'nome': colaboradores.nome,
                'cpf': colaboradores.cpf,
                'email': colaboradores.email,
                'celular': colaboradores.celular,
                'empresa': colaboradores.empresa,
                'cargo': colaboradores.cargo,
                'endereco': colaboradores.endereco,
                'cep': colaboradores.cep,
                'ativo': colaboradores.ativo,
            }
        except AttributeError:
            response = {
                'status': 'error',
                'message': 'Colaborador não encontrado'
            }
        return response

    def put(self, id):
        colaboradores = ModeloColaborador.query.filter_by(
            id=id).first()
        dados = request.json

        if 'id' in dados:
            colaboradores.id = dados['id']

        if 'nome' in dados:
            colaboradores.nome = dados['nome']

        if 'cpf' in dados:
            colaboradores.cpf = dados['cpf']

        if 'email' in dados:
            colaboradores.email = dados['email']

        if 'celular' in dados:
            colaboradores.celular = dados['celular']

        if 'empresa' in dados:
            colaboradores.empresa = dados['empresa']

        if 'cargo' in dados:
            colaboradores.cargo = dados['cargo']

        if 'endereco' in dados:
            colaboradores.endereco = dados['endereco']

        if 'cep' in dados:
            colaboradores.cep = dados['cep']

        if 'ativo' in dados:
            colaboradores.ativo = dados['ativo']

        colaboradores.save()
        response = {
            'id': colaboradores.id,
            'nome': colaboradores.nome,
            'cpf': colaboradores.cpf,
            'email': colaboradores.email,
            'celular': colaboradores.celular,
            'empresa': colaboradores.empresa,
            'cargo': colaboradores.cargo,
            'endereco': colaboradores.endereco,
            'cep': colaboradores.cep,
            'ativo': colaboradores.ativo
        }

        return response

    def delete(self, id):
        colaboradores = ModeloColaborador.query.filter_by(id=id).first()
        mensagem = 'Colaborador {} excluido com sucesso'.format(colaboradores)
        colaboradores.delete()
        return {'status': 'sucesso', 'mensagem': mensagem}

class ListaColaborador(Resource):

    def get(self):
        colaborador = ModeloColaborador.query.all()
        response = [{
            'id': i.id,
            'nome': i.nome,
            'cpf': i.cpf,
            'email': i.email,
            'celular': i.celular,
            'empresa': i.empresa,
            'cargo': i.cargo,
            'endereco': i.endereco,
            'cep': i.cep,
            'ativo': i.ativo

        } for i in colaborador]
        return response

    def post(self):
        dados = request.json
        retorno = []
        retorno.append(- dados['id'])
        for i in range(0, int(dados['id'])):
            retorno.append(dados['nome'])
        colaborador = ModeloColaborador(
            nome=dados['nome'],
            cpf=dados['cpf'],
            email=dados['email'],
            celular=dados['celular'],
            empresa=dados['empresa'],
            cargo=dados['cargo'],
            endereco=dados['endereco'],
            cep=dados['cep'],
            ativo=dados['ativo']
        )

        colaborador.save()
        response = {
            'id': colaborador.id,
            'nome': colaborador.nome,
            'cpf': colaborador.cpf,
            'email': colaborador.email,
            'celular': colaborador.celular,
            'empresa': colaborador.empresa,
            'cargo': colaborador.cargo,
            'endereco': colaborador.endereco,
            'cep': colaborador.cep,
            'ativo': colaborador.ativo
        }
        return response

class Ponto(Resource):
    def get(self, id):
        ponto = ModeloPonto.query.filter_by(
            id=id).first()
        try:
            response = {
                'id': ponto.id,
                'tipo_de_ponto': ponto.tipo_de_ponto,
                'created_at': str(ponto.created_at),
                'last_modified_at': str(ponto.last_modified_at),
                'colaborador_id': ponto.colaborador_id,
            }
        except AttributeError:
            response = {
                'status': 'error',
                'message': 'Ponto não encontrado'
            }
        return response

    def put(self, id):
        ponto = ModeloPonto.query.filter_by(
            id=id).first()
        dados = request.json

        if 'id' in dados:
            ponto.id = dados['id']

        if 'tipo_de_ponto' in dados:
            ponto.tipo_de_ponto = dados['tipo_de_ponto']

        if 'created_at' in dados:
            ponto.created_at = dados['created_at']

        if 'last_modified_at' in dados:
            ponto.last_modified_at = dados['last_modified_at']

        if 'colaborador_id' in dados:
            ponto.colaborador_id = dados['colaborador_id']

        ponto.save()
        response = {
            'id': ponto.id,
            'tipo_de_ponto': ponto.tipo_de_ponto,
            'created_at': ponto.created_at,
            'last_modified_at': ponto.last_modified_at,
            'colaborador_id': ponto.colaborador_id,
        }

        return response

    def delete(self, id):
        ponto = ModeloPonto.query.filter_by(id=id).first()
        mensagem = 'Colaborador {} excluido com sucesso'.format(ponto)
        ponto.delete()
        return {'status': 'sucesso', 'mensagem': mensagem}

class ListaPonto(Resource):

    def get(self):
        ponto = ModeloPonto.query.all()
        response = [{
            'id': i.id,
            'tipo_de_ponto': i.tipo_de_ponto,
            'created_at': i.created_at,
            'last_modified_at': i.last_modified_at,
            'colaborador_id': i.colaborador_id,
        } for i in ponto]
        return response

    def post(self):
        dados = request.json
        retorno = []
        retorno.append(- dados['id'])
        for i in range(0, int(dados['id'])):
            retorno.append(dados['colaborador_id'])
        ponto = ModeloPonto(
            tipo_de_ponto=dados['tipo_de_ponto'],
            created_at=dados['created_at'],
            last_modified_at=dados['last_modified_at'],
            colaborador_id=dados['colaborador_id'],
        )

        ponto.save()
        response = {
            'id': ponto.id,
            'tipo_de_ponto': ponto.tipo_de_ponto,
            'created_at': ponto.created_at,
            'last_modified_at': ponto.last_modified_at,
            'colaborador_id': ponto.colaborador_id,
        }
        return response

api.add_resource(Colaborador, '/colaborador/<int:id>/')
api.add_resource(ListaColaborador, '/listacolaborador/')
api.add_resource(Ponto, '/ponto/<int:id>/')
api.add_resource(ListaPonto, '/listaponto/')

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


""" Constants module: Usefull constants computation functions for gradient
descent algorithm. """
# Authors: Hamza Cherkaoui <hamza.cherkaoui@inria.fr>
# License: BSD (3-clause)

import numpy as np
import numba

from .atlas import get_indices_from_roi

def _precompute_sum_ztz_sum_ztz_y(uz_roi, X_roi, n_times_atom, factor):
    """ Precompute usefull constants for _grad_v_scaled_hrf

    Parameters
    ----------
    X_roi : array, shape (n_voxels_roi, n_times), fMRI data
    uz_roi : array, shape (n_voxels_roi, n_times_valid), neural activty signals

    Return
    ------
    sum_ztz : array, shape (2 * n_times_valid - 1,), sum of auto-convolution
        for each temporal activation (for each voxels)
    sum_ztz_y : array or None, shape (n_times_atom,), sum of auto-convolution
        for each temporal activation convolved with x_j (for each voxels)
    """
    n_voxels_roi, n_times_valid = uz_roi.shape
    sum_ztz = np.zeros((2 * n_times_atom - 1,))
    sum_ztz_y = np.zeros((n_times_atom,))
    uz_roi_padded = np.zeros(2 * (n_times_atom - 1) + n_times_valid)
    for j in range(n_voxels_roi):
        _indices = slice(n_times_atom - 1, n_times_atom - 1 + n_times_valid)
        uz_roi_padded[_indices] = uz_roi[j, :]
        sum_ztz += np.convolve(uz_roi[j, :][::-1], uz_roi_padded, 'valid')
        # In case we try to recover the cost-function from the gradient,
        # the formula needs a '2 *' before 'the X_roi[j, :]', thus the factor
        # option (equal to 1.0 in case of gradient compuation)
        sum_ztz_y += np.convolve(uz_roi[j, :][::-1], factor * X_roi[j, :],
                                 'valid')
    return sum_ztz, sum_ztz_y

@numba.jit((numba.float64[:, :], numba.float64[:, :], numba.int64[:, :]),
           nopython=True, cache=True, fastmath=True)
def _precompute_uvtuv(u, v, rois_idx):  # pragma: no cover
    """ Pre-compute uvtuv.

    Parameters
    ----------
    u : array, shape (n_atoms, n_voxels), spatial maps
    v : array, shape (n_hrf_rois, n_times_atom), HRFs
    rois_idx: array, shape (n_hrf_rois, max_indices_per_rois), HRF ROIs

    Return
    ------
    uvtuv : array, shape (n_atoms, n_atoms, 2 * n_times_atom - 1), precomputed
        operator
    """
    n_atoms, n_voxels = u.shape
    _, n_times_atom = v.shape
    vtv = np.empty((n_voxels, 2 * n_times_atom - 1))
    for m in range(rois_idx.shape[0]):
        indices = get_indices_from_roi(m, rois_idx)
        vtv[indices, :] = np.convolve(v[m, :], v[m, ::-1])
    uvtuv = np.zeros((n_atoms, n_atoms, 2 * n_times_atom - 1))
    for k0 in range(n_atoms):
        for k in range(n_atoms):
            _sum = u[k0, 0] * u[k, 0] * vtv[0]
            for j in range(1, n_voxels):
                _sum += u[k0, j] * u[k, j] * vtv[j, :]
            uvtuv[k0, k, :] = _sum
    return uvtuv

def _precompute_B_C(X, z, H, rois_idx):
    """ Compute list of B, C from givem H (and X, z).

    Parameters
    ----------
    X : array, shape (n_voxels, n_times), fMRI data
    z : array, shape (n_atoms, n_times_valid), temporal components
    H : array, shape (n_hrf_rois, n_times_valid, n_times), Toeplitz matrices
    rois_idx: array, shape (n_hrf_rois, max_indices_per_rois),

    Return
    ------
    B : array, shape (n_atoms, n_voxels), precomputed operator
    C : array, shape (n_atoms, n_atoms), precomputed operator
    """
    n_hrf_rois, _ = rois_idx.shape
    n_voxels, _ = X.shape
    n_atoms, _ = z.shape
    B = np.empty((n_hrf_rois, n_atoms, n_voxels))
    C = np.empty((n_hrf_rois, n_atoms, n_atoms))
    for m in range(n_hrf_rois):
        indices = get_indices_from_roi(m, rois_idx)
        H_m = H[m, :, :]
        z_H_m_t = z.dot(H_m.T)
        C[m, :, :] = z_H_m_t.dot(z_H_m_t.T)
        B[m, :, indices] = X[indices, :].dot(z_H_m_t.T)
    return B, C

@numba.jit((numba.float64[:, :], numba.float64[:, :], numba.float64[:, :, :]),
           nopython=True, cache=True, fastmath=True)
def _precompute_d_basis_constant(X, uz, H):  # pragma: no cover
    """ Precompute AtA and AtX.

    Parameters
    ----------
    X : array, shape (n_voxels, n_times), fMRI data
    uz : array, shape (n_voxels, n_times_valid), neural activity
    H : array, shape (n_hrf_rois, n_times_valid, n_times), Toeplitz matrices

    Return
    ------
    AtA : array, shape (n_atoms_hrf, n_atoms_hrf), precomputed operator
    AtX : array, shape (n_atoms_hrf), precomputed operator
    """
    n_voxels_rois, n_times = X.shape
    n_atoms_hrf, _, _ = H.shape
    A_d = np.empty((n_atoms_hrf, n_voxels_rois * n_times))
    AtX = np.empty(n_atoms_hrf)
    AtA = np.empty((n_atoms_hrf, n_atoms_hrf))
    for d in range(n_atoms_hrf):
        A_d[d, :] = uz.dot(H[d, :, :].T).ravel()
    X_ravel = X.ravel()
    for d in range(n_atoms_hrf):
        AtX[d] = A_d[d, :].dot(X_ravel)
        for d1 in range(d, n_atoms_hrf):
            AtA[d, d1] = AtA[d1, d] = A_d[d, :].dot(A_d[d1, :])
    return AtA, AtX


import argparse
import os
import pandas as pd
import numpy as np
from sklearn.preprocessing import StandardScaler
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error, r2_score
from joblib import dump, load

# inference functions ---------------
def model_fn(model_dir):
    clf = load(os.path.join(model_dir, "model.joblib"))
    return clf

if __name__ == "__main__":
    print('Training started')
    parser = argparse.ArgumentParser()
    parser.add_argument('--train', type=str, default=os.environ.get('SM_CHANNEL_TRAIN'))
    parser.add_argument('--model-dir', type=str, default=os.environ.get('SM_MODEL_DIR'))
    parser.add_argument('--output-data-dir', type=str, default=os.environ.get('SM_OUTPUT_DATA_DIR'))

    parser.add_argument('--n_estimators', type=int, default=100)
    parser.add_argument('--max_depth', type=int, default=15)
    parser.add_argument('--max_features', type=str, default='auto')
    parser.add_argument('--min_samples_split', type=int, default=10)
    parser.add_argument('--min_samples_leaf', type=int, default=20)
    parser.add_argument('--currency', type=str)
    args, _ = parser.parse_known_args()

    print('Args', args)
    currency = args.currency
    input_data = os.path.join(args.train, f'USD_{currency}_features.csv')
    df = pd.read_csv(input_data, parse_dates=[0], index_col=0)

    df['1d_future'] = df['1d_delta'].shift(-1)
    df.dropna(inplace=True)

    ## Prepare data for model training

    scaler = StandardScaler();
    data = scaler.fit_transform(df)
    data = pd.DataFrame(data, index=df.index, columns = df.columns)
    features = np.array(data.columns).tolist()
    features.remove('1d_future')
    x = data[features]
    y = data['1d_future']
    train_size = int(0.7 * y.shape[0])
    x_train = x[:train_size]
    y_train = y[:train_size]
    x_test = x[train_size:]
    y_test = y[train_size:]

    rf_model = RandomForestRegressor(
        n_estimators=args.n_estimators,
        max_depth=args.max_depth,
        max_features=args.max_features,
        min_samples_split=args.min_samples_split,
        min_samples_leaf=args.min_samples_leaf,
        random_state=42
    )

    rf_model.fit(x_train, y_train)
    y_pred = rf_model.predict(x_test)
    mse = mean_squared_error(y_test, y_pred)
    r2 = r2_score(y_test, y_pred)

    print(f'Mean Squared Error: {mse:.6f}')
    print(f'R2 Score: {r2:.6f}')
    dump(rf_model, os.path.join(args.model_dir, 'model.joblib'))


import asyncio
import time

import pandas as pd
import pytest
import numpy as np

from collections import Counter

from ...... import dataframe as md
from ...... import tensor as mt
from ......config import Config
from ......core import TileContext
from ......core.context import get_context
from ......core.graph import TileableGraph, TileableGraphBuilder, ChunkGraphBuilder
from ......core.operand import ShuffleFetchType
from ......lib.aio.isolation import new_isolation, stop_isolation
from ......resource import Resource
from ......serialization import serialize
from ......tests.core import require_ray, mock
from ......utils import lazy_import, get_chunk_params
from .....context import ThreadedServiceContext
from .....subtask import Subtask
from ....analyzer import GraphAnalyzer
from ....core import new_task_id, Task
from ..config import RayExecutionConfig
from ..context import (
    RayExecutionContext,
    RayExecutionWorkerContext,
    RayRemoteObjectManager,
    _RayRemoteObjectContext,
)
from ..executor import (
    execute_subtask,
    OrderedSet,
    RayTaskExecutor,
    RayTaskState,
    _RayChunkMeta,
    _RaySubtaskRuntime,
    _RaySlowSubtaskChecker,
)
from ..fetcher import RayFetcher
from ..shuffle import ShuffleManager

ray = lazy_import("ray")

def _gen_subtask_chunk_graph(t):
    graph = TileableGraph([t.data])
    next(TileableGraphBuilder(graph).build())
    return next(ChunkGraphBuilder(graph, fuse_enabled=False).build())

def _gen_subtask_graph(t):
    tileable_graph = t.build_graph(tile=False)
    chunk_graph = next(ChunkGraphBuilder(tileable_graph).build())
    bands = [(f"address_{i}", "numa-0") for i in range(4)]
    band_resource = dict((band, Resource(num_cpus=1)) for band in bands)
    task = Task("mock_task", "mock_session", tileable_graph)
    analyzer = GraphAnalyzer(
        chunk_graph,
        band_resource,
        task,
        Config(),
        dict(),
        shuffle_fetch_type=ShuffleFetchType.FETCH_BY_INDEX,
    )
    subtask_graph = analyzer.gen_subtask_graph()
    return chunk_graph, subtask_graph

class MockRayTaskExecutor(RayTaskExecutor):
    def __init__(self, *args, **kwargs):
        self._set_attrs = Counter()
        self._monitor_tasks = []
        super().__init__(*args, **kwargs)

    @classmethod
    async def _get_apis(cls, session_id: str, address: str):
        return None, None

    @staticmethod
    def _get_ray_executor():
        # Export remote function once.
        return None

    async def get_available_band_resources(self):
        return {}

    async def execute_subtask_graph(self, *args, **kwargs):
        self._monitor_tasks.clear()
        return await super().execute_subtask_graph(*args, **kwargs)

    async def _update_progress_and_collect_garbage(self, *args, **kwargs):
        # Infinite loop to test monitor task cancel.
        self._monitor_tasks.append(asyncio.current_task())
        return await super()._update_progress_and_collect_garbage(*args, **kwargs)

    def monitor_tasks(self):
        return self._monitor_tasks

    def set_attr_counter(self):
        return self._set_attrs

    def __setattr__(self, key, value):
        super().__setattr__(key, value)
        self._set_attrs[key] += 1

class MockTileContext(TileContext):
    def get_all_progress(self) -> float:
        return 1.0

@require_ray
@pytest.mark.asyncio
@mock.patch("mars.services.task.execution.ray.executor.RayTaskState.create")
@mock.patch("mars.services.task.execution.ray.context.RayExecutionContext.init")
@mock.patch("ray.get")
async def test_ray_executor_create(
    mock_ray_get, mock_execution_context_init, mock_task_state_actor_create
):
    task = Task("mock_task", "mock_session", TileableGraph([]))

    # Create RayTaskState actor as needed by default.
    mock_config = RayExecutionConfig.from_execution_config({"backend": "ray"})
    executor = await MockRayTaskExecutor.create(
        mock_config,
        session_id="mock_session_id",
        address="mock_address",
        task=task,
        tile_context=TileContext(),
    )
    assert isinstance(executor, MockRayTaskExecutor)
    assert mock_task_state_actor_create.call_count == 0
    ctx = get_context()
    assert isinstance(ctx, RayExecutionContext)
    ctx.create_remote_object("abc", lambda: None)
    assert mock_ray_get.call_count == 1
    assert mock_task_state_actor_create.call_count == 1

@require_ray
@pytest.mark.asyncio
async def test_ray_executor_destroy():
    task = Task("mock_task", "mock_session", TileableGraph([]))
    mock_config = RayExecutionConfig.from_execution_config({"backend": "ray"})
    executor = MockRayTaskExecutor(
        config=mock_config,
        task=task,
        tile_context=TileContext(),
        task_context={},
        task_chunks_meta={},
        lifecycle_api=None,
        meta_api=None,
    )
    counter = executor.set_attr_counter()
    assert len(counter) > 0
    keys = executor.__dict__.keys()
    assert counter.keys() >= keys
    counter.clear()
    executor.destroy()
    keys = set(keys) - {"_set_attrs", "_monitor_tasks"}
    assert counter.keys() == keys, "Some keys are not reset in destroy()."
    for k, v in counter.items():
        assert v == 1
    assert await executor.get_progress() == 1.0

@require_ray
@mock.patch("ray.get_runtime_context")
def test_ray_execute_subtask_basic(_):
    raw = np.ones((10, 10))
    raw_expect = raw + 1
    a = mt.ones((10, 10), chunk_size=10)
    b = a + 1

    subtask_id = new_task_id()
    subtask_chunk_graph = _gen_subtask_chunk_graph(b)
    r = execute_subtask(subtask_id, serialize(subtask_chunk_graph), 0, False)
    np.testing.assert_array_equal(r, raw_expect)
    test_get_meta_chunk = subtask_chunk_graph.result_chunks[0]
    r = execute_subtask(subtask_id, serialize(subtask_chunk_graph), 1, False)
    assert len(r) == 2
    meta_dict, r = r
    assert len(meta_dict) == 1
    assert meta_dict[test_get_meta_chunk.key][0] == get_chunk_params(
        test_get_meta_chunk
    )
    np.testing.assert_array_equal(r, raw_expect)

@require_ray
@pytest.mark.asyncio
async def test_ray_fetcher(ray_start_regular_shared2):
    pd_value = pd.DataFrame(
        {
            "col1": [str(i) for i in range(10)],
            "col2": np.random.randint(0, 100, (10,)),
        }
    )
    pd_object_ref = ray.put(pd_value)
    np_value = np.asarray([1, 3, 6, 2, 4])
    np_object_ref = ray.put(np_value)
    # Test RayFetcher to fetch mixed values.
    fetcher = RayFetcher()
    await fetcher.append("pd_key", {"object_refs": [pd_object_ref]})
    await fetcher.append("np_key", {"object_refs": [np_object_ref]})
    await fetcher.append("pd_key", {"object_refs": [pd_object_ref]}, [slice(1, 3, 1)])
    await fetcher.append("np_key", {"object_refs": [np_object_ref]}, [slice(1, 3, 1)])
    results = await fetcher.get()
    pd.testing.assert_frame_equal(results[0], pd_value)
    np.testing.assert_array_equal(results[1], np_value)
    pd.testing.assert_frame_equal(results[2], pd_value.iloc[1:3])
    np.testing.assert_array_equal(results[3], np_value[1:3])

@require_ray
@pytest.mark.asyncio
async def test_ray_remote_object(ray_start_regular_shared2):
    class _TestRemoteObject:
        def __init__(self, i):
            self._i = i

        def value(self):
            return self._i

        def foo(self, a, b):
            return self._i + a + b

        async def bar(self, a, b):
            return self._i * a * b

    # Test RayTaskState reference
    state = RayTaskState.create()
    await state.create_remote_object.remote("aaa", _TestRemoteObject, 123)
    assert await state.call_remote_object.remote("aaa", "value") == 123
    state = RayTaskState.create()
    assert await state.call_remote_object.remote("aaa", "value") == 123

    # Test RayRemoteObjectManager
    name = "abc"
    manager = RayRemoteObjectManager()
    manager.create_remote_object(name, _TestRemoteObject, 2)
    r = await manager.call_remote_object(name, "foo", 3, 4)
    assert r == 9
    r = await manager.call_remote_object(name, "bar", 3, 4)
    assert r == 24
    manager.destroy_remote_object(name)
    with pytest.raises(KeyError):
        await manager.call_remote_object(name, "foo", 3, 4)

    # Test _RayRemoteObjectContext
    context = _RayRemoteObjectContext(lambda: RayTaskState.create())
    context.create_remote_object(name, _TestRemoteObject, 2)
    remote_object = context.get_remote_object(name)
    r = remote_object.foo(3, 4)
    assert r == 9
    r = remote_object.bar(3, 4)
    assert r == 24
    context.destroy_remote_object(name)
    with pytest.raises(KeyError):
        remote_object.foo(3, 4)

    class MyException(Exception):
        pass

    class _ErrorRemoteObject:
        def __init__(self):
            raise MyException()

    with pytest.raises(MyException):
        context.create_remote_object(name, _ErrorRemoteObject)

    handle = RayTaskState.get_handle()
    assert handle is not None

@require_ray
def test_ray_execution_context(ray_start_regular_shared2):
    value = 123
    o = ray.put(value)

    def fake_init(self):
        pass

    async def fake_get_chunks_meta_from_service(
        self, data_keys, fields=None, error="raise"
    ):
        mock_meta = {"meta_1": {fields[0]: 1}, "meta_3": {fields[0]: 3}}
        return [mock_meta[k] for k in data_keys]

    with mock.patch.object(
        ThreadedServiceContext, "__init__", new=fake_init
    ), mock.patch.object(
        RayExecutionContext,
        "_get_chunks_meta_from_service",
        new=fake_get_chunks_meta_from_service,
    ):
        mock_config = RayExecutionConfig.from_execution_config({"backend": "ray"})
        mock_worker_addresses = ["mock_worker_address"]
        isolation = new_isolation("test", threaded=True)
        try:
            context = RayExecutionContext(
                mock_config, {"abc": o}, {}, mock_worker_addresses, lambda: None
            )
            context._loop = isolation.loop
            r = context.get_chunks_result(["abc"])
            assert r == [value]

            r = context.get_worker_addresses()
            assert r == mock_worker_addresses

            r = context.get_chunks_meta(["meta_1"], fields=["memory_size"])
            assert r == [{"memory_size": 1}]

            context._task_chunks_meta["meta_1"] = _RayChunkMeta(memory_size=2)
            r = context.get_chunks_meta(["meta_1", "meta_3"], fields=["memory_size"])
            assert r == [{"memory_size": 2}, {"memory_size": 3}]
        finally:
            stop_isolation("test")

def test_ray_execution_worker_context():
    context = RayExecutionWorkerContext(lambda: None)
    with pytest.raises(NotImplementedError):
        context.set_running_operand_key("mock_session_id", "mock_op_key")
    with pytest.raises(NotImplementedError):
        context.register_custom_log_path(
            "mock_session_id",
            "mock_tileable_op_key",
            "mock_chunk_op_key",
            "mock_worker_address",
            "mock_log_path",
        )

    assert context.set_progress(0.1) is None
    assert context.new_custom_log_dir() is None
    assert context.get_storage_info("mock_address") == {}

@require_ray
@pytest.mark.asyncio
async def test_ray_execution_config(ray_start_regular_shared2):
    t1 = mt.random.randint(10, size=(100, 10), chunk_size=100)
    chunk_graph, subtask_graph = _gen_subtask_graph(t1)

    real_executor = RayTaskExecutor._get_ray_executor()

    class MockExecutor:
        opt = {}

        @classmethod
        def options(cls, *args, **kwargs):
            cls.opt = kwargs
            return real_executor.options(*args, **kwargs)

    task = Task("mock_task", "mock_session", TileableGraph([]))
    mock_config = RayExecutionConfig.from_execution_config(
        {
            "backend": "ray",
            "ray": {
                "monitor_interval_seconds": 0,
                "subtask_max_retries": 4,
                "subtask_num_cpus": 0.8,
                "subtask_memory": 1001,
                "n_cpu": 1,
                "n_worker": 1,
            },
        }
    )
    tile_context = MockTileContext()
    executor = MockRayTaskExecutor(
        config=mock_config,
        task=task,
        tile_context=tile_context,
        task_context={},
        task_chunks_meta={},
        lifecycle_api=None,
        meta_api=None,
    )
    executor._ray_executor = MockExecutor
    async with executor:
        await executor.execute_subtask_graph(
            "mock_stage", subtask_graph, chunk_graph, tile_context
        )

    assert MockExecutor.opt["num_cpus"] == 0.8
    assert MockExecutor.opt["max_retries"] == 4
    assert MockExecutor.opt["memory"] == 1001

@require_ray
@pytest.mark.asyncio
@pytest.mark.parametrize("gc_method", ["submitted", "completed"])
async def test_executor_context_gc(ray_start_regular_shared2, gc_method):
    popped_seq = []

    class MockTaskContext(dict):
        def pop(self, k, d=None):
            popped_seq.append(k)
            return super().pop(k, d)

    t1 = mt.random.randint(10, size=(100, 10), chunk_size=100)
    t2 = mt.random.randint(10, size=(100, 10), chunk_size=50)
    t3 = t2 + t1
    t4 = t3.sum(0)
    chunk_graph, subtask_graph = _gen_subtask_graph(t4)
    task = Task("mock_task", "mock_session", TileableGraph([]), fuse_enabled=True)
    mock_config = RayExecutionConfig.from_execution_config(
        {
            "backend": "ray",
            "ray": {
                "monitor_interval_seconds": 0,
                "log_interval_seconds": 0,
                "subtask_max_retries": 0,
                "n_cpu": 1,
                "n_worker": 1,
                "gc_method": gc_method,
            },
        }
    )
    tile_context = MockTileContext()
    task_context = MockTaskContext()
    executor = MockRayTaskExecutor(
        config=mock_config,
        task=task,
        tile_context=tile_context,
        task_context=task_context,
        task_chunks_meta={},
        lifecycle_api=None,
        meta_api=None,
    )
    executor._ray_executor = RayTaskExecutor._get_ray_executor()

    original_execute_subtask_graph = executor._execute_subtask_graph

    async def _wait_gc_execute_subtask_graph(*args, **kwargs):
        # Mock _execute_subtask_graph to wait the monitor task done.
        await original_execute_subtask_graph(*args, **kwargs)
        await executor.monitor_tasks()[0]

    with mock.patch.object(
        executor, "_execute_subtask_graph", _wait_gc_execute_subtask_graph
    ):
        async with executor:
            await executor.execute_subtask_graph(
                "mock_stage", subtask_graph, chunk_graph, tile_context
            )
            await asyncio.sleep(0)
            assert len(executor.monitor_tasks()) == 1
            assert executor.monitor_tasks()[0].done()

    assert len(task_context) == 1
    assert len(popped_seq) == 6
    subtasks = list(subtask_graph.topological_iter())
    chunk_keys1 = set(
        map(
            lambda c: c.key,
            (
                subtasks[0].chunk_graph.results
                + subtasks[1].chunk_graph.results
                + subtasks[3].chunk_graph.results
            ),
        )
    )
    chunk_keys2 = set(
        map(
            lambda c: c.key,
            (subtasks[2].chunk_graph.results + subtasks[4].chunk_graph.results),
        )
    )
    assert chunk_keys1 == set(popped_seq[0:4])
    assert chunk_keys2 == set(popped_seq[4:])

    task_context.clear()

    original_update_progress_and_collect_garbage = (
        executor._update_progress_and_collect_garbage
    )

    async def infinite_update_progress_and_collect_garbage(*args, **kwargs):
        # Mock _update_progress_and_collect_garbage that never done.
        await original_update_progress_and_collect_garbage(*args, **kwargs)
        while True:
            await asyncio.sleep(0)

    with mock.patch("logging.Logger.info") as log_patch, mock.patch.object(
        executor,
        "_update_progress_and_collect_garbage",
        infinite_update_progress_and_collect_garbage,
    ):
        async with executor:
            await executor.execute_subtask_graph(
                "mock_stage2", subtask_graph, chunk_graph, tile_context
            )
            await asyncio.sleep(0)
            assert len(executor.monitor_tasks()) == 1
            assert executor.monitor_tasks()[0].done()
        assert log_patch.call_count > 0
        args = [c.args[0] for c in log_patch.call_args_list]
        assert any("Submitted [%s/%s]" in a for a in args)
        assert any("Completed [%s/%s]" in a for a in args)

    assert len(task_context) == 1

    task_context.clear()

    # Test the monitor aiotask is done even an exception is raised.
    async def _raise_load_subtask_inputs(*args, **kwargs):
        # Mock _load_subtask_inputs to raise an exception.
        await asyncio.sleep(0)
        1 / 0

    with mock.patch.object(
        executor, "_load_subtask_inputs", _raise_load_subtask_inputs
    ):
        async with executor:
            with pytest.raises(ZeroDivisionError):
                await executor.execute_subtask_graph(
                    "mock_stage3", subtask_graph, chunk_graph, tile_context
                )
            await asyncio.sleep(0)
            assert len(executor.monitor_tasks()) == 1
            assert executor.monitor_tasks()[0].done()

@require_ray
@pytest.mark.asyncio
@pytest.mark.parametrize("gc_method", ["submitted", "completed"])
async def test_execute_shuffle(ray_start_regular_shared2, gc_method):
    chunk_size, n_rows = 10, 50
    df = md.DataFrame(
        pd.DataFrame(np.random.rand(n_rows, 3), columns=list("abc")),
        chunk_size=chunk_size,
    )
    df2 = df.groupby(["a"]).apply(lambda x: x)
    chunk_graph, subtask_graph = _gen_subtask_graph(df2)
    task = Task("mock_task", "mock_session", TileableGraph([]), fuse_enabled=True)

    class MockRayExecutor:
        @staticmethod
        def options(**kwargs):
            num_returns = kwargs["num_returns"]

            class _Wrapper:
                @staticmethod
                def remote(*args):
                    args = [
                        ray.get(a) if isinstance(a, ray.ObjectRef) else a for a in args
                    ]
                    r = execute_subtask(*args)
                    assert len(r) == num_returns
                    return [ray.put(i) for i in r]

            return _Wrapper

    mock_config = RayExecutionConfig.from_execution_config(
        {
            "backend": "ray",
            "ray": {
                "monitor_interval_seconds": 0,
                "subtask_max_retries": 0,
                "n_cpu": 1,
                "n_worker": 1,
                "gc_method": gc_method,
            },
        }
    )
    tile_context = MockTileContext()
    task_context = {}
    executor = MockRayTaskExecutor(
        config=mock_config,
        task=task,
        tile_context=tile_context,
        task_context=task_context,
        task_chunks_meta={},
        lifecycle_api=None,
        meta_api=None,
    )
    executor._ray_executor = MockRayExecutor

    # Test ShuffleManager.remove_object_refs
    sm = ShuffleManager(subtask_graph)
    sm._mapper_output_refs[0].fill(1)
    sm.remove_object_refs(next(iter(sm._reducer_indices.keys())))
    assert pd.isnull(sm._mapper_output_refs[0][:, 0]).all()
    sm._mapper_output_refs[0].fill(1)
    sm.remove_object_refs(next(iter(sm._mapper_indices.keys())))
    assert pd.isnull(sm._mapper_output_refs[0][0]).all()
    with pytest.raises(ValueError):
        sm.remove_object_refs(None)

    original_execute_subtask_graph = executor._execute_subtask_graph

    async def _wait_gc_execute_subtask_graph(
        stage_id, subtask_graph, chunk_graph, monitor_context
    ):
        # Mock _execute_subtask_graph to wait the monitor task done.
        await original_execute_subtask_graph(
            stage_id, subtask_graph, chunk_graph, monitor_context
        )
        await executor.monitor_tasks()[0]
        assert pd.isnull(monitor_context.shuffle_manager._mapper_output_refs[0]).all()

    with mock.patch.object(
        executor, "_execute_subtask_graph", _wait_gc_execute_subtask_graph
    ), mock.patch("ray.get_runtime_context"):
        async with executor:
            await executor.execute_subtask_graph(
                "mock_stage", subtask_graph, chunk_graph, tile_context
            )
        await asyncio.sleep(0)
        assert len(executor.monitor_tasks()) == 1
        assert executor.monitor_tasks()[0].done()

    assert len(task_context) == len(chunk_graph.results)

@require_ray
@pytest.mark.asyncio
async def test_slow_subtask_checker():
    subtasks = [
        Subtask(str(i), logic_key=f"logic_key1", logic_parallelism=5) for i in range(5)
    ]
    for s in subtasks:
        s.runtime = _RaySubtaskRuntime()
    submitted = OrderedSet()
    completed = OrderedSet()
    now = time.time()
    checker = _RaySlowSubtaskChecker(5, submitted, completed)
    updater = checker.update()
    for s in subtasks:
        submitted.add(s)
    for _ in updater:
        break
    assert all(s.runtime.start_time >= now for s in subtasks)
    await asyncio.sleep(0.01)
    assert not any(checker.is_slow(s) for s in subtasks)
    completed.add(subtasks[0])
    completed.add(subtasks[1])
    for _ in updater:
        break
    await asyncio.sleep(0.01)
    completed.add(subtasks[2])
    assert not any(checker.is_slow(s) for s in subtasks[3:])
    completed.add(subtasks[3])
    for _ in updater:
        break
    assert not checker.is_slow(subtasks[4])
    await asyncio.sleep(0.1)
    assert checker.is_slow(subtasks[4])

@require_ray
@pytest.mark.asyncio
async def test_execute_slow_task(ray_start_regular_shared2):
    t1 = mt.random.randint(10, size=(100, 10), chunk_size=10)
    t2 = mt.random.randint(10, size=(100, 10), chunk_size=30)
    t3 = t2 + t1
    t4 = t3.sum(0)
    chunk_graph, subtask_graph = _gen_subtask_graph(t4)
    task = Task("mock_task", "mock_session", TileableGraph([]), fuse_enabled=True)
    mock_config = RayExecutionConfig.from_execution_config(
        {
            "backend": "ray",
            "ray": {
                "monitor_interval_seconds": 0,
                "log_interval_seconds": 0,
                "check_slow_subtasks_iqr_ratio": 1.5,
                "check_slow_subtasks_interval_seconds": 0,
                "subtask_max_retries": 0,
                "n_cpu": 1,
                "n_worker": 1,
            },
        }
    )
    tile_context = MockTileContext()
    executor = MockRayTaskExecutor(
        config=mock_config,
        task=task,
        tile_context=tile_context,
        task_context={},
        task_chunks_meta={},
        lifecycle_api=None,
        meta_api=None,
    )
    slow_subtask_id = list(subtask_graph)[-1].subtask_id

    def mock_execute_subtask(subtask_id, *args):
        if subtask_id == slow_subtask_id:
            time.sleep(2)
        return execute_subtask(subtask_id, *args)

    executor._ray_executor = ray.remote(mock_execute_subtask)

    with mock.patch("logging.Logger.info") as log_patch:
        async with executor:
            await executor.execute_subtask_graph(
                "mock_stage2", subtask_graph, chunk_graph, tile_context
            )
            await asyncio.sleep(0)
            assert len(executor.monitor_tasks()) == 1
            assert executor.monitor_tasks()[0].done()
        assert log_patch.call_count > 0
        slow_ray_object_refs = set()
        for c in log_patch.call_args_list:
            if c.args[0] == "Slow tasks(%s): %s":
                count, object_refs = c.args[1:]
                assert count >= 1
                slow_ray_object_refs.update(object_refs)
        assert len(slow_ray_object_refs) >= 1


from . import user
from main import db
from schemas.users import (
    UserSignUpSchema,
    UserSchema
)
from models.users import Role

from flask import (
    request
)
from loguru import logger

@user.post('/')
def sign_up():
    schema = UserSignUpSchema()
    data = request.get_json(force=True)
    try:
        new_user = schema.load(data)
        # persist to database
        try:
            db.session.add(new_user)
            db.session.commit()

            resp = UserSchema().dump(new_user)

        except Exception as e:
            logger.error(e)
            db.session.rollback()
        finally:
            db.session.close()

    except Exception as e:
        return {
            'status': 'error',
            'message': str(e),
            'data': {}
        }

    return {
        'status': 'success',
        'message': 'sign up successful',
        'data': resp
    }, 201

@user.patch('/')
def update_info():
    pass


from sys import exit
from itertools import product

inputArr = [
    '3, 1000',
    '2, 5, 4',
    '3, 7, 8, 9',
    '5, 5, 7, 8, 9, 10'
]

K, M = map(int, inputArr[0].split(', '))

# exit if the int is larger than the parameters
# 1 <= K <= 7
# 1 <= M <= 1000
if (K > 7 or K < 1):
    exit("K is bigger than parameters, the integor needs to be between 1 and 7.")

if (M > 1000 or M < 1):
    exit("M is bigger than parameters, the integor needs to be between 1 and 1,000.")

# print('K: ',K, ' M: ', M)

arrayN = []

# append to arrayN an array of the string items between inputArr's [1:]
for i in range(1, K + 1):
    arrayN.append([int(x) for x in inputArr[i].split(', ')])

# print(arrayN)

# function: sum of square of possible combinations, modulo if greater than 1000
def squaredSums(nums):
    return sum(x**2 for x in nums) % M

# put combinations of ints in own array
possible_combinations = list(product(*arrayN))

# get squared sum of possible combinations through squaredSum()
possible_sums = list(map(squaredSums, possible_combinations))

# print the max value
print("""The sum of the maximum values from each arrays,
         after being squared, is {}.""".format(max(possible_sums)))


import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import matplotlib as mpl

def graph2d(x_values, y_values, x_axis, y_axis):
    """Graphs a curve in 2D, with the values passed as parameters."""
    plt.plot(x_values, y_values)
    plt.xlabel(x_axis)
    plt.ylabel(y_axis)
    plt.show()

def graph3d(x_values, y_values, z_values, x_axis, y_axis, z_axis):
    """Graphs a parametric curve with the values passed as parameters."""
    mpl.rcParams['legend.fontsize'] = 10

    fig = plt.figure()
    ax = fig.gca(projection='3d')

    ax.plot(x_values, y_values, z_values, label='parametric curve')
    ax.legend()

    ax.set_xlabel(x_axis)
    ax.set_ylabel(y_axis)
    ax.set_zlabel(z_axis)

    plt.show()


from publisher.models import LayoutTemplate, Book, BookDownload, BookRating, BookPage, BookImage, BookComment, Translation,  Publisher
from serializers import LayoutTemplateSerializer, BookSerializer, BookPageSerializer, TranslationSerializer
from serializers import ImageSerializer, CommentSerializer
from rest_framework.decorators import detail_route, list_route, permission_classes
from rest_framework import status
from rest_framework.parsers import FormParser, MultiPartParser
from django.http import HttpResponse
from django.core.files.base import ContentFile
from rest_framework import viewsets
from django.http import FileResponse
from rest_framework.views import APIView
from rest_framework.response import Response
from rest_framework import authentication, permissions
from django.http import JsonResponse, FileResponse, Http404
from rest_framework import generics
from rest_framework.decorators import api_view, parser_classes
from django.db.models import Count, Max, Q
from django.conf import settings
from polli import constants
from users.models import User
import csv
import string
import random
import urllib2
import json
import time
import boto3
import json

# Layout Template - API
class LayoutTemplateViewSet(viewsets.ModelViewSet):
    queryset = LayoutTemplate.objects.all()
    serializer_class = LayoutTemplateSerializer

    def create(self, request):
        serializer = LayoutTemplateSerializer(data=request.data)
        if serializer.is_valid():
            user = self.request.user
            publisher = user.publisher if user.user_type == User.PUBLISHER else None
            serializer.save(publisher=publisher)
            return Response(serializer.data)

# Book - API Views
class BookViewSet(viewsets.ModelViewSet):
    queryset = Book.objects.all()
    serializer_class = BookSerializer

    def create(self, request):
        serializer = BookSerializer(data=request.data)
        if serializer.is_valid():
            user = self.request.user
            publisher = user.publisher if user.user_type == User.PUBLISHER else None
            book = serializer.save(publisher=publisher)
            translation = request.data['translation']
            if translation:
                Translation.objects.create(book=book, language=translation)

            # Send Notification Email
            subject = 'Book Uploaded: ({}) {}'.format(book.id, book.name)
            heading = 'New Book Uploaded'
            book.send_upload_notification(subject=subject, heading=heading, request=request)
            return Response(serializer.data)
        else:
            return Response(serializer.errors)

    def list(self, request):
        books = Book.objects.filter(translations__status='published')
        page = self.paginate_queryset(books)
        serializer = BookSerializer(page, many=True)
        return self.get_paginated_response(serializer.data)

    @detail_route(methods=['get'])
    def get_book_bundle(self, request, pk=None):
            # Record Download
            primary_language = request.query_params['primary_language']
            secondary_language = request.query_params['secondary_language']
            BookDownload.objects.create(
                user=request.user,
                book_id=pk,
                primary_language=primary_language,
                secondary_language=secondary_language
            )

            # Get Bundle
            book = Book.objects.get(pk=pk)
            bundle_file_path = book.generate_bundle(primary_language=primary_language, secondary_language=secondary_language)
            return FileResponse(open(bundle_file_path, 'rb'), content_type='application/octet-stream')

    @detail_route(methods=['post'])
    def rate_book(self, request, pk):
            book_rating, created = BookRating.objects.get_or_create(user=request.user, book_id=pk)
            book_rating.rating = float(request.data['rating'])
            book_rating.save()
            return Response('Rate Book Complete')

    @detail_route(methods=['post'])
    def process_book(self, request, pk=None):
        book = Book.objects.get(pk=pk)
        book.process_book_pdf()
        serializer = BookSerializer(instance=book)
        return Response(serializer.data)

    @detail_route(methods=['post'])
    def get_book_processing_status(self, request, pk):
        book = Book.objects.get(pk=pk)
        serializer = BookSerializer(instance=book)
        return Response(serializer.data)

    @detail_route(methods=['post'])
    def update_book_cover(self, request, pk=None):
        book = Book.objects.get(pk=pk)
        book.set_cover_image(image_url=request.data['image_url'])
        serializer = BookSerializer(instance=book)
        return Response(serializer.data)

# Page - API Views
class PageViewSet(viewsets.ModelViewSet):
    queryset = BookPage.objects.all()
    serializer_class = BookPageSerializer

    def create(self, request):
        serializer = BookPageSerializer(data=request.data)
        if serializer.is_valid():
            book_id = request.data['book']
            order = Book.objects.get(pk=book_id).pages.aggregate(Max('order'))['order__max'] + 1
            serializer.save(order=order)
            return Response(serializer.data)
        else:
            return Response(serializer.errors)

    @list_route(methods=['post'])
    def update_page_sorting(self, request):
        for page in json.loads(request.data['pages']):
            BookPage.objects.filter(id=int(page['id'])).update(order=int(page['order']))
        return Response('Page Sorting Updated')

# Translation - API Views
class TranslationViewSet(viewsets.ModelViewSet):
    queryset = Translation.objects.all()
    serializer_class = TranslationSerializer

    def create(self, request):
        serializer = TranslationSerializer(data=request.data)
        if serializer.is_valid():
            serializer.save()
            return Response(serializer.data)
        else:
            return Response(serializer.errors)

    @detail_route(methods=['get'])
    def get_translation_csv(self, request, pk):
        translation_csv = HttpResponse(content_type='text/csv')
        translation_csv['Content-Disposition'] = 'attachment; filename="translation_{}.csv"'.format(pk)
        translation = Translation.objects.get(pk=pk)
        translation.generate_translation_csv(translation_csv)
        return translation_csv

    @detail_route(methods=['post'])
    def upload_translation_csv(self, request, pk=None):
        translation = Translation.objects.get(pk=pk)
        csv_data = ContentFile(request.FILES['csv'].read())
        translation.set_translation_csv(csv_data)
        translation.process_translation_csv()
        return Response('upload translation complete')

    @detail_route(methods=['post'])
    def update_translation_review(self, request, pk):
        data = request.data
        translation = Translation.objects.get(pk=pk)
        translation.update_translation_review(page_num=data['page_num'], line_num=data['line_num'], l2=data['l2'], complete=data['complete'])
        return Response('update translation review complete')

    @detail_route(methods=['post'])
    def update_blend_review(self, request, pk):
        data = request.data
        translation = Translation.objects.get(pk=pk)
        translation.update_blend_review(page_num=data['page_num'], line_num=data['line_num'], blend=data['blend'], complete=data['complete'])
        return Response('update blend review complete')

    @detail_route(methods=['post'])
    def create_blend_request(self, request, pk):
        # Update Translation
        translation = Translation.objects.get(pk=pk)
        translation.status = Translation.BLEND
        translation.blend_status = Translation.BLEND_PENDING
        translation.save()

        # Add Blend Request to SQS
        sqs = boto3.resource('sqs', region_name='us-east-1', aws_access_key_id=settings.AWS_ACCESS_KEY_ID, aws_secret_access_key=settings.AWS_SECRET_ACCESS_KEY)
        queue = sqs.get_queue_by_name(QueueName=settings.AWS_SQS_NAME)
        queue.send_message(MessageBody=json.dumps({'translation_id': pk}))
        return Response('create blend request')

    @list_route(permission_classes=[permissions.AllowAny], methods=['post'])
    def generate_blends(self, request):
        translation = Translation.objects.get(pk=request.data['translation_id'])
        translation.generate_blends()
        return Response('Generate Blends')

# Images - API Views
class ImageViewSet(viewsets.ModelViewSet):
    queryset = BookImage.objects.all()
    serializer_class = ImageSerializer

    def create(self, request):
        serializer = ImageSerializer(data=request.data)
        if serializer.is_valid():
            image = serializer.save()

            # Send Notification Email
            book = image.book
            subject = 'Image Uploaded: ({}) {}'.format(book.id, book.name)
            heading = 'New Image Uploaded'
            book.send_upload_notification(subject=subject, heading=heading, request=request)
            return Response(serializer.data)
        else:
            return Response(serializer.errors)

# Comment - API Views
class CommentViewSet(viewsets.ModelViewSet):
    queryset = BookComment.objects.all()
    serializer_class = CommentSerializer

    def create(self, request):
        serializer = CommentSerializer(data=request.data)
        if serializer.is_valid():
            serializer.save(creator=request.user)
            return Response(serializer.data)
        else:
            return Response(serializer.errors)

    @detail_route(methods=['post'])
    def mark_comment_read(self, request, pk):
        comments = BookComment.objects.filter(Q(id=pk)|Q(parent_comment_id=pk))
        for c in comments:
            c.read_by.add(request.user)
        return Response('Comment Read')

    @detail_route(methods=['post'])
    def resolve_comment(self, request, pk):
        BookComment.objects.filter(Q(id=pk) | Q(parent_comment_id=pk)).update(status=BookComment.CLOSED_STATUS)
        comments.update(status=BookComment.CLOSED_STATUS)
        return Response('Comment Resolved')


from django import forms
from django.contrib.auth.forms import UserCreationForm
from .models import DeployInfo

class DeployForm(forms.ModelForm):
    class Meta:
        model = DeployInfo
        fields = ['lsn']
        widgets = {
            'lsn': forms.NumberInput(attrs={'class': 'form-control'}),
        }
        labels = {
            'lsn': '',
        }


from direct.showbase.DirectObject import DirectObject
import panda3d.core as p3d

from .mapgen.dungeon import Dungeon
from .rangeindicator import RangeIndicator

class GameState(DirectObject):
    def __init__(self):
        super().__init__()
        self.root_node = p3d.NodePath('State Root')
        self.root_node.reparent_to(base.render)

    def cleanup(self):
        self.ignoreAll()
        self.root_node.remove_node()
        self.root_node = None

    def update(self, dt):
        pass

class MainState(GameState):
    PLAYER_SPEED = 15
    CAM_MOVE_BORDER = 0.8
    CAM_MOVE_SPEED = 50

    DUNGEON_SX = 50
    DUNGEON_SY = 50

    def __init__(self):
        super().__init__()

        self.accept('toggle-debug-cam', self.toggle_debug_cam)
        self.accept('move', self.move_player)
        self.accept('ability1', self.toggle_range, [0])
        self.accept('ability2', self.toggle_range, [1])
        self.accept('ability3', self.toggle_range, [2])
        self.accept('ability4', self.toggle_range, [3])

        self.mapgen = 'static'
        dungeon = Dungeon(self.mapgen, self.DUNGEON_SX, self.DUNGEON_SY)
        dungeon.model_root.reparent_to(self.root_node)

        dlight = p3d.DirectionalLight('sun')
        dlight.set_color(p3d.LVector3(0.2, 0.2, 0.2))
        dlight.set_shadow_caster(True, 4096, 4096)
        dlnp = self.root_node.attach_new_node(dlight)
        dlnp.set_z(10)
        dlnp.set_p(-90)
        lens = dlight.get_lens()
        lens.set_film_size(60)
        lens.set_near(1)
        lens.set_far(100)
        self.root_node.set_light(dlnp)

        dlight2 = p3d.DirectionalLight('ground')
        dlight2.set_color(p3d.LVector3(0.1, 0.1, 0.1))
        dlnp2 = self.root_node.attach_new_node(dlight2)
        self.root_node.set_light(dlnp2)

        loader = p3d.Loader.get_global_ptr()
        player = p3d.NodePath(loader.load_sync('dungeon.bam')).find('**/MonsterSpawn').node()
        playernp = self.root_node.attach_new_node(player)
        playernp.set_pos(dungeon.player_start)
        playernp.set_z(1.5)
        self.player_ranges = [
            RangeIndicator('box', length=5, width=1),
            RangeIndicator('circle', radius=2),
            RangeIndicator('circle', radius=3),
            RangeIndicator('circle', radius=4),
        ]
        for rangeindicator in self.player_ranges:
            rangeindicator.graphics.reparent_to(playernp)
            rangeindicator.visible = False

        # self.root_node.ls()
        # self.root_node.analyze()

        self.dungeons = [dungeon]
        self.dungeon_idx = 0
        self.dungeon = dungeon
        self.player = playernp
        self.last_tele_loc = None
        self.target = self.player.get_pos()
        self.debug_cam = False
        self.reset_camera()

    def toggle_debug_cam(self):
        if not self.debug_cam:
            mat = p3d.LMatrix4(base.camera.get_mat())
            mat.invert_in_place()
            base.mouseInterfaceNode.set_mat(mat)
            base.enableMouse()
            self.debug_cam = True
        else:
            base.disableMouse()
            self.reset_camera()
            self.debug_cam = False

    def toggle_range(self, ability):
        rangeindicator = self.player_ranges[ability]
        rangeindicator.visible = not rangeindicator.visible

    def reset_camera(self):
        campos = self.player.get_pos()
        campos.z += 25
        campos.y -= 25
        base.camera.set_mat(p3d.LMatrix4.ident_mat())
        base.cam.set_pos(campos)
        base.cam.look_at(self.player.get_pos())

    def update(self, dt):
        # Update player position
        movvec = self.target - self.player.get_pos()
        newpos = self.player.get_pos()
        if movvec.length_squared() < 0.4:
            newpos.set_x(self.target.x)
            newpos.set_y(self.target.y)
        else:
            movvec.normalize()
            movvec *= self.PLAYER_SPEED * dt
            newpos.set_x(self.player.get_x() + movvec.x)
            newpos.set_y(self.player.get_y() + movvec.y)
            self.player.look_at(newpos)

        if self.dungeon.is_walkable(*newpos.xy):
            self.player.set_pos(newpos)

        if self.dungeon.is_exit(*newpos.xy):
            next_didx = self.dungeon_idx + 1
            if next_didx >= len(self.dungeons):
                # Create a new dungeon
                self.dungeons.append(Dungeon(self.mapgen, self.DUNGEON_SX, self.DUNGEON_SY))
            self.switch_to_dungeon(next_didx)

        if self.last_tele_loc is not None:
            if (self.last_tele_loc - self.player.get_pos()).length_squared() > 3:
                self.last_tele_loc = None

        if self.last_tele_loc is None:
            teleloc = self.dungeon.get_tele_loc(*newpos.xy)
            if teleloc is not None:
                newpos.x, newpos.y = teleloc
                self.last_tele_loc = p3d.LVector3(newpos)
                self.target = p3d.LVector3(newpos)
                self.player.set_pos(newpos)
                self.reset_camera()

        if not self.debug_cam and base.mouseWatcherNode.has_mouse():
            mousex, mousey = base.mouseWatcherNode.get_mouse()
            border = self.CAM_MOVE_BORDER
            camdelta = p3d.LVector2(0, 0)
            if mousex < -border:
                camdelta.x -= 1
            elif mousex > border:
                camdelta.x += 1

            if mousey < -border:
                camdelta.y -= 1
            elif mousey > border:
                camdelta.y += 1

            camdelta.normalize()
            camdelta *= self.CAM_MOVE_SPEED * dt

            campos = base.cam.get_pos()
            campos.x += camdelta.x
            campos.y += camdelta.y
            base.cam.set_pos(campos)

    def move_player(self):
        if not base.mouseWatcherNode.has_mouse() or self.debug_cam:
            return

        mousevec = base.mouseWatcherNode.get_mouse()
        near = p3d.LPoint3()
        far = p3d.LPoint3()
        base.camLens.extrude(mousevec, near, far)
        near = self.root_node.get_relative_point(base.cam, near)
        far = self.root_node.get_relative_point(base.cam, far)

        plane = p3d.Plane(p3d.LVector3.up(), p3d.LPoint3())
        worldpos = p3d.LPoint3()
        if plane.intersects_line(worldpos, near, far):
            self.target.set_x(worldpos.x)
            self.target.set_y(worldpos.y)

    def switch_to_dungeon(self, dungeon_idx):
        # Switch to next layer
        print("Switching to dungeon layer {}".format(dungeon_idx))
        self.dungeon.model_root.detach_node()
        self.dungeon = self.dungeons[dungeon_idx]
        self.dungeon.model_root.reparent_to(self.root_node)
        self.dungeon_idx = dungeon_idx
        self.player.set_pos(self.dungeon.player_start)
        self.player.set_z(1.5)
        self.target = self.player.get_pos()
        self.reset_camera()


# training / eval for burger equation

"""
This is the main file to train the GNN on the burger equation

The burger equation is a simple PDE that can be solved with a GNN (and a pinn loss).
The PDE is :
    u_t + u*u_x = mu * u_xx

With the following boundary conditions:
    u(-1, t) = 0
    u(1, t) = 0
    u(x, 0) = sin(pi * x)

The goal is to train a GNN to solve this PDE.
"""

import sys
sys.path.append("/app/")

from discretize_pinn_loss.pdebenchmark_darcy import Darcy2DPDEDataset
from discretize_pinn_loss.models_fno import FNO2d
from discretize_pinn_loss.losses_darcy import DarcyFlowOperator, DarcyLoss

import pytorch_lightning as pl

import torch
from torch_geometric.data import Data, DataLoader, Batch

import hashlib
import datetime
import math

import os

import wandb

import matplotlib.pyplot as plt

import einops

import numpy as np

torch.manual_seed(42)

class FnoFull(pl.LightningModule):
    def __init__(self, model_fno, loss_function):
        super().__init__()

        self.model = model_fno
        self.loss_function = loss_function

        self.loss_mse = torch.nn.MSELoss()

    def forward(self, graph):

        # get the input node and reshape it because the model expects a batch
        nodes = graph.x

        batch_size = graph.ptr.shape[0] - 1

        size_image = math.sqrt(int(nodes.shape[0] / batch_size))

        nodes = einops.rearrange(nodes, "(b n) d -> b n d", b=batch_size)

        images = einops.rearrange(nodes, "b (s sb) d -> b s sb d", s=int(size_image), sb=int(size_image))

        result = self.model(images)

        result = einops.rearrange(result, "b s sb n -> b (s sb) n", s=int(size_image), sb=int(size_image))

        result = einops.rearrange(result, "b n d -> (b n) d")

        return result

    def training_step(self, batch, batch_idx):

        graph_a_x = batch
        edge_index = graph_a_x.edge_index
        edges = graph_a_x.edge_attr

        # forward pass
        nodes_pred = self.forward(graph_a_x)

        # we create the two graphs
        graph_pred = Data(x=nodes_pred, edge_index=edge_index, edge_attr=edges)

        # compute loss
        relative_loss = self.loss_function(graph_pred, graph_a_x)

        loss = self.loss_mse(relative_loss, torch.zeros_like(relative_loss))

        self.log("train_loss", loss)

        return loss

    def validation_step(self, batch, batch_idx):

        a_x = batch

        mask = batch.mask

        # forward pass
        u_x = self.forward(a_x)

        # we create the two graphs
        graph_pred = Data(x=u_x, edge_index=a_x.edge_index, edge_attr=a_x.edge_attr)

        print("graph_pred", graph_pred)
        print("a_x", a_x)

        # compute loss
        relative_loss = torch.abs(self.loss_function(graph_pred, a_x, mask))

        print("relative_loss", relative_loss.shape)

        loss = self.loss_mse(relative_loss, torch.zeros_like(relative_loss))

        self.log("val_loss", loss)

        # now we can save the result image to wandb for visualization
        # we need to reshape the image to be able to plot it
        batch_size = a_x.ptr.shape[0] - 1

        # we should also log the loss with the generated target value coming form the solver (data)
        graph_target= Data(x=batch.target, edge_index=batch.edge_index, edge_attr=batch.edge_attr)

        relative_loss_solver = torch.abs(self.loss_function(graph_target, batch, mask=mask))

        loss_solver = self.loss_mse(relative_loss_solver, torch.zeros_like(relative_loss_solver))

        self.log("val_loss_solver", loss_solver)

        size_image = math.sqrt(int(a_x.x.shape[0] / batch_size))

        a_x = einops.rearrange(a_x.x, "(b n) d -> b n d", b=batch_size)
        a_x = einops.rearrange(a_x, "b (s sb) d -> b s sb d", s=int(size_image), sb=int(size_image))
        u_x = einops.rearrange(u_x, "(b n) d -> b n d", b=batch_size)
        u_x = einops.rearrange(u_x, "b (s sb) d -> b s sb d", s=int(size_image), sb=int(size_image))

        # same operation for the target
        target = batch.target

        target = einops.rearrange(target, "(b n) d -> b n d", b=batch_size)
        target = einops.rearrange(target, "b (s sb) d -> b s sb d", s=int(size_image), sb=int(size_image))

        #same thing on relative loss
        relative_loss = einops.rearrange(relative_loss.squeeze(), "(b n) -> b n", b=batch_size)
        relative_loss = einops.rearrange(relative_loss, "b (s sb) -> b s sb", s=int(size_image), sb=int(size_image))

        #same thing on relative loss solver
        relative_loss_solver = einops.rearrange(relative_loss_solver.squeeze(), "(b n) -> b n", b=batch_size)
        relative_loss_solver = einops.rearrange(relative_loss_solver, "b (s sb) -> b s sb", s=int(size_image), sb=int(size_image))

        # now we can save the result image to wandb for visualization
        # we need to reshape the image to be able to plot it
        # save the image into png format
        fig, axs = plt.subplots(1, 5, figsize=(15, 5))

        pos0 = axs[0].imshow(a_x[0, :, :, 0].detach().cpu().numpy())
        axs[0].set_title("a_x")

        pos1 = axs[1].imshow(u_x[0, :, :, 0].detach().cpu().numpy(), vmin=0, vmax=0.30)
        axs[1].set_title("u_x")

        pos2 = axs[2].imshow(target[0, :, :, 0].detach().cpu().numpy(), vmin=0, vmax=0.30)
        axs[2].set_title("target")

        pos3 = axs[3].imshow(relative_loss[0, :, :].detach().cpu().numpy(), vmin=0, vmax=4)
        axs[3].set_title("relative loss pred")

        pos4 = axs[4].imshow(relative_loss_solver[0, :, :].detach().cpu().numpy(), vmin=0, vmax=4)
        axs[4].set_title("relative loss solver")

        fig.colorbar(pos0, ax=axs[0])
        fig.colorbar(pos1, ax=axs[1])
        fig.colorbar(pos2, ax=axs[2])
        fig.colorbar(pos3, ax=axs[3])
        fig.colorbar(pos4, ax=axs[4])

        # save into png
        fig.savefig("tmp.png")

        self.logger.log_image("a_x u_x target", ["tmp.png"])

        return loss

    # on validation epoch end
    def validation_epoch_end(self, outputs):
        pass

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.model.parameters(), lr=1e-3)
        return optimizer

    def test_train_step(self, dataloader):

        # init optimizer
        optimizer = torch.optim.Adam(self.model.parameters(), lr=1e-5)

        with torch.autograd.detect_anomaly():

            for batch in dataloader:
                loss = self.training_step(batch, 0)
                print("loss", loss)

                loss.backward()
                optimizer.step()
                optimizer.zero_grad()

def train():

    # create domain
    nb_space = 128

    delta_x = 1.0 / nb_space
    delta_y = delta_x

    batch_size = 1

    # init dataset and dataloader
    path = "/app/data/2D_DarcyFlow_beta1.0_Train.hdf5"
    dataset = Darcy2DPDEDataset(path_hdf5=path, delta_x=delta_x, delta_y=delta_y)

    # divide into train and test
    train_size = int(0.95 * len(dataset))
    test_size = len(dataset) - train_size

    train_dataset, test_dataset = torch.utils.data.random_split(dataset, [train_size, test_size], generator=torch.Generator().manual_seed(42))

    dataloader_train = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=0)
    dataloader_test = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=0)

    # init model
    input_dim = 3
    modes = 32
    width = 32

    path_init = "/app/discretize_pinn_loss/data_init/init_solution.pt"

    # we create the model
    model = FNO2d(modes1=modes, modes2=modes,  width=width, input_dim=input_dim, path_init=path_init)

    # we create the burger function
    darcy_loss = DarcyLoss(index_derivative_node=0, index_derivative_x=0, index_derivative_y=1, delta_y=delta_y, delta_x=delta_x)

    # we create the trainer
    gnn_full = FnoFull(model, darcy_loss)

    # we create the trainer with the logger
    # init wandb key
    # read wandb_key/key.txt file to get the key
    with open("/app/wandb_key/key.txt", "r") as f:
        key = f.read()

    os.environ['WANDB_API_KEY'] = key
    wandb.init(project='2D_Darcy', entity='forbu14')

    wandb_logger = pl.loggers.WandbLogger(project="2D_Darcy", name="2D_Darcy")
    trainer = pl.Trainer(max_epochs=10, logger=wandb_logger, gradient_clip_val=0.5, accumulate_grad_batches=2)

    # test with test_train_step
    # gnn_full.test_train_step(dataloader_train)

    # we train
    trainer.fit(gnn_full, dataloader_train, dataloader_test)

    # we save the model
    torch.save(trainer.model.state_dict(), "models_params/darcy_model.pt")

    return gnn_full

if __name__ == "__main__":
    train()


import sqlite3
import openpyxl
import datetime
import os
import settings
from openpyxl.styles.fonts import Font
from openpyxl.styles.borders import Border, Side
from openpyxl.styles import Alignment
from openpyxl.styles import PatternFill

rss_new_url = settings.rss_new_url
rss_blog = settings.rss_blog
rss_blog_url = settings.rss_blog_url
dbname = settings.sqlite_dbname
now = datetime.datetime.now()
filename = settings.output_path + 'aws-rss_' + now.strftime('%Y%m%d.%H%M%S') + '.xlsx'

def main():
    rssawsnew_excel()
    if rss_blog == "true":
        rssawsblog_excel()

def rssawsnew_excel():
    workbook = openpyxl.Workbook()
    sheet1 = workbook.active
    sheet1.title = 'Work(Whats New)'

    border = Border(top=Side(style='thin', color='000000'),
        bottom=Side(style='thin', color='000000'),
        left=Side(style='thin', color='000000'),
        right=Side(style='thin', color='000000')
    )

    fill = PatternFill(patternType='solid',fgColor='00FFFF')

    sheet1['A2'].value = 'No.'
    sheet1['B2'].value = '投稿日'
    sheet1['C2'].value = 'サービス'
    sheet1['D2'].value = 'アップデートタイトル'
    sheet1['E2'].value = '概要'
    sheet1['F2'].value = '調査対象'
    sheet1.auto_filter.ref = "A2:F2"

    for abc in ['A','B','C','D','E','F']:
        sheet1[abc + '2'].font = Font(name='Meiryo',size='10',bold=True)
        sheet1[abc + '2'].border = border
        sheet1[abc + '2'].alignment = Alignment(horizontal='center', vertical='center')
        sheet1[abc + '2'].fill = fill

    db = sqlite3.connect(dbname)
    cur = db.cursor()
    cur.execute('SELECT no,published,category,title,url,summary FROM rss_new_table order by published desc')
    for i, row in enumerate (cur):
        sheet1['A' + str(i+3)].value =  (i+1)
        sheet1['B' + str(i+3)].value = row[1]
        sheet1['C' + str(i+3)].value = row[2]
        sheet1['D' + str(i+3)].value = row[3]
        sheet1['D' + str(i+3)].hyperlink = row[4]
        sheet1['D' + str(i+3)].style = "Hyperlink"
        sheet1['E' + str(i+3)].value = row[5]

    cur.close()

    for rows0 in sheet1['A3':'A1000']:
        for cell0 in rows0:
            cell0.font = Font(name='Meiryo',size='9')
            cell0.border = border
            cell0.alignment = Alignment(horizontal='center', vertical='center')

    for rows1 in sheet1['B3':'C1000']:
        for cell1 in rows1:
            cell1.font = Font(name='Meiryo',size='9')
            cell1.border = border
            cell1.alignment = Alignment(horizontal='left', vertical='center', wrapText=True)

    for rows2 in sheet1['D3':'D1000']:
        for cell2 in rows2:
            cell2.font = Font(name='Meiryo',size='9',color='0000FF')
            cell2.border = border
            cell2.alignment = Alignment(horizontal='left', vertical='top', wrapText=True)

    for rows3 in sheet1['E3':'E1000']:
        for cell3 in rows3:
            cell3.font = Font(name='Meiryo',size='9')
            cell3.border = border
            cell3.alignment = Alignment(horizontal='left', vertical='top', wrapText=True)

    for rows4 in sheet1['F3':'F1000']:
        for cell4 in rows4:
            cell4.font = Font(name='Meiryo',size='9')
            cell4.border = border
            cell4.alignment = Alignment(horizontal='center', vertical='center')

    sheet1.column_dimensions['A'].width = 4
    sheet1.column_dimensions['B'].width = 12
    sheet1.column_dimensions['C'].width = 20
    sheet1.column_dimensions['D'].width = 55
    sheet1.column_dimensions['E'].width = 82
    sheet1.column_dimensions['F'].width = 12
    sheet1.freeze_panes = 'A3'

    workbook.save(filename)

def rssawsblog_excel():
    workbook = openpyxl.load_workbook(filename)
    workbook.create_sheet( title="Work(AWS Blog)" )
    sheet2 = workbook['Work(AWS Blog)']

    border = Border(top=Side(style='thin', color='000000'),
        bottom=Side(style='thin', color='000000'),
        left=Side(style='thin', color='000000'),
        right=Side(style='thin', color='000000')
    )

    fill = PatternFill(patternType='solid',fgColor='00FFFF')

    sheet2['A2'].value = 'No.'
    sheet2['B2'].value = '投稿日'
    sheet2['C2'].value = 'サービス'
    sheet2['D2'].value = 'アップデートタイトル'
    sheet2['E2'].value = '概要'
    sheet2['F2'].value = '調査対象'
    sheet2.auto_filter.ref = "A2:F2"

    for abc in ['A','B','C','D','E','F']:
        sheet2[abc + '2'].font = Font(name='Meiryo',size='10',bold=True)
        sheet2[abc + '2'].border = border
        sheet2[abc + '2'].alignment = Alignment(horizontal='center', vertical='center')
        sheet2[abc + '2'].fill = fill

    db = sqlite3.connect(dbname)
    cur = db.cursor()
    cur.execute('SELECT no,published,category,title,url,summary FROM rss_blog_table order by published desc')
    for i, row in enumerate (cur):
        sheet2['A' + str(i+3)].value =  (i+1)
        sheet2['B' + str(i+3)].value = row[1]
        sheet2['C' + str(i+3)].value = row[2]
        sheet2['D' + str(i+3)].value = row[3]
        sheet2['D' + str(i+3)].hyperlink = row[4]
        sheet2['D' + str(i+3)].style = "Hyperlink"
        sheet2['E' + str(i+3)].value = row[5]

    cur.close()

    for rows0 in sheet2['A3':'A1000']:
        for cell0 in rows0:
            cell0.font = Font(name='Meiryo',size='9')
            cell0.border = border
            cell0.alignment = Alignment(horizontal='center', vertical='center')

    for rows1 in sheet2['B3':'C1000']:
        for cell1 in rows1:
            cell1.font = Font(name='Meiryo',size='9')
            cell1.border = border
            cell1.alignment = Alignment(horizontal='left', vertical='center', wrapText=True)

    for rows2 in sheet2['D3':'D1000']:
        for cell2 in rows2:
            cell2.font = Font(name='Meiryo',size='9',color='0000FF')
            cell2.border = border
            cell2.alignment = Alignment(horizontal='left', vertical='top', wrapText=True)

    for rows3 in sheet2['E3':'E1000']:
        for cell3 in rows3:
            cell3.font = Font(name='Meiryo',size='9')
            cell3.border = border
            cell3.alignment = Alignment(horizontal='left', vertical='top', wrapText=True)

    for rows4 in sheet2['F3':'F1000']:
        for cell4 in rows4:
            cell4.font = Font(name='Meiryo',size='9')
            cell4.border = border
            cell4.alignment = Alignment(horizontal='center', vertical='center')

    sheet2.column_dimensions['A'].width = 4
    sheet2.column_dimensions['B'].width = 12
    sheet2.column_dimensions['C'].width = 20
    sheet2.column_dimensions['D'].width = 55
    sheet2.column_dimensions['E'].width = 82
    sheet2.column_dimensions['F'].width = 12
    sheet2.freeze_panes = 'A3'

    workbook.save(filename)

if __name__ == "__main__":
    main()