__author__ = "https://github.com/kdha0727"

import os
import functools
import contextlib

import torch
import torch.distributed as dist
from torch.cuda import is_available as _cuda_available

RANK = 0
WORLD_SIZE = 1

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#                                       Setup Tools                                       #
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

def is_initialized():
    # if pytorch isn't compiled with c10d, is_initialized is omitted from namespace.
    # this function wraps
    """
    Returns c10d (distributed) runtime is initialized.
    """
    return dist.is_available() and getattr(dist, "is_initialized", lambda: False)()

def setup_dist(temp_dir, rank, world_size):
    """
    Set up a distributed process group.
    """

    if is_initialized():
        return True

    init_file = os.path.abspath(os.path.join(temp_dir, '.torch_distributed_init'))
    if os.name == 'nt':
        init_method = 'file:///' + init_file.replace('\\', '/')
        dist.init_process_group(
            backend='gloo', init_method=init_method, rank=rank, world_size=world_size)
    else:
        init_method = f'file://{init_file}'
        dist.init_process_group(
            backend='nccl', init_method=init_method, rank=rank, world_size=world_size)

    global RANK, WORLD_SIZE
    RANK = rank
    WORLD_SIZE = world_size

    torch.cuda.set_device(dev())
    torch.cuda.empty_cache()

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#                                      General Tools                                      #
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

@functools.lru_cache(maxsize=None)
def get_rank(group=None):
    if group is not None and is_initialized():
        return dist.get_rank(group=group)
    return RANK

@functools.lru_cache(maxsize=None)
def get_world_size(group=None):
    if group is not None and is_initialized():
        return dist.get_world_size(group=group)
    return WORLD_SIZE

def barrier(*args, **kwargs):
    if is_initialized():
        return dist.barrier(*args, **kwargs)

@contextlib.contextmanager
def synchronized_ops():
    barrier()
    yield
    barrier()
    return

@functools.lru_cache(maxsize=None)
def dev(group=None):
    """
    Get the device to use for torch.distributed.
    """
    if _cuda_available():
        return torch.device(get_rank(group))
    return torch.device("cpu")

def load_state_dict(local_or_remote_path, **kwargs):
    """
    Load a PyTorch file.
    """
    with open(local_or_remote_path, "rb") as f:
        return torch.load(f, **kwargs)

def broadcast(tensor, src=0, group=None, async_op=False):
    """
    Synchronize a Tensor across ranks from {src} rank. (default=0)
    :param tensor: torch.Tensor.
    :param src: source rank to sync params from. default is 0.
    :param group:
    :param async_op:
    """
    if not is_initialized():
        return
    with torch.no_grad():
        dist.broadcast(tensor, src, group=group, async_op=async_op)

def sync_params(params, src=0, group=None, async_op=False):
    """
    Synchronize a sequence of Tensors across ranks from {src} rank. (default=0)
    :param params: Sequence of torch.Tensor.
    :param src: source rank to sync params from. default is 0.
    :param group:
    :param async_op:
    """
    if not is_initialized():
        return
    for p in params:
        broadcast(p, src, group=group, async_op=async_op)


import boto3
import json
import uuid

print('Loading function')
def lambda_handler(event, context):
    bucketName = event['Records'][0]['s3']['bucket']['name']
    fileName = event['Records'][0]['s3']['object']['key']

    return detect_labels_and_put_dynamoDB(fileName, bucketName)

def detect_labels_and_put_dynamoDB(photo, bucket):
    rekognitionClient=boto3.client('rekognition', 'us-east-2')
    dynamoClient = boto3.client('dynamodb')

    response = rekognitionClient.detect_labels(Image={'S3Object':{'Bucket':bucket,'Name':photo}},
        MaxLabels=10)

    print('Detected labels for ' + photo)
    for label in response['Labels']:
        dynamoClient.put_item(
            TableName='RekognitionDetails',
            Item= {
                'ID' : {
                    'S': str(uuid.uuid4())
                },
                'Filename': {
                    'S': photo
                },
                'Category': {
                    'S' : label['Name']
                },
                'Confidence': {
                   'N': str(label['Confidence'])
                }
            })

def main():
    photo=''
    bucket=''
    label_count=detect_labels(photo, bucket)
    print("Labels detected: " + str(label_count))


# pylint: disable=W0611, E0401
"""
Main goal of this module is to scrape and parse data from "visityerevan.am" website
"""
import logging
import sys

from dataclasses import dataclass
from urllib.parse import urljoin

from httpx import Client
from selectolax.parser import HTMLParser, Node

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.WARNING)
stream_handler = logging.StreamHandler(sys.stdout)
formatter = logging.Formatter(
    "[%(asctime)s] %(levelname)s:%(name)s:%(lineno)d:%(message)s")
stream_handler.setFormatter(formatter)
logger.addHandler(stream_handler)

HEADERS = {"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) " +
           "AppleWebKit/537.36 (KHTML, like Gecko) " +
           "Chrome/109.0.0.0 Safari/537.36"}

@dataclass
class Event:
    """ Class contains all info about event """
    title: str
    description: str
    url_to_original: str
    time: str
    price: str
    img: str

@dataclass
class Response:
    """ Class contains html of page and info about existing of the next page """
    body_html: HTMLParser
    status_code: int

def serialize_event(event):
    """ Resulting format for each event """
    return {
        "id": "work in progress...",
        "type": "parsed_v1",
        "parserName": "visityerevan",
        "title": event.title,
        "description": event.description,
        "date": event.time,
        "durationInSeconds": 0,
        "location": {
                "country": "Armenia",
                "city": "Erevan",
        },
        "image": event.img,
        "price": {
            "amount": event.price,
            "currency": "AMD"
        },
        "timezone": {
            "timezoneName": "AMT",
            "timezoneOffset": "UTC +4",
        },
        "url": event.url_to_original,
    }

def get_page(client: Client, url: str) -> Response:
    """ Scrape html from page and check if next pages appears """
    resp = client.get(url, headers=HEADERS)
    html = HTMLParser(resp.text)

    return Response(body_html=html, status_code=resp.status_code)

def get_pages_amount(client: Client, url: str) -> int:
    """ func to get number of pages with events """
    resp = client.get(url, headers=HEADERS)
    html = HTMLParser(resp.text)

    pages_amount = html.css("ul[class='pagination justify-content-center'] >" +
                            "li[class='page-item']")[-1:][0].text()

    return int(pages_amount)

def is_valid(data):
    """ Helps us to catch website's structure changes """
    if data is None:
        logger.warning(
            "Seems that website changed structure. Please recheck code and website")

        return False
    else:
        return True

def parse_detail(blocks: list) -> list:
    """ Clean and prepare all data that we need """
    result = []

    # In this loop we will extract all
    #  Info that we can from each event's div
    for block in blocks:
        # Extract and prepare "time"
        month_day = block.css_first(
            "div[class='col-12 mt-n1'] > div")
        # Need validate data each parsing attempt
        if is_valid(month_day):
            month_day = month_day.text().replace('\n', '').strip()

        time = block.css_first(
            "div[class='text-grey text-md mb-2']")

        if is_valid(time):
            time = time.text().replace('\n', '').strip().split(' ')
            cleaned_time = f"{month_day} {time[-1:][0]}"
        else:
            cleaned_time = None
        # Extract and prepare "description"
        description = block.css_first("p")
        if is_valid(description):
            description = description.text().strip()
        # Clean and prepare "url"
        url = block.css_first("a").attrs["href"]
        if is_valid(url):
            url = "https://www.visityerevan.am" + url
        # Extract price
        price = ''
        cards = block.css("p.card-text > span")

        if len(cards) == 0:
            logger.warning(
                "Seems that website changed structure. Please recheck code and website")
        else:
            for card in cards:
                card = card.text()
                if "AMD" in card:
                    price = card.replace("AMD", "").strip()
                else:
                    price = "no info"
        # Extract img link
        img = block.css_first("img").attrs["src"]

        if is_valid(img):
            img = "https://www.visityerevan.am" + img
        # There is not need in cleaning "title"
        # With data we have create a new event object
        event = Event(
            title=block.css_first("h5").text(),
            description=description,
            url_to_original=url,
            time=cleaned_time,
            price=price,
            img=img
        )

        result.append(serialize_event(event))

    return result

def scrape_blocks(html: HTMLParser) -> list:
    """ Getting all divs with information from page """
    blocks = html.css("div[class='row px-lg-7']" +
                      " > div")

    return blocks

def pagination_loop(client: Client) -> list:
    """ Loop through all pages """
    url = "https://www.visityerevan.am/browse/things-to-do-events/ru/"
    # How many pages we will scrape
    pages_amount = get_pages_amount(client, url)
    # Blocks contains all divs that we need
    blocks = []
    # Iterating through all pages
    for page_number in range(1, pages_amount + 1):
        # Mutating a url to get page with current page number
        url = urljoin(url, f"?sel_filters=&current_page={page_number}")
        # Get object with scraped html markup from current page
        page = get_page(client, url)
        # Grad all divs with events data and append to list
        blocks += scrape_blocks(page.body_html)

    # Scraping is done, time to close session
    client.close()

    return blocks

async def scrape_website() -> list:
    """ Main function which contains all logic """
    # Start a new session
    client = Client()
    # Create list with all divs which contain info about events
    all_blocks = pagination_loop(client)
    # Parsing data from divs
    parsed_data = parse_detail(all_blocks)

    return parsed_data


import firebase_admin
import googleapiclient
from firebase_admin import credentials
from firebase_admin import db
import os
from os.path import join, dirname
from dotenv import load_dotenv
from XmlParser import XmlParser

class FirebaseService:
    dotenv_path = join(dirname(__file__), '.env')
    load_dotenv(dotenv_path)
    DATABASE_URL = os.environ.get("DATABASE_URL")
    if not firebase_admin._apps:
        print('初期化')
        cred = credentials.Certificate('hologram-test-firebase-adminsdk.json')

        firebase_admin.initialize_app(cred, {
            'databaseURL': DATABASE_URL,
        })

    ref_db = db.reference('/video')

    def __init__(self, video_item):
        self.video_item = video_item

    # FirestoreのドキュメントIDを取得
    def get_db_id(self):
        print("FirebaseService", 'get_db_id')
        id_list = []
        key_val = self.ref_db.get()
        # DB上に書き込まれたアイテムのvideoIdを取得
        for key, val in key_val.items():
            id_list.append(key)
        return id_list

    #
    def write_video_item(self):
        print("FirebaseService", "write_video_item")
        self.ref_db.update(self.video_item)

    def delete_video_item(self, update_db_items, xml_video_ids, error_channel_ids):
        print("FirebaseService", 'update_db_items', len(update_db_items), update_db_items)
        print("FirebaseService", 'xml_video_ids', len(xml_video_ids), xml_video_ids)
        print("FirebaseService", 'error_channel_ids', len(error_channel_ids), error_channel_ids)

        # 一週間以上まえのアイテムのみ抽出
        # （更新後のDB上のアイテム）-（XMLで取得したアイテム）
        last_week_ids = set(update_db_items).difference(set(xml_video_ids))
        print('last_week_ids', last_week_ids)
        for single_id in last_week_ids:
            db_channelId = self.ref_db.child(f'{single_id}').child('channelId').get()
            print('db_channel_id', db_channelId)
            # dbから取得したアイテムのチャンネルIDがエラーが発生したチャンネルIDリストの中に含まれていなければ削除
            # xml_parseで取得できなかったチャンネルの動画情報が削除されてしまうため
            if db_channelId not in set(error_channel_ids):
                print('delete', f'{single_id}')
                self.ref_db.child(f'{single_id}').delete()


import os
import sys

# 修改工作目录为程序所在目录，这样通过注册表实现开机自动启动时也能获取到正确的工作目录
# PS: 放到这个地方，是确保在所有其他初始化代码之前先修改掉工作目录
dirpath = os.path.dirname(os.path.realpath(sys.argv[0]))
old_path = os.getcwd()
os.chdir(dirpath)

import argparse
import datetime
import time
from multiprocessing import freeze_support

import psutil

import ga
from check_first_run import check_first_run_async
from config import config, load_config
from db_def import try_migrate_db
from first_run import is_weekly_first_run
from log import color, log_directory, logger
from main_def import (
    auto_send_cards,
    check_all_skey_and_pskey,
    check_djc_role_binding,
    check_proxy,
    check_update,
    get_user_buy_info,
    print_update_message_on_first_run_new_version,
    run,
    sas,
    show_ask_message_box,
    show_buy_info,
    show_extra_infos,
    show_lottery_status,
    show_multiprocessing_info,
    show_notices,
    show_pay_info,
    show_recommend_reward_tips,
    try_auto_update,
    try_auto_update_ignore_permission_on_special_case,
    try_join_xinyue_team,
    try_load_old_version_configs_from_user_data_dir,
    try_report_usage_info,
    try_save_configs_to_user_data_dir,
    try_take_dnf_helper_chronicle_task_awards_again_after_all_accounts_run_once,
    try_take_xinyue_team_award,
)
from pool import close_pool, init_pool
from qq_login import QQLogin
from show_usage import show_usage
from update import notify_manual_check_update_on_release_too_long
from usage_count import increase_counter
from util import (
    MiB,
    async_call,
    async_message_box,
    change_console_window_mode_async,
    change_title,
    clean_dir_to_size,
    disable_pause_after_run,
    disable_quick_edit_mode,
    is_run_in_github_action,
    kill_other_instance_on_start,
    pause,
    remove_old_version_portable_chrome_files,
    show_head_line,
    show_unexpected_exception_message,
)
from version import author, now_version, ver_time

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--no_max_console", default=False, action="store_true", help="是否不将窗口调整为最大化")
    parser.add_argument(
        "--wait_for_pid_exit",
        default=0,
        type=int,
        help="启动后是否等待对应pid的进程结束后再启动，主要用于使用配置工具启动小助手的情况，只有配置工具退出运行，自动更新才能正常进行",
    )
    parser.add_argument("--max_wait_time", default=5, type=int, help="最大等待时间")
    args = parser.parse_args()

    return args

def prepare_env():
    args = parse_args()

    # 最大化窗口
    if not args.no_max_console:
        logger.info("尝试调整窗口显示模式，打包exe可能会运行的比较慢")
        change_console_window_mode_async()

    if args.wait_for_pid_exit != 0:
        # 通过配置工具打开
        increase_counter(ga_category="open_by", name="config_tool", ga_misc_params={"dr": "config_tool"})
        logger.info(f"等待pid为{args.wait_for_pid_exit}的配置工具退出运行，从而确保可能有的自动更新能够正常进行，最大将等待{args.max_wait_time}秒")

        wait_time = 0.0
        retry_time = 0.1
        while wait_time <= args.max_wait_time:
            if not psutil.pid_exists(args.wait_for_pid_exit):
                logger.info("配置工具已成功退出，将开始运行小助手~")
                break

            time.sleep(retry_time)
            wait_time += retry_time
    else:
        # 直接打开
        increase_counter(ga_category="open_by", name="directly", ga_misc_params={"dr": "directly"})

def main():
    try_migrate_db()

    increase_counter(name="run/begin", ga_type=ga.GA_REPORT_TYPE_PAGE_VIEW)

    prepare_env()

    # 启动时检查是否需要同步本机数据目录备份的旧版本配置
    try_load_old_version_configs_from_user_data_dir()

    change_title()

    print_update_message_on_first_run_new_version()

    logger.warning(f"开始运行DNF蚊子腿小助手，ver={now_version} {ver_time}，powered by {author}")
    logger.warning(color("fg_bold_cyan") + "如果觉得我的小工具对你有所帮助，想要支持一下我的话，可以帮忙宣传一下或打开付费指引/支持一下.png，扫码打赏哦~")

    # 读取配置信息
    load_config("config.toml", "config.toml.local")
    cfg = config()

    if len(cfg.account_configs) == 0:
        raise Exception("未找到有效的账号配置，请检查是否正确配置。ps：多账号版本配置与旧版本不匹配，请重新配置")

    try_auto_update_ignore_permission_on_special_case(cfg)

    notify_manual_check_update_on_release_too_long(cfg.common)

    check_proxy(cfg)

    try_report_usage_info(cfg)

    if cfg.common.disable_cmd_quick_edit:
        disable_quick_edit_mode()

    show_notices()

    if cfg.common.allow_only_one_instance:
        logger.info("当前仅允许单个实例运行，将尝试干掉其他实例~")
        async_call(kill_other_instance_on_start)
    else:
        logger.info("当前允许多个实例同时运行~")

    pool_size = cfg.get_pool_size()
    init_pool(pool_size)

    change_title(multiprocessing_pool_size=pool_size, enable_super_fast_mode=cfg.common.enable_super_fast_mode)

    show_multiprocessing_info(cfg)

    account_names = []
    for account_cfg in cfg.account_configs:
        account_names.append(account_cfg.name)

    logger.info(f"当前共配置{len(account_names)}个账号，具体如下：{account_names}")

    clean_dir_to_size(log_directory, cfg.common.max_logs_size * MiB, cfg.common.keep_logs_size * MiB)
    clean_dir_to_size(f"utils/{log_directory}", cfg.common.max_logs_size * MiB, cfg.common.keep_logs_size * MiB)

    current_chrome_version = QQLogin(cfg.common).get_chrome_major_version()
    remove_old_version_portable_chrome_files(current_chrome_version)

    show_ask_message_box(cfg)

    # 检查是否有更新，用于提示未购买自动更新的朋友去手动更新~
    if cfg.common.check_update_on_start:
        check_update(cfg)

    check_all_skey_and_pskey(cfg)

    check_djc_role_binding()

    # 确保道聚城绑定OK后在活动运行同时进行异步的弹窗提示
    check_first_run_async(cfg)

    # 挪到所有账号都登陆后再尝试自动更新，从而能够判定是否已购买DLC
    try_auto_update(cfg)

    # 查询付费信息供后面使用
    show_head_line("查询付费信息")
    logger.warning("开始查询付费信息，请稍候~")
    user_buy_info = get_user_buy_info(cfg.get_qq_accounts())
    show_buy_info(user_buy_info, cfg, need_show_message_box=False)

    sas(cfg, "启动时展示账号概览", user_buy_info)

    # 预先尝试创建和加入固定队伍，从而每周第一次操作的心悦任务也能加到队伍积分中
    try_join_xinyue_team(cfg, user_buy_info)

    # 正式进行流程
    run(cfg, user_buy_info)

    try_take_dnf_helper_chronicle_task_awards_again_after_all_accounts_run_once(cfg, user_buy_info)

    # 尝试领取心悦组队奖励
    try_take_xinyue_team_award(cfg, user_buy_info)

    # # 尝试派赛利亚出去打工
    # try_xinyue_sailiyam_start_work(cfg)

    # 活动开启关闭时调这个开关即可
    enable_card_lottery = True

    if enable_card_lottery:
        auto_send_cards(cfg)

    show_extra_infos(cfg)
    sas(cfg, "运行完毕展示账号概览", user_buy_info)

    if enable_card_lottery:
        show_lottery_status("卡片赠送完毕后展示各账号抽卡卡片以及各礼包剩余可领取信息", cfg, need_show_tips=True)

    show_pay_info(cfg)

    show_recommend_reward_tips(user_buy_info)

    # 显示小助手的使用概览
    if cfg.common._show_usage:
        show_usage()

    # 运行结束展示下多进程信息
    show_multiprocessing_info(cfg)

    # 检查是否有更新，用于提示未购买自动更新的朋友去手动更新~
    if cfg.common.check_update_on_end:
        check_update(cfg)

    # 运行完毕备份配置到本机数据目录
    try_save_configs_to_user_data_dir()

    increase_counter(name="run/end", ga_type=ga.GA_REPORT_TYPE_PAGE_VIEW)

    show_head_line("运行完毕")

def main_wrapper():
    freeze_support()

    logger.info(color("bold_green") + f"已将工作目录设置为小助手所在目录：{dirpath}，之前为：{old_path}")

    try:
        run_start_time = datetime.datetime.now()
        main()
        total_used_time = datetime.datetime.now() - run_start_time
        logger.warning(color("fg_bold_yellow") + f"运行完成，共用时{total_used_time}")

        # 如果总用时太高的情况时，尝试提示开启多进程和超快速模式
        cfg = config()
        if total_used_time > datetime.timedelta(minutes=10) and (
            not cfg.common.enable_multiprocessing or not cfg.common.enable_super_fast_mode
        ):
            msg = (
                f"当前累计用时似乎很久({total_used_time})，是否要尝试多进程和超快速模式？\n"
                "多进程模式下，将开启多个进程并行运行不同账号的领取流程\n"
                "额外开启超快速模式，会进一步将不同账号的不同活动都异步领取，进一步加快领取速度\n"
                "\n"
                "如果需要开启，请打开配置工具，在【公共配置】tab中勾选【是否启用多进程功能】和【是否启用超快速模式（并行活动）】"
            )
            logger.warning(color("bold_cyan") + msg)
            if is_weekly_first_run("用时过久提示"):
                async_message_box(msg, "用时过久", print_log=False)

        # 按照分钟级别来统计使用时长
        total_minutes = int(total_used_time.total_seconds()) // 60
        increase_counter(ga_category="run_used_time_minutes", name=total_minutes)
    except Exception as e:
        show_unexpected_exception_message(e)
        # 如果在github action，则继续抛出异常
        if is_run_in_github_action():
            raise e
    finally:
        # 暂停一下，方便看结果
        if not disable_pause_after_run() and not is_run_in_github_action():
            async_call_close_pool_after_some_time()
            pause()
        close_pool()

def async_call_close_pool_after_some_time():
    def _close():
        wait_time = 10 * 60
        logger.info(f"{wait_time} 秒后将自动关闭进程池，方便有足够时间查看进程池中触发的弹窗信息")
        time.sleep(wait_time)
        close_pool()

    async_call(_close)

if __name__ == "__main__":
    main_wrapper()


import got3
import pymongo
from vaderSentiment.vaderSentiment import SentimentIntensityAnalyzer

# connect to mongo deamon
connection = pymongo.MongoClient("mongodb://localhost")

# connect to the collection called uber_tweets in the kubrick db
db = connection.kubrick.uberban_tweets

count = 0

try:
    while True:
        tweetCriteria = got3.manager.TweetCriteria().setSince("2017-09-22").setQuerySearch("uberban")
        #tweetCriteria = got3.manager.TweetCriteria().setQuerySearch("uberban")
        tweet = got3.manager.TweetManager.getTweets(tweetCriteria)[count]
        sent = SentimentIntensityAnalyzer().polarity_scores(tweet.text)['compound']
        print(tweet.text)
        print(sent)
        print(tweet.date)
        db.insert_many([{"tweet": tweet.text, "sentiment": sent}])
        count += 1
except:
    print("tweet scrape ended with {no_tweets} tweets".format(no_tweets = count))

Old UberBan Tweets Scrape + Sentiment Analysis/old_tweets.py


import joblib

model = None

def init_model(
    db,
    model_themes_path='./flaskr/model/log_reg_themes',
    model_cats_path='./flaskr/model/log_reg_cats'
):
    global model

    cur = db.cursor()
    query = """
    select id from theme order by id;
    """
    cur.execute(query)
    theme_ids = [id[0] for id in cur.fetchall()]
    cur.close()

    cur = db.cursor()
    query = """
    select id from category order by id;
    """
    cur.execute(query)
    cats_ids = [id[0] for id in cur.fetchall()]
    cur.close()

    model = Model(theme_ids, cats_ids, model_themes_path, model_cats_path)

def get_model():
    return model

class Model:
    def __init__(self, theme_ids, cats_id, model_themes_path, model_cats_path):
        self.model_themes = joblib.load(model_themes_path)
        self.model_cats = joblib.load(model_cats_path)
        self.theme_ids = theme_ids
        self.cats_ids = cats_id

    def analyse_theme(self, text: str, detailed_text:str = None, probs_count=3):
        probs = self.model_themes.predict_proba([text])[0]
        if detailed_text:
            d_probs = self.model_themes.predict_proba([detailed_text])[0]
            probs = probs * d_probs

        most_likely_probs = None
        if probs_count <= 0:
            most_likely_probs = [self.theme_ids[id]
                             for id in probs.argsort().tolist()[:][::-1]]
        else:
            most_likely_probs = [self.theme_ids[id]
                             for id in probs.argsort().tolist()[-probs_count:][::-1]]

        return most_likely_probs

    def analyse_cat(self, text: str, detailed_text:str = None):
        probs = self.model_cats.predict_proba([text])[0]

        if detailed_text:
            d_probs = self.model_cats.predict_proba([detailed_text])[0]
            probs = probs * d_probs

        most_likely_probs = [self.cats_ids[id]
                             for id in probs.argsort().tolist()[:][::-1]]
        return most_likely_probs


#!/usr/bin/env python3

import random
from typing import Tuple

from functions.aes import AESCipher, pkcs7_pad, get_blocks, gen_random_bytes

def _encryption_oracle(bytes_: bytes) -> Tuple[bytes, str]:
    key = gen_random_bytes(16)
    iv = gen_random_bytes(16)
    prefix = gen_random_bytes(random.randint(5, 10))
    suffix = gen_random_bytes(random.randint(5, 10))
    pt = prefix + bytes_ + suffix

    cbc_mode = random.choice([True, False])

    if cbc_mode:
        cbc = AESCipher(AESCipher.MODE_CBC, key, iv=iv)
        ct = cbc.encrypt(pkcs7_pad(pt))
        answer = "cbc"
    else:
        ecb = AESCipher(AESCipher.MODE_ECB, key)
        ct = ecb.encrypt(pkcs7_pad(pt))
        answer = "ecb"

    return ct, answer

def challenge11() -> bool:
    pt = bytes(gen_random_bytes(1) * random.randint(100, 200))
    ct, answer = _encryption_oracle(pt)
    blocks = get_blocks(ct)
    unique_blocks = len(set(blocks))

    guess = "cbc" if len(blocks) == unique_blocks else "ecb"

    return True if guess == answer else False

if __name__ == "__main__":
    for _ in range(100):
        assert challenge11(), "The result does not match the expected value"

    print("Ok")


from nose.tools import eq_
from mock import patch, Mock, sentinel

from noderunner.process import open_process

@patch("subprocess.Popen", return_value=sentinel.proc)
def test_open_process(p):
    ret = open_process(sentinel.fd,
                       sentinel.secret,
                       nodepath=sentinel.node_path)

    eq_(ret, sentinel.proc)

    p.assert_called_once()


"""
@author:  Yuhao Cheng
@contact: yuhao.cheng[at]outlook.com
"""
#!!!!! ignore the warning messages
import warnings
warnings.filterwarnings('ignore')
import os
import pickle
import math
import torch
import time
import numpy as np
from PIL import Image
from collections import OrderedDict
import torchvision.transforms as T
import torchvision.transforms.functional as tf
from torch.utils.data import DataLoader

from pyanomaly.core.utils import AverageMeter, flow_batch_estimate, tensorboard_vis_images, make_info_message, ParamSet
from pyanomaly.datatools.evaluate.utils import psnr_error
from ..abstract.base_engine import BaseTrainer, BaseInference
from ..engine_registry import ENGINE_REGISTRY

__all__ = ['MEMAETrainer', 'MEMAEInference']

@ENGINE_REGISTRY.register()
class MEMAETrainer(BaseTrainer):
    NAME = ["MEMAE.TRAIN"]
    def custom_setup(self):

        # basic meter
        self.loss_meter_MemAE = AverageMeter(name='loss_memae')

    def train(self,current_step):
        # Pytorch [N, C, D, H, W]
        # initialize
        start = time.time()
        self.set_requires_grad(self.MemAE, True)
        self.MemAE.train()
        writer = self.kwargs['writer_dict']['writer']
        global_steps = self.kwargs['writer_dict']['global_steps_{}'.format(self.kwargs['model_type'])]

        # get the data
        data, anno, meta = next(self._train_loader_iter)  # the core for dataloader
        self.data_time.update(time.time() - start)

        input_data = data.cuda()

        # True Process =================Start===================
        output_rec, att = self.MemAE(input_data)
        loss_rec = self.rec_loss(output_rec, input_data)
        loss_mem = self.mem_loss(att)
        loss_memae_all = self.loss_lamada['rec_loss'] * loss_rec + self.loss_lamada['mem_loss'] * loss_mem
        # loss_memae_all = self.loss_lamada['rec_loss'] * loss_rec
        self.optim_MemAE.zero_grad()
        # with torch.autograd.set_detect_anomaly(True):
        loss_memae_all.backward()
        self.optim_MemAE.step()
        self.loss_meter_MemAE.update(loss_memae_all.detach())

        if self.config.TRAIN.adversarial.scheduler.use:
            self.lr_memae.step()
        # ======================End==================

        self.batch_time.update(time.time() - start)

        if (current_step % self.steps.param['log'] == 0):
            msg = make_info_message(current_step, self.steps.param['max'], self.kwargs['model_type'], self.batch_time,
                                    self.config.TRAIN.batch_size, self.data_time, [self.loss_meter_MemAE])
            self.logger.info(msg)
        writer.add_scalar('Train_loss_MemAE', self.loss_meter_MemAE.val, global_steps)

        if (current_step % self.steps.param['vis'] == 0):
            vis_objects = OrderedDict({
                'train_output_rec_memeae': output_rec.detach(),
                'train_input':  input_data.detach()
            })
            tensorboard_vis_images(vis_objects, writer, global_steps, self.normalize.param['train'])
        global_steps += 1

        # reset start
        start = time.time()

        # self.saved_model = {'MemAE':self.MemAE}
        self.saved_model['MemAE'] = self.MemAE
        # self.saved_optimizer = {'optim_MemAE': self.optim_MemAE}
        self.saved_optimizer['optimizer_MemAE']= self.optim_MemAE
        # self.saved_loss = {'loss_MemAE':self.loss_meter_MemAE.val}
        self.saved_loss['loss_MemAE'] = self.loss_meter_MemAE.val
        self.kwargs['writer_dict']['global_steps_{}'.format(self.kwargs['model_type'])] = global_steps

@ENGINE_REGISTRY.register()
class MEMAEInference(BaseInference):
    NAME = ["MEMAE.INFERENCE"]
    def inference(self):
        for h in self._hooks:
            h.inference()


import requests
import pandas as pd
import arrow
import warnings
import io
from email.mime.text import MIMEText
from email.mime.multipart import MIMEMultipart
import smtplib
import logging
warnings.filterwarnings('ignore', 'Unverified HTTPS request')
url = "https://protect.cylance.com/Reports/ThreatDataReportV1/memoryprotection/"
token = "Token"
fullurl = (url + token)
path = 'Filepath'
logfile = 'FilePath'
nv = arrow.now()
date = nv.shift(days=-1).format('M/D/YYYY')
date2 = nv.shift(days=-1).format('YYYYMD')
def email_send(email_data):
    from_addr = "EmailFrom"
    to_addr = "EmailTo"
    to_list = ["To_List"]
    msg = MIMEMultipart()
    msg['From'] = from_addr
    msg['To'] = to_addr
    msg['Subject'] = "Cylance Exploit Attempts for %s" %(date)
    part2 = MIMEText(email_data, 'html')
    msg.attach(part2)
    server = smtplib.SMTP("smtpRelay", 25)
    server.sendmail(from_addr,to_list,msg.as_string())
    server.quit()

if __name__ == '__main__':

    logging.basicConfig(filename=logfile, level=logging.INFO, format='%(asctime)s %(levelname)-8s %(message)s')
    logging.info('Requesting MEM TDR')
    urlData = requests.get(fullurl).content
    rawData = pd.read_csv(io.StringIO(urlData.decode('utf-8')))
    logging.info('Creating dataframe')
    df2 = pd.DataFrame(rawData)
    logging.info('Dropping Serial Column')
    df3 = df2.drop(["Serial Number",], axis = 1)
    logging.info('Filtering Data by date')
    test3 = (df3[df3['ADDED'].str.contains(date)])
    logging.info('Selecting Column Headers')
    output = (test3[["Device Name","ADDED",'PROCESS NAME','ACTION','TYPE','USER NAME']])
    print(output)

    if output.empty:
        logging.info('No Memory Exploit for %s' % (date))

    else:
        logging.info('Creating CSV')
        output.to_csv(path + date2 + "mem.csv", index=False)
        logging.info('CSV Created')
        logging.info('Converting Data to HTML')
        email_data = output.to_html(index = False)
        logging.info('Preparing Email')
        email_send(email_data)
        logging.info('Email Sent')


import tensorflow as tf
from utils.nn import linear
from .tdnn import TDNN

def embed_characters(input, vocab_size, embed_dim=40, scope=None, reuse=None,
             use_batch_norm=True, use_highway=True, highway_layers=2):
    """ Character-level embedding """

    with tf.variable_scope(scope or 'Embedder') as scope:
        if reuse: scope.reuse_variables()
        input = tf.unpack(tf.transpose(input, [1, 0, 2]))  # L * [N, W]
        embedding = tf.get_variable('embedding', [vocab_size, embed_dim])

        embedded = []
        for word in input:
            embed = tf.nn.embedding_lookup(embedding, word) # [N, W, d]
            conved = TDNN(embed, embed_dim)

            if use_batch_norm:
                conved = batch_norm(conved)

            if use_highway:
                conved = highway(conved, conved.get_shape()[1], highway_layers, 0)

            embedded.append(conved)
            scope.reuse_variables()

    return embedded

def batch_norm(x, epsilon=1e-5):
    shape = x.get_shape().as_list()

    with tf.variable_scope('BatchNorm'):
        gamma = tf.get_variable("gamma", [shape[-1]],
            initializer=tf.random_normal_initializer(1., 0.02))
        beta = tf.get_variable("beta", [shape[-1]],
            initializer=tf.constant_initializer(0.))

        mean, variance = tf.nn.moments(x, [0, 1])

        return tf.nn.batch_norm_with_global_normalization(
            x, mean, variance, beta, gamma, epsilon,
            scale_after_normalization=True)

def highway(input_, size, layer_size=1, bias=-2, f=tf.nn.relu):
    """Highway Network (cf. http://arxiv.org/abs/1505.00387).

    t = sigmoid(Wy + b)
    z = t * g(Wy + b) + (1 - t) * y
    where g is nonlinearity, t is transform gate, and (1 - t) is carry gate.
    """
    with tf.variable_scope('Highway'):
        output = input_
        for idx in range(layer_size):
            output = f(linear(output, size, 0, scope='output_lin_%d' % idx, init='he'))

            transform_gate = tf.sigmoid(linear(input_, size, 0, scope='transform_lin_%d' % idx) + bias)
            carry_gate = 1. - transform_gate

            output = transform_gate * output + carry_gate * input_

    return output


""" https://adventofcode.com/2020/day/10 """

from typing import Dict, List
from collections import defaultdict

Adapters = List[int]

def part1(adapters: Adapters) -> int:
    """ O(nLogn) solution """

    jolts = 0
    diffs: Dict[int, int] = defaultdict(int)
    for adapter in sorted(adapters):
        diffs[adapter - jolts] += 1
        jolts = adapter

    return diffs[1] * (diffs[3] + 1)

def part2(adapters: Adapters) -> int:
    """ O(nLogn) solution """

    adapters = sorted(adapters)
    adapters = [0] + adapters + [max(adapters)+3]
    paths = {adapters[0]: 1}
    for x in adapters[1:]:
        paths[x] = sum(paths[x - y] for y in range(1, 4) if x - y in paths)

    return paths[adapters[-1]]

if __name__ == "__main__":
    TEST1 = [int(line.strip()) for line in open("tests/d10.txt", "r")]
    TEST2 = [int(line.strip()) for line in open("tests/d10_2.txt", "r")]
    PUZZLE = [int(line.strip()) for line in open("puzzles/d10.txt", "r")]

    assert part1(TEST1) == 35
    assert part1(TEST2) == 220
    assert part2(TEST1) == 8
    assert part2(TEST2) == 19208

    print(f"Part 1: {part1(PUZZLE)}")
    print(f"Part 2: {part2(PUZZLE)}")


import rest_framework.authentication
from drf_yasg.views import get_schema_view
from drf_yasg import openapi
from rest_framework import permissions

from user.auth.auth import JwtQueryParamsAuthentication

schema_view = get_schema_view(
    openapi.Info(
        title="接口文档",
        default_version="1.0",
        terms_of_service='',
        contact=openapi.Contact(name="Andy Z Wright", email="andyzwright021@gmail.com"),
        license=openapi.License(name="MIT LICENCE"),
    ),
    public=True,
    permission_classes=(permissions.AllowAny,),
    # authentication_classes=(JwtQueryParamsAuthentication,)
    authentication_classes=(),
)


import numpy as np
import pandas as pd
import time
from metric import SampleScore,EventScore, AdjustedMutualInfoScore
from joblib import Parallel, delayed

class Experiment:

    def __init__(self,algorithms:list, configurations:list, thresholds = np.linspace(0,1,101),njobs=1,verbose = True) -> None:

        """Initialization

        Args:
            algorithms (list): list of algorithm classes
            configurations (list): list of list of configurations as dictionnaries for each algorithms classes
            thresholds (np.ndarray, optional): numpy array of thresholds to consider for the event based metric. Defaults to numpy.linspace(0,1,101).
        """
        self.algorithms = algorithms
        self.configurations = configurations
        self.thresholds = thresholds
        self.njobs = njobs
        self.verbose = verbose

    def compute_scores(self,label,prediction):

        single_pred = np.clip(np.sum(prediction,axis=0),0,1).reshape(1,-1)
        single_label = np.clip(np.sum(label,axis=0),0,1).reshape(1,-1)

        scores = []

        #single sample score
        p,r,f = SampleScore().score(single_label,single_pred)
        scores.append(["sss-precision",p])
        scores.append(["sss-recall",r])
        scores.append(["sss-fscore",f])

        #sample score
        p,r,f = SampleScore().score(label,prediction)
        scores.append(["ss-precision",p])
        scores.append(["ss-recall",r])
        scores.append(["ss-fscore",f])

        # weigthed sample score
        p,r,f = SampleScore(averaging="weighted").score(label,prediction)
        scores.append(["w-ss-precision",p])
        scores.append(["w-ss-recall",r])
        scores.append(["w-ss-fscore",f])

        #single event score
        lp,lr,lf = EventScore().score(single_label,single_pred,self.thresholds)
        for t,p,r,f in zip(self.thresholds,lp,lr,lf):
            scores.append([f"ses-precision_{np.round(t,2)}",p])
            scores.append([f"ses-recall_{np.round(t,2)}",r])
            scores.append([f"ses-fscore_{np.round(t,2)}",f])
        scores.append(["ses-auc-precision",np.mean(lp)])
        scores.append(["ses-auc-recall",np.mean(lr)])
        scores.append(["ses-auc-fscore",np.mean(lf)])

        #event score
        lp,lr,lf = EventScore().score(label,prediction,self.thresholds)
        for t,p,r,f in zip(self.thresholds,lp,lr,lf):
            scores.append([f"es-precision_{np.round(t,2)}",p])
            scores.append([f"es-recall_{np.round(t,2)}",r])
            scores.append([f"es-fscore_{np.round(t,2)}",f])
        scores.append(["es-auc-precision",np.mean(lp)])
        scores.append(["es-auc-recall",np.mean(lr)])
        scores.append(["es-auc-fscore",np.mean(lf)])

        # weighted event score
        lp,lr,lf = EventScore(averaging="weighted").score(label,prediction,self.thresholds)
        for t,p,r,f in zip(self.thresholds,lp,lr,lf):
            scores.append([f"w-es-precision_{np.round(t,2)}",p])
            scores.append([f"w-es-recall_{np.round(t,2)}",r])
            scores.append([f"w-es-fscore_{np.round(t,2)}",f])
        scores.append(["w-es-auc-precision",np.mean(lp)])
        scores.append(["w-es-auc-recall",np.mean(lr)])
        scores.append(["w-es-auc-fscore",np.mean(lf)])

        #ajusted mutual information
        scores.append(["amis",AdjustedMutualInfoScore().score(label,prediction)])

        return scores

    def signal_algo_class_experiement(self,signal_idx,signal,label,algo_class,config,config_idx):
        "Return a DF"
        #keep only labels row that are activated by the signal
        label = label[label.sum(axis=1)>0]

        #update the number of patterns to predict if required
        t_config = config.copy()
        if ("n_patterns" in t_config.keys()):
            if (isinstance(t_config["n_patterns"],int)):
                t_config["n_patterns"] = label.shape[0]
            else:
                t_config["n_patterns"] = None

        try:
            #get predictions
            algo = algo_class(**t_config)
            start = time.time()
            algo.fit(signal)
            end = time.time()

            #compute scores
            scores = self.compute_scores(label,algo.prediction_mask_)

            tdf = pd.DataFrame(scores,columns=["metric","score"])
            tdf["algorithm"] = algo_class.__name__
            tdf["config_idx"] = config_idx
            tdf["execution_time"] = end - start
            tdf["signal_idx"] = signal_idx
            tdf["n_patterns"] = label.shape[0]
            tdf["predicted_n_patterns"] = algo.prediction_mask_.shape[0]

            if self.verbose:
                s1 = np.round(tdf[tdf["metric"] == "es-auc-fscore"].score.values[0],2)
                s2 = np.round(tdf[tdf["metric"] == "amis"].score.values[0],2)
                print(f"signal_id: {signal_idx}, algo: {algo_class.__name__}, config_id: {config_idx}, f-auc: {s1}, ami: {s2}")

            return tdf

        except:
            s= f"signal_id: {signal_idx}, algo: {algo_class.__name__}, config_id: {config_idx} failed to fit."
            if self.verbose:
                print(s)
            if self.logs_path_ is not None:
                with open(self.logs_path_,"a") as f:
                    f.write(s +"\n")

    def run_experiment(self,dataset:np.ndarray,labels:np.ndarray,backup_path = None,batch_size=10,logs_path = None,verbose = True)->np.ndarray:
        """_summary_

        Args:
            dataset (np.ndarray): array of signals, signal shape (L,), variable length allowed
            labels (np.ndarray): array of labels, label shape (L,), variable length allowed
            signal_configs (pd.DataFrame, optional): Dataframe containing the configuration of the synthetic generator for each signals.
            backup_path (str, optional): Path to store df in case of big experiment. If None no saving. Defaults to None.
            batch_size (int, optional)
            verbose (bool, optional): verbose. Defaults to True.

        Returns:
            pd.DataFrame: scores_df
        """
        self.logs_path_ = logs_path

        n_signals = len(dataset)
        n_configs = np.sum([len(conf) for conf in self.configurations])
        total = n_signals*n_configs

        if backup_path != None:
            n_batches  = n_signals//batch_size
            if n_batches >0:
                batches =[zip(dataset[i*batch_size:(i+1)*batch_size],labels[i*batch_size:(i+1)*batch_size]) for i in range(n_batches)]
            else:
                batches = []
            if n_signals % batch_size !=0:
                batches.append(zip(dataset[n_batches*batch_size:],labels[n_batches*batch_size:]))
        else:
            batches = [zip(dataset,labels)]

        self.df_ = pd.DataFrame()

        counts = 0
        for batch in batches:
            results = Parallel(n_jobs=self.njobs)(
                delayed(self.signal_algo_class_experiement)(counts+id_s,signal,label,algo,config,id_c)
                for id_s,(signal,label) in enumerate(batch)
                for id_a,algo in enumerate(self.algorithms)
                for id_c,config in enumerate(self.configurations[id_a])
                )
            counts = min(counts+batch_size,n_signals)
            self.df_= pd.concat((self.df_,*results)).reset_index(drop = True)
            self.df_ = self.df_.astype({'metric':str, "score":float, "algorithm":str,'config_idx':int,"signal_idx":int, "n_patterns":int, "predicted_n_patterns":int})

            if backup_path != None:
                self.df_.to_csv(backup_path)

            if verbose:
                print(f"Achieved [{counts*n_configs}/{total}]")

        return self.df_


from typing import List

class Solution:
    def findPeakElement(self, nums: List[int]) -> int:
        n = len(nums)
        left = 0
        right = n - 1

        def get_num(i):
            if i == -1 or i == n:
                return float('-inf')
            return nums[i]

        ans = -1
        while right >= left:
            mid = left + (right - left) // 2
            if get_num(mid - 1) < get_num(mid) > get_num(mid + 1):
                return mid
            elif get_num(mid) < get_num(mid + 1):
                left = mid + 1
            else:
                right = mid - 1
        return ans


import boto3
import traceback
import datetime
import os

from botocore.exceptions import ClientError
from ..models.bucket import Bucket
from ..util.preprocessor import preprocess

"""
S3 functions
"""

def get_active_bucket_or_create_new(username):
    """
    Returns the user's current active bucket. If there are no buckets,
    or all the users buckets are full, a new one will be created first.

    returns: Bucket object
    """

    try:
        # queries database for non-full buckets
        range_cond = Bucket.sort.startswith("BUCKET_")
        buckets = Bucket.query(hash_key=username, range_key_condition=range_cond)

        # return existing if not full or create new
        for buck in buckets:
            if not buck.full:
                return buck
        else:
            bucket = create_bucket(username)
            return bucket
    except Exception as e:
        traceback.print_exc()
        response_object = {
            "status": "failed",
            "message": "Could not query buckets in DB. {}".format(e),
        }
        return response_object, 500

def create_bucket(username, region="eu-central-1"):
    """
    Creates an S3 bucket in S3.
    Naming format: 'flasktextapi-{ENV}-{USERNAME}-BUCKET{BUCKET_ID}'
    IMPORTANT: underscores in usernames are converted to dashes.

    returns: bucket
    """
    # create S3 bucket
    try:
        bucket_id = _generate_bucket_id(username)
        username_conv = username.replace("_", "-")
        bucket_name = "flasktextapi-{env}-{username}-bucket{id}".format(
            env=os.getenv("BOILERPLATE_ENV"), username=username_conv, id=bucket_id
        )
        bucket = boto3.resource("s3").Bucket(bucket_name)
        location = {"LocationConstraint": region}
        response = bucket.create(CreateBucketConfiguration=location)
        if response["ResponseMetadata"]["HTTPStatusCode"] == 200:
            try:
                db_bucket = _create_db_bucket(username, bucket_id, bucket_name)
                return db_bucket
            except Exception:
                traceback.print_exc()
                response_object = {
                    "status": "failed",
                    "message": "Bucket created successfully but bucket details could not be stored to database.",
                }
                return response_object, 500
        else:
            response_object = {"status": "failed", "message": "could not create bucket"}
            return response_object, 500
    except Exception as e:
        traceback.print_exc()
        response_object = {
            "status": "failed",
            "message": "Could not create bucket. {}".format(e),
        }
        return response_object, 500

def add_file(username, input, bucket_name, id, region="eu-central-1"):
    """
    Adds a text to an S3 bucket.

    Naming format of file: 'unprocessed_{username}_{text.public_id}'

    @return: Name of the file as String.
    """

    # check input type
    if not isinstance(input, str):
        raise ValueError("Text needs to be a String.")

    bucket = boto3.resource("s3").Bucket(bucket_name)
    key = "unprocessed_{id}_{username}".format(id=id, username=username)
    bucket.put_object(Body=bytes(input, "utf-8"), Key=key)
    return key

def add_preprocessed_file(username, input, bucket_name, id, region="eu-central-1"):
    """
    Adds a text to an S3 bucket.

    Naming format of file: 'unprocessed_{username}_{text.public_id}'

    @return: Name of the file as String.
    """

    # check input type
    if not isinstance(input, str):
        raise ValueError("Text needs to be a String.")

    # preprocess input
    prepr_input = preprocess(input)

    bucket = boto3.resource("s3").Bucket(bucket_name)
    key = "preprocessed_{id}_{username}".format(id=id, username=username)
    bucket.put_object(Body=bytes(prepr_input, "utf-8"), Key=key)
    return key

def get_object(bucket_name, key):
    """
    Fetches an object from S3.

    returns: String
    """
    s3 = boto3.resource("s3")
    object = s3.Object(bucket_name, key)
    return object.get()["Body"].read().decode("utf-8")

def delete_objects(username, bucket_id, objects):
    """
    Deletes an object from an s3 bucket.

    Returns: List of deleted objects
    """
    db_bucket = Bucket.get(hash_key=username, range_key="BUCKET_{}".format(bucket_id))
    bucket = boto3.resource("s3").Bucket(db_bucket.bucket_name)

    delete_dict = {"Objects": [{"Key": name} for name in objects]}

    response = bucket.delete_objects(Delete=delete_dict)

    if response["ResponseMetadata"]["HTTPStatusCode"] == 200:
        deleted_items = []
        for item in response["Deleted"]:
            deleted_items.append(item["Key"])
        return deleted_items
    else:
        deleted_items = []
        return deleted_items

"""
DB Functions related to S3
"""

def _create_db_bucket(username, id, bucket_name):
    new_bucket = Bucket(
        username=username,
        sort="BUCKET_{}".format(id),
        bucket_name=bucket_name,
        created_date=datetime.datetime.utcnow(),
        public_id=id,
        full=False,
    )
    new_bucket.save()
    return new_bucket

"""
Helper functions
"""

def _generate_bucket_id(username):
    full_buckets = Bucket.query(
        hash_key=username, range_key_condition=Bucket.sort.startswith("BUCKET_")
    )
    new_id = 0

    for buck in full_buckets:
        if buck.public_id > new_id:
            new_id = buck.public_id + 1

    return new_id

def _bucket_full(bucket_name):
    bucket = boto3.resource("s3").Bucket(bucket_name)
    size = sum([object.size for object in bucket.objects.all()])

    if size > 4990000000000:
        return True
    else:
        return False


from flask import Flask,request,render_template,redirect, url_for
from flask import jsonify
import requests
from cassandra.cluster import Cluster
from collections import OrderedDict

app = Flask(__name__)

KEYSPACE = "twitterkeyspace"

@app.route('/', methods=['GET'])
def home():
	if request.method == 'GET':
		return render_template('lab.html',flag=0)

# Lab Query 1
@app.route('/labquery1', methods=['GET','POST'])
def labquery1():
	if request.method == 'GET':
		return render_template('lab.html',flag=0)
	elif request.method == 'POST':
		date = request.form['date1']
		# print (date)

		# Connecting cassandra session
		cluster = Cluster(['127.0.0.1'])
		session = cluster.connect()
		session.default_timeout = 60
		session.set_keyspace(KEYSPACE)

		query = "SELECT author_id FROM tweet_table1 WHERE date = '{}'".format(date)
		rows = session.execute(query)
		session.execute("DROP TABLE IF EXISTS tweet_table_author_frequency")
		session.execute("""
        CREATE TABLE tweet_table_author_frequency (
            author_id text,
            frequency counter,
            PRIMARY KEY (author_id)
        )
        """)
		for row in rows:
			if row.author_id:
				query = "UPDATE tweet_table_author_frequency SET frequency = frequency + 1 WHERE author_id = '{}'".format(row.author_id)
				session.execute(query)
		rows = session.execute("SELECT * FROM tweet_table_author_frequency")
		pop = {}
		for row in rows:
			pop.update({row.author_id : row.frequency})
		res = sorted(pop,key=pop.get,reverse=True)
		result = []
		k = 0
		for r in res:
			temp = []
			temp.append(date)
			temp.append(r)
			temp.append(pop[r])
			result.append(temp)
			k += 1
			print (temp)
		return render_template('lab.html',result=result,count=k,flag=1)

# Lab Query 2
@app.route('/labquery2', methods=['GET','POST'])
def labquery2():
	if request.method == 'GET':
		return render_template('lab.html',flag=0)
	elif request.method == 'POST':
		date = request.form['date2']
		# print (date)

		# Connecting cassandra session
		cluster = Cluster(['127.0.0.1'])
		session = cluster.connect()
		session.default_timeout = 60
		session.set_keyspace(KEYSPACE)

		query = "SELECT hashtag, location FROM tweet_table2 WHERE date = '{}'".format(date)
		rows = session.execute(query)
		session.execute("DROP TABLE IF EXISTS tweet_table_hashtag_location")
		session.execute("""
        CREATE TABLE tweet_table_hashtag_location (
            hashtag text,
            location text,
            frequency counter,
            PRIMARY KEY ((hashtag,location))
        )
        """)
		for row in rows:
			if row.hashtag and row.location:
				print (row.hashtag,row.location)
				query = "UPDATE tweet_table_hashtag_location SET frequency = frequency + 1 WHERE hashtag = '{}'".format(row.hashtag) + "AND location = '{}'".format(row.location)
				session.execute(query)
		rows = session.execute("SELECT * FROM tweet_table_hashtag_location")
		pop = {}
		for row in rows:
			pop.update({(row.hashtag,row.location) : row.frequency})
		res = sorted(pop,key=pop.get,reverse=True)
		result = []
		k = 0
		for r in res:
			temp = []
			temp.append(date)
			temp.append(r[0])
			temp.append(r[1])
			temp.append(pop[r])
			result.append(temp)
			print (temp)
			k += 1
		return render_template('lab.html',result=result,count=k,flag=2)
	else:
		return render_template('lab.html',flag=0)

if __name__ == '__main__':
	app.run(host='127.0.0.1',port=5000,debug=True)


import logging
import math
import threading
import time

import torch

#import support.kernels as kernel_factory
from ...support.kernels import factory
from ...core import default
from ...core.model_tools.deformations.exponential import Exponential
from ...core.models.abstract_statistical_model import AbstractStatisticalModel
from ...core.models.model_functions import initialize_control_points, initialize_momenta
from ...core.observations.deformable_objects.deformable_multi_object import DeformableMultiObject
from ...in_out.array_readers_and_writers import *
from ...in_out.dataset_functions import create_template_metadata
#import support.utilities as utilities
from ...support.utilities import move_data, get_best_device
from .abstract_statistical_model import process_initial_data
import torch.nn.functional as F

logger = logging.getLogger(__name__)

# def _subject_attachment_and_regularity(arg):
#     """
#     Auxiliary function for multithreading (cannot be a class method).
#     """
#     from .abstract_statistical_model import process_initial_data
#     if process_initial_data is None:
#         raise RuntimeError('process_initial_data is not set !')
#
#     # Read arguments.
#     freeze_sparse_matrix = False
#     (deformable_objects, multi_object_attachment, objects_noise_variance,
#      freeze_template, freeze_control_points, freeze_momenta,
#      exponential, sobolev_kernel, use_sobolev_gradient, tensor_scalar_type, gpu_mode) = process_initial_data
#     (i, template, template_data, control_points, momenta, with_grad, momenta_t, sparse_matrix, alpha) = arg
#
#     # start = time.perf_counter()
#     device, device_id = get_best_device(gpu_mode=gpu_mode)
#     # device, device_id = ('cpu', -1)
#     if device_id >= 0:
#         torch.cuda.set_device(device_id)
#
#     # convert np.ndarrays to torch tensors. This is faster than transferring torch tensors to process.
#     template_data = {key: move_data(value, device=device, dtype=tensor_scalar_type,
#                                               requires_grad=with_grad and not freeze_template)
#                      for key, value in template_data.items()}
#     template_points = {key: move_data(value, device=device, dtype=tensor_scalar_type,
#                                                 requires_grad=with_grad and not freeze_template)
#                        for key, value in template.get_points().items()}
#     control_points = move_data(control_points, device=device, dtype=tensor_scalar_type,
#                                          requires_grad=with_grad and not freeze_control_points)
#     momenta = move_data(momenta, device=device, dtype=tensor_scalar_type,
#                                   requires_grad=with_grad and not freeze_momenta)
#
#     assert torch.device(
#         device) == control_points.device == momenta.device, 'control_points and momenta tensors must be on the same device. ' \
#                                                             'device=' + device + \
#                                                             ', control_points.device=' + str(control_points.device) + \
#                                                             ', momenta.device=' + str(momenta.device)
#
#     attachment, regularity = DeterministicAtlasHypertemplate._deform_and_compute_attachment_and_regularity(
#         exponential, template_points, control_points, momenta,
#         template, template_data, multi_object_attachment,
#         deformable_objects[i], objects_noise_variance, alpha,
#         device)
#
#     res = DeterministicAtlasHypertemplate._compute_gradients(
#         attachment, regularity,
#         freeze_template, momenta_t,
#         freeze_control_points, control_points,
#         freeze_momenta, momenta, freeze_sparse_matrix, sparse_matrix,
#         with_grad)
#     # elapsed = time.perf_counter() - start
#     # logger.info('pid=' + str(os.getpid()) + ', ' + torch.multiprocessing.current_process().name +
#     #       ', device=' + device + ', elapsed=' + str(elapsed))
#     return i, res

class DeterministicAtlasWithModule(AbstractStatisticalModel):
    """
    Deterministic atlas object class.

    """

    ####################################################################################################################
    ### Constructor:
    ####################################################################################################################

    def __init__(self, template_specifications, number_of_subjects,

                 dimension=default.dimension,
                 tensor_scalar_type=default.tensor_scalar_type,
                 tensor_integer_type=default.tensor_integer_type,
                 dense_mode=default.dense_mode,
                 number_of_processes=default.number_of_processes,

                 deformation_kernel_type=default.deformation_kernel_type,
                 deformation_kernel_width=default.deformation_kernel_width,
                 deformation_kernel_device=default.deformation_kernel_device,

                 shoot_kernel_type=default.shoot_kernel_type,
                 number_of_time_points=default.number_of_time_points,
                 use_rk2_for_shoot=default.use_rk2_for_shoot, use_rk2_for_flow=default.use_rk2_for_flow,

                 freeze_template=default.freeze_template,
                 use_sobolev_gradient=default.use_sobolev_gradient,
                 smoothing_kernel_width=default.smoothing_kernel_width,

                 initial_control_points=default.initial_control_points,
                 freeze_control_points=default.freeze_control_points,
                 initial_cp_spacing=default.initial_cp_spacing,

                 initial_momenta=default.initial_momenta,
                 freeze_momenta=default.freeze_momenta,

                 gpu_mode=default.gpu_mode,
                 process_per_gpu=default.process_per_gpu,

                 **kwargs):

        AbstractStatisticalModel.__init__(self, name='DeterministicAtlas', number_of_processes=number_of_processes,
                                          gpu_mode=gpu_mode)

        # Global-like attributes.
        self.dimension = dimension
        self.tensor_scalar_type = tensor_scalar_type
        self.tensor_integer_type = tensor_integer_type
        self.dense_mode = dense_mode

        # Declare model structure.
        self.fixed_effects['template_data'] = None
        self.fixed_effects['hypertemplate_data'] = None
        self.fixed_effects['control_points'] = None
        self.fixed_effects['momenta'] = None
        self.fixed_effects['momenta_t'] = None
        self.fixed_effects['module_intensities'] = None
        self.fixed_effects['module_positions'] = None
        self.fixed_effects['module_variances'] = None

        self.freeze_template = freeze_template
        self.freeze_control_points = freeze_control_points
        self.freeze_momenta = freeze_momenta
        self.freeze_sparse_matrix = False

        self.alpha = 1

        # Deformation.
        self.exponential = Exponential(
            dense_mode=dense_mode,
            kernel=factory(deformation_kernel_type,
                                          gpu_mode=gpu_mode,
                                          kernel_width=deformation_kernel_width),
            shoot_kernel_type=shoot_kernel_type,
            number_of_time_points=number_of_time_points,
            use_rk2_for_shoot=use_rk2_for_shoot, use_rk2_for_flow=use_rk2_for_flow)

        self.exponential_t = Exponential(
            dense_mode=dense_mode,
            kernel=factory(deformation_kernel_type,
                                          gpu_mode=gpu_mode,
                                          kernel_width=deformation_kernel_width),
            shoot_kernel_type=shoot_kernel_type,
            number_of_time_points=number_of_time_points,
            use_rk2_for_shoot=use_rk2_for_shoot, use_rk2_for_flow=use_rk2_for_flow)

        # Template.
        (object_list, self.objects_name, self.objects_name_extension,
         self.objects_noise_variance, self.multi_object_attachment) = create_template_metadata(template_specifications,
                                                                                self.dimension)

        self.template = DeformableMultiObject(object_list)
        self.hypertemplate = DeformableMultiObject(object_list)
        # self.template.update()

        self.number_of_objects = len(self.template.object_list)

        self.use_sobolev_gradient = use_sobolev_gradient
        self.smoothing_kernel_width = smoothing_kernel_width
        if self.use_sobolev_gradient:
            self.sobolev_kernel = factory(deformation_kernel_type,
                                                         gpu_mode=gpu_mode,
                                                         kernel_width=smoothing_kernel_width)

        # Template data.
        self.fixed_effects['template_data'] = self.template.get_data()
        self.fixed_effects['hypertemplate_data'] = self.hypertemplate.get_data()

        # Control points.
        self.fixed_effects['control_points'] = initialize_control_points(
            initial_control_points, self.template, initial_cp_spacing, deformation_kernel_width,
            self.dimension, self.dense_mode)
        self.number_of_control_points = len(self.fixed_effects['control_points'])

        # Momenta.
        self.fixed_effects['momenta'] = initialize_momenta(
            initial_momenta, self.number_of_control_points, self.dimension, number_of_subjects)
        self.fixed_effects['momenta'] = 0.0001*np.ones(self.fixed_effects['momenta'].shape)
        self.number_of_subjects = number_of_subjects

        self.fixed_effects['momenta_t'] = initialize_momenta(
            None, self.number_of_control_points, self.dimension, 1)

        # initial_cp = initialize_control_points(None, self.template, 40, None, self.dimension, False)
        # self.nb_modules = initial_cp.shape[0]
        # self.fixed_effects['module_positions'] = np.array([initial_cp,]*self.number_of_subjects)
        #self.fixed_effects['module_positions'] = np.array([[np.array(self.template.get_points()['image_points'].shape[:-1])/2]*initial_cp.shape[0]]*self.number_of_subjects)
        for i in range(object_list[0].bounding_box.shape[0]):
            object_list[0].bounding_box[i,0] = object_list[0].bounding_box[i,1]/2 - 10/2
            object_list[0].bounding_box[i, 1] = object_list[0].bounding_box[i, 1] / 2 + 10 / 2
        t = DeformableMultiObject(object_list)
        initial_cp = initialize_control_points(None, t, 3, None, self.dimension, False)
        self.nb_modules = initial_cp.shape[0]
        self.fixed_effects['module_positions'] = np.array([initial_cp, ] * self.number_of_subjects)

        # k = 0
        # j = 0
        # add = 1
        # while k < self.nb_modules - 1:
        #     self.fixed_effects['module_positions'][:,k,j] += add
        #     self.fixed_effects['module_positions'][:, k+1, j] -= add
        #     k += 2
        #     if j == self.fixed_effects['module_positions'].shape[2] - 1:
        #         j = 0
        #         add += 1
        #     else:
        #         j+=1

        self.fixed_effects['module_intensities'] = 0*np.ones([self.number_of_subjects, self.nb_modules])
        self.fixed_effects['module_variances'] = 5*np.ones([self.number_of_subjects, self.nb_modules, self.dimension])
        self.process_per_gpu = process_per_gpu

        self.regu_var_m = 5
        self.regu_var_m_ortho = 10

    def initialize_noise_variance(self, dataset, device='cpu'):
        if np.min(self.objects_noise_variance) < 0:
            hypertemplate_data, hypertemplate_points, template_data, template_points, control_points, momenta, momenta_t \
                = self._fixed_effects_to_torch_tensors(False, device=device)
            targets = dataset.deformable_objects
            targets = [target[0] for target in targets]

            residuals_torch = []
            self.exponential.set_initial_template_points(template_points)
            self.exponential.set_initial_control_points(control_points)
            for i, target in enumerate(targets):
                self.exponential.set_initial_momenta(momenta[i])
                self.exponential.update()
                deformed_points = self.exponential.get_template_points()
                deformed_data = self.template.get_deformed_data(deformed_points, template_data)
                residuals_torch.append(self.multi_object_attachment.compute_distances(
                    deformed_data, self.template, target))

            residuals = np.zeros((self.number_of_objects,))
            for i in range(len(residuals_torch)):
                residuals += residuals_torch[i].detach().cpu().numpy()

            # Initialize the noise variance hyper-parameter as a 1/100th of the initial residual.
            for k, obj in enumerate(self.objects_name):
                if self.objects_noise_variance[k] < 0:
                    nv = 0.01 * residuals[k] / float(self.number_of_subjects)
                    self.objects_noise_variance[k] = nv
                    logger.info('>> Automatically chosen noise std: %.4f [ %s ]' % (math.sqrt(nv), obj))

    ####################################################################################################################
    ### Encapsulation methods:
    ####################################################################################################################

    # Template data ----------------------------------------------------------------------------------------------------
    def get_template_data(self):
        return self.fixed_effects['template_data']

    def set_template_data(self, td):
        self.fixed_effects['template_data'] = td
        self.template.set_data(td)

    def get_hypertemplate_data(self):
        return self.fixed_effects['hypertemplate_data']

    # Control points ---------------------------------------------------------------------------------------------------
    def get_control_points(self):
        return self.fixed_effects['control_points']

    def set_control_points(self, cp):
        self.fixed_effects['control_points'] = cp
        # self.number_of_control_points = len(cp)

    # Momenta ----------------------------------------------------------------------------------------------------------
    def get_momenta(self):
        return self.fixed_effects['momenta']

    def set_momenta(self, mom):
        self.fixed_effects['momenta'] = mom

    def get_momenta_t(self):
        return self.fixed_effects['momenta_t']

    def set_momenta_t(self, mom):
        self.fixed_effects['momenta_t'] = mom

    def set_module_intensities(self, w):
        self.fixed_effects['module_intensities'] = w

    def get_module_intensities(self):
        return self.fixed_effects['module_intensities']

    def set_module_positions(self, c):
        self.fixed_effects['module_positions'] = c

    def get_module_positions(self):
        return self.fixed_effects['module_positions']

    def set_module_variances(self, sigma):
        self.fixed_effects['module_variances'] = sigma

    def get_module_variances(self):
        return self.fixed_effects['module_variances']

    # Full fixed effects -----------------------------------------------------------------------------------------------
    def get_fixed_effects(self):
        out = {}
        if not self.freeze_template:
            out['momenta_t'] = self.fixed_effects['momenta_t']
        if not self.freeze_control_points:
            out['control_points'] = self.fixed_effects['control_points']
        if not self.freeze_momenta:
            out['momenta'] = self.fixed_effects['momenta']
        if not self.freeze_sparse_matrix:
            out['module_positions'] = self.fixed_effects['module_positions']
            out['module_intensities'] = self.fixed_effects['module_intensities']
            out['module_variances'] = self.fixed_effects['module_variances']
        return out

    def set_fixed_effects(self, fixed_effects):
        if not self.freeze_template:
            device, _ = get_best_device(self.gpu_mode)
            hypertemplate_data, hypertemplate_points, template_data, template_points, control_points, momenta, momenta_t, module_intensities, module_positions, module_variances \
                = self._fixed_effects_to_torch_tensors(False, device=device)
            self.exponential_t.set_initial_template_points(hypertemplate_points)
            self.exponential_t.set_initial_control_points(control_points)
            self.exponential_t.set_initial_momenta(momenta_t[0])
            self.exponential_t.move_data_to_(device=device)
            self.exponential_t.update()
            template_points = self.exponential_t.get_template_points()
            template_data = self.hypertemplate.get_deformed_data(template_points, hypertemplate_data)
            template_data = {key: value.detach().cpu().numpy() for key, value in template_data.items()}
            self.set_momenta_t(fixed_effects['momenta_t'])
            self.set_template_data(template_data)
        if not self.freeze_control_points:
            self.set_control_points(fixed_effects['control_points'])
        if not self.freeze_momenta:
            self.set_momenta(fixed_effects['momenta'])
        if not self.freeze_sparse_matrix:
            self.set_module_positions(fixed_effects['module_positions'])
            self.set_module_variances((fixed_effects['module_variances']))
            self.set_module_intensities(fixed_effects['module_intensities'])

    ####################################################################################################################
    ### Public methods:
    ####################################################################################################################

    def setup_multiprocess_pool(self, dataset):
        self._setup_multiprocess_pool(initargs=([target[0] for target in dataset.deformable_objects],
                                                self.multi_object_attachment,
                                                self.objects_noise_variance,
                                                self.freeze_template, self.freeze_control_points, self.freeze_momenta,
                                                self.exponential, self.sobolev_kernel, self.use_sobolev_gradient,
                                                self.tensor_scalar_type, self.gpu_mode))

    # Compute the functional. Numpy input/outputs.
    def compute_log_likelihood(self, dataset, population_RER, individual_RER, mode='complete', with_grad=False):
        """
        Compute the log-likelihood of the dataset, given parameters fixed_effects and random effects realizations
        population_RER and indRER.

        :param fixed_effects: Dictionary of fixed effects.
        :param population_RER: Dictionary of population random effects realizations.
        :param individual_RER: Dictionary of individual random effects realizations.
        :param mode: Indicates which log_likelihood should be computed, between 'complete', 'model', and 'class2'.
        :param with_grad: Flag that indicates wether the gradient should be returned as well.
        :return:
        """

        if self.number_of_processes > 1:
            targets = [target[0] for target in dataset.deformable_objects]
            (deformable_objects, multi_object_attachment, objects_noise_variance,
             freeze_template, freeze_control_points, freeze_momenta,
            exponential, sobolev_kernel, use_sobolev_gradient, tensor_scalar_type, gpu_mode) = process_initial_data
            device, device_id = get_best_device(gpu_mode=gpu_mode)

            self.exponential_t.set_initial_template_points(self.hypertemplate.get_points())
            self.exponential_t.set_initial_control_points(self.fixed_effects['control_points'])
            self.exponential_t.set_initial_momenta(self.fixed_effects['momenta_t'])
            self.exponential_t.move_data_to_(device=device)
            self.exponential_t.update()
            template_points = self.exponential_t.get_template_points()
            template_data = self.hypertemplate.get_deformed_data(template_points, self.fixed_effects['hypertemplate_data'])

            args = [(i, self.template,
                     template_data,
                     self.fixed_effects['control_points'],
                     self.fixed_effects['momenta'][i],
                     with_grad) for i in range(len(targets))]

            start = time.perf_counter()

            results = self.pool.map(_subject_attachment_and_regularity, args, chunksize=1)  # TODO: optimized chunk size
            # results = self.pool.imap_unordered(_subject_attachment_and_regularity, args, chunksize=1)
            # results = self.pool.imap(_subject_attachment_and_regularity, args, chunksize=int(len(args)/self.number_of_processes))
            logger.debug('time taken for deformations : ' + str(time.perf_counter() - start))

            # Sum and return.
            if with_grad:
                attachment = 0.0
                regularity = 0.0

                gradient = {}
                if not self.freeze_template:
                    gradient['momenta_t'] = np.zeros(self.fixed_effects['momenta_t'].shape)
                if not self.freeze_control_points:
                    gradient['control_points'] = np.zeros(self.fixed_effects['control_points'].shape)
                if not self.freeze_momenta:
                    gradient['momenta'] = np.zeros(self.fixed_effects['momenta'].shape)

                for result in results:
                    i, (attachment_i, regularity_i, gradient_i) = result
                    attachment += attachment_i
                    regularity += regularity_i
                    for key, value in gradient_i.items():
                        if key == 'momenta':
                            gradient[key][i] = value
                        else:
                            gradient[key] += value
                return attachment, regularity, gradient
            else:
                attachment = 0.0
                regularity = 0.0
                for result in results:
                    i, (attachment_i, regularity_i) = result
                    attachment += attachment_i
                    regularity += regularity_i
                return attachment, regularity

        else:
            device, device_id = get_best_device(gpu_mode=self.gpu_mode)
            hypertemplate_data, hypertemplate_points, template_data, template_points, control_points, momenta, \
            momenta_t, module_intensities, module_positions, module_variances = self._fixed_effects_to_torch_tensors(with_grad,device=device)
            sparse_matrix = self.construct_sparse_matrix(template_points['image_points'], module_positions, module_variances, module_intensities)
            return self._compute_attachment_and_regularity(dataset, hypertemplate_data, hypertemplate_points, control_points,
                                                           momenta, momenta_t, sparse_matrix, module_intensities, module_positions, module_variances, with_grad, device=device)

    ####################################################################################################################
    ### Private methods:
    ####################################################################################################################

    @staticmethod
    def _deform_and_compute_attachment_and_regularity(exponential, template_points, control_points, momenta, module_positions, module_variances, module_intensities, sparse_matrix,
                                                      template, template_data,
                                                      multi_object_attachment, deformable_objects,
                                                      objects_noise_variance, regu_var_m, regu_var_m_ortho,
                                                      device='cpu'):

        # Deform.
        exponential.set_initial_template_points(template_points)
        exponential.set_initial_control_points(control_points)
        exponential.set_initial_momenta(momenta)
        exponential.move_data_to_(device=device)
        exponential.update()

        # Compute attachment and regularity.
        deformed_points = exponential.get_template_points()
        deformed_data = template.get_deformed_data(deformed_points, template_data)
        deformed_data['image_intensities'] += sparse_matrix
        attachment = -multi_object_attachment.compute_weighted_distance(deformed_data, template, deformable_objects,
                                                                        objects_noise_variance)
        regularity = - exponential.get_norm_squared()

        # x = template_points['image_points'].cpu().detach()
        # final_cp = exponential.control_points_t[-1].cpu().detach()
        # regu = torch.zeros(x.shape[:-1], dtype=torch.float64)
        # for k in range(final_cp.shape[0]):
        #     m = momenta[k].clone().cpu().detach()
        #     if m.numpy().any():
        #         m /= torch.norm(m)
        #         e = torch.randn(m.size(), dtype=torch.float64)
        #         e -= e.dot(m) * m
        #         e /= torch.norm(e)
        #         if m.size()[0] == 2:
        #             y = x - final_cp[k]
        #             regu += torch.exp(-torch.mm(y.view(-1,2), m.view(2,1)) ** 2 / (2 * 1) - torch.mm(y.view(-1,2), e.view(2,1)) ** 2 / (2 * 10)).view(y.shape[:-1])
        #         else:
        #             e2 = torch.cross(m, e)
        #             y = x - final_cp[k]
        #             regu += torch.exp(
        #                 -torch.dot(y, m) ** 2 / (2 * 1) - torch.dot(y, e) ** 2 / (2 * 10) - torch.dot(y, e2) ** 2 / (2 * 10))
        #
        # dim = template_data['image_intensities'].shape
        # regu2 = torch.zeros(dim).double()
        # for k in range(module_positions.shape[0]):
        #     x_norm = torch.mul(x ** 2, 1 / module_variances[k] ** 2).sum(-1).view(-1, 1)
        #     y_norm = torch.mul(module_positions[k] ** 2, 1 / module_variances[k] ** 2).sum().view(-1, 1)
        #     points_divided = torch.mul(x, 1 / module_variances[k] ** 2)
        #     dist = (x_norm + y_norm - 2.0 * torch.mul(points_divided, module_positions[k]).sum(-1).view(-1,
        #                                                                                1)).reshape(
        #             dim)
        #     regu2 += torch.exp(-dist)*torch.abs(module_intensities[k].detach())
        #
        # regularity -= 200.*torch.sum(torch.mul(torch.tensor(regu, dtype=torch.float64),regu2))

        final_cp = exponential.control_points_t[-1]
        final_momenta = exponential.momenta_t[-1]
        for k in range(final_cp.shape[0]):
            if momenta[k].detach().numpy().any():
                m = final_momenta[k] / torch.norm(momenta[k])
                e = torch.randn(m.shape[0], dtype=torch.float64)
                e = e - torch.dot(e, m) * m
                e = e / torch.norm(e)
                for l in range(module_positions.shape[0]):
                    if m.size()[0] == 2:
                        y = module_positions[l] - final_cp[k]
                        regularity += - 10000*torch.sum(torch.exp(- torch.mm(y.view(-1,2), m.view(2,1)) ** 2 / (2 * regu_var_m) - torch.mm(y.view(-1,2), e.view(2,1)) ** 2 / (2 * regu_var_m_ortho)).view(y.shape[:-1]))
                    else:
                        e2 = torch.cross(m, e)
                        y = module_positions[l] - final_cp[k]
                        regularity += - torch.sum(torch.exp(
                            -torch.dot(y, m) ** 2 / (2 * regu_var_m) - torch.dot(y, e) ** 2 / (2 * regu_var_m_ortho) - torch.dot(y, e2) ** 2 / (2 * regu_var_m_ortho)).view(y.shape[:-1]))

        x_norm = (module_positions ** 2).sum(1).view(-1, 1)
        dist = x_norm + x_norm.view(1, -1) - 2.0 * torch.mm(module_positions, torch.transpose(module_positions, 0, 1))

        # dist = torch.zeros([16,16], dtype=torch.float64)
        # for k in range(16):
        #     dist[k,:] = torch.norm(module_positions - module_positions[k], dim=1)**2/400
        regularity -= -torch.sum(torch.exp(-dist))

        assert torch.device(
            device) == attachment.device == regularity.device, 'attachment and regularity tensors must be on the same device. ' \
                                                               'device=' + device + \
                                                               ', attachment.device=' + str(attachment.device) + \
                                                               ', regularity.device=' + str(regularity.device)

        return attachment, regularity

    @staticmethod
    def _compute_gradients(attachment, regularity,
                           freeze_template, momenta_t,
                           freeze_control_points, control_points,
                           freeze_momenta, momenta, freeze_sparse_matrix, module_intensities, module_positions, module_variances,
                           with_grad=False):
        if with_grad:
            total_for_subject = attachment + regularity
            total_for_subject.backward()

            gradient = {}
            if not freeze_template:
                assert momenta_t.grad is not None, 'Gradients have not been computed'
                gradient['momenta_t'] = momenta_t.grad.detach().cpu().numpy()
            if not freeze_control_points:
                assert control_points.grad is not None, 'Gradients have not been computed'
                gradient['control_points'] = control_points.grad.detach().cpu().numpy()
            if not freeze_momenta:
                assert momenta.grad is not None, 'Gradients have not been computed'
                gradient['momenta'] = momenta.grad.detach().cpu().numpy()
            if not freeze_sparse_matrix:
                gradient['module_intensities'] = module_intensities.grad.detach().cpu().numpy()
                gradient['module_positions'] = module_positions.grad.detach().cpu().numpy()
                gradient['module_variances'] = module_variances.grad.detach().cpu().numpy()

            res = attachment.detach().cpu().numpy(), regularity.detach().cpu().numpy(), gradient

        else:
            res = attachment.detach().cpu().numpy(), regularity.detach().cpu().numpy()

        return res

    def _compute_attachment_and_regularity(self, dataset, hypertemplate_data, hypertemplate_points, control_points, momenta, momenta_t, sparse_matrix, module_intensities, module_positions, module_variances,
                                           with_grad=False, device='cpu'):
        """
        Core part of the ComputeLogLikelihood methods. Torch input, numpy output.
        Single-thread version.
        """

        # Initialize.
        targets = [target[0] for target in dataset.deformable_objects]
        attachment = 0.
        regularity = 0.

        self.exponential_t.set_initial_template_points(hypertemplate_points)
        self.exponential_t.set_initial_control_points(control_points)
        self.exponential_t.set_initial_momenta(momenta_t[0])
        self.exponential_t.move_data_to_(device=device)
        self.exponential_t.update()
        template_points = self.exponential_t.get_template_points()
        template_data = self.hypertemplate.get_deformed_data(template_points, hypertemplate_data)
        self.set_template_data({key: value.detach().cpu().numpy() for key, value in template_data.items()})

        regularity -= self.exponential_t.get_norm_squared()

        # loop for every deformable object
        # deform and update attachment and regularity
        for i, target in enumerate(targets):
            new_attachment, new_regularity = DeterministicAtlasWithModule._deform_and_compute_attachment_and_regularity(
                self.exponential, template_points, control_points, momenta[i], module_positions[i], module_variances[i], module_intensities[i], sparse_matrix[i],
                self.template, template_data, self.multi_object_attachment,
                target, self.objects_noise_variance, self.regu_var_m, self.regu_var_m_ortho,
                device=device)

            attachment += new_attachment
            regularity += new_regularity

        # Compute gradient.
        return self._compute_gradients(attachment, regularity,
                                       self.freeze_template, momenta_t,
                                       self.freeze_control_points, control_points,
                                       self.freeze_momenta, momenta, self.freeze_sparse_matrix, module_intensities, module_positions, module_variances,
                                       with_grad)

    ####################################################################################################################
    ### Private utility methods:
    ####################################################################################################################

    def construct_sparse_matrix(self, points, module_centers, module_variances, module_intensities):
        dim = (self.number_of_subjects,) + self.fixed_effects['template_data']['image_intensities'].shape
        sparse_matrix = torch.zeros(dim).double()
        for i in range(dim[0]):
            for k in range(self.nb_modules):
                x_norm = torch.mul(points ** 2, 1/module_variances[i,k]**2).sum(-1).view(-1, 1)
                y_norm = torch.mul(module_centers[i,k] ** 2, 1/module_variances[i,k]**2).sum().view(-1, 1)
                points_divided = torch.mul(points, 1/module_variances[i,k]**2)
                dist = (x_norm + y_norm - 2.0 * torch.mul(points_divided, module_centers[i,k]).sum(-1).view(-1,1)).reshape(dim[1:])
                sparse_matrix[i] += torch.exp(-dist)*module_intensities[i,k]
                # sparse_matrix[i] += 70/81*(1-dist)**2*(dist<1).double()*module_intensities[i,k]

                # x_norm = (points ** 2)
                # y_norm = (module_centers[i,k] ** 2)
                # dist = (x_norm + y_norm - 2.0 * torch.mul(points, module_centers[i, k])).view(-1, 2)
                # rect = ((dist[:,0] < module_variances[i,k,0]) * (dist[:,1] < module_variances[i,k,1])).double()
                # f = torch.exp(-dist.sum(1)/10)
                # conv2 = torch.nn.Conv2d(1, 1, kernel_size=1, stride=1, padding=1, bias=False)
                # conv2.weights = rect.reshape([1, 1, 100, 100])
                # sparse_matrix[i] += conv2(f.float().reshape([1,1,100,100]))[0,0,1:-1,1:-1].double()*module_intensities[i,k]

                #sparse_matrix[i] += module_intensities[i,k] * (1/(1+torch.exp(-2*2*(points[:,:,0] - module_centers[i,k,0] + module_variances[i,k,0]/2))) - 1/(1+torch.exp(-2*2*(points[:,:,0] - module_centers[i,k,0] - module_variances[i,k,0]/2)))) * (1/(1+ torch.exp(-2*2*(points[:,:,1] - module_centers[i,k,1] + module_variances[i,k,1]/2))) - 1/(1+torch.exp(-2*2*(points[:,:,1] - module_centers[i,k,1] - module_variances[i,k,1]/2))))

                # x_norm = torch.mul(points ** 2, 1/module_variances[i,k]**2).sum(2).view(-1, 1)
                # y_norm = torch.mul(module_centers[i,k] ** 2, 1/module_variances[i,k]**2).sum().view(-1, 1)
                # points_divided = torch.mul(points, 1/module_variances[i,k]**2)
                # dist = (x_norm + y_norm - 2.0 * torch.mul(points_divided, module_centers[i,k]).sum(2).view(-1,1)).reshape(dim[1:])
                # sparse_matrix[i] += module_intensities[i,k] * (1/(1+torch.exp(-2*100*(dist+0.1))) - 1/(1+torch.exp(-2*100*(dist - 1))))
        return sparse_matrix

    def _fixed_effects_to_torch_tensors(self, with_grad, device='cpu'):
        """
        Convert the fixed_effects into torch tensors.
        """
        # Template data.
        template_data = self.fixed_effects['template_data']
        template_data = {key: move_data(value, device=device, dtype=self.tensor_scalar_type,
                                                  requires_grad=False)
                         for key, value in template_data.items()}

        # Template points.
        template_points = self.template.get_points()
        template_points = {key: move_data(value, device=device, dtype=self.tensor_scalar_type,
                                                    requires_grad=False)
                           for key, value in template_points.items()}

        hypertemplate_data = self.fixed_effects['hypertemplate_data']
        hypertemplate_data = {key: move_data(value, device=device, dtype=self.tensor_scalar_type,
                                                  requires_grad=False)
                         for key, value in hypertemplate_data.items()}

        # Template points.
        hypertemplate_points = self.hypertemplate.get_points()
        hypertemplate_points = {key: move_data(value, device=device, dtype=self.tensor_scalar_type,
                                                    requires_grad=False)
                           for key, value in hypertemplate_points.items()}

        momenta_t = self.fixed_effects['momenta_t']
        momenta_t = move_data(momenta_t, device=device, dtype=self.tensor_scalar_type,
                                                 requires_grad=(not self.freeze_template and with_grad))

        # Control points.
        if self.dense_mode:
            assert (('landmark_points' in self.template.get_points().keys()) and
                    ('image_points' not in self.template.get_points().keys())), \
                'In dense mode, only landmark objects are allowed. One at least is needed.'
            control_points = template_points['landmark_points']
        else:
            control_points = self.fixed_effects['control_points']
            control_points = move_data(control_points, device=device, dtype=self.tensor_scalar_type,
                                                 requires_grad=(not self.freeze_control_points and with_grad))
            # control_points = Variable(torch.from_numpy(control_points).type(self.tensor_scalar_type),
            #                           requires_grad=(not self.freeze_control_points and with_grad))
        # Momenta.
        momenta = self.fixed_effects['momenta']
        momenta = move_data(momenta, device=device, dtype=self.tensor_scalar_type,
                                      requires_grad=(not self.freeze_momenta and with_grad))

        module_intensities = self.fixed_effects['module_intensities']
        module_intensities = move_data(module_intensities, device=device, dtype=self.tensor_scalar_type,
                                      requires_grad=(not self.freeze_sparse_matrix and with_grad))

        module_positions = self.fixed_effects['module_positions']
        module_positions = move_data(module_positions, device=device, dtype=self.tensor_scalar_type,
                                             requires_grad=(not self.freeze_sparse_matrix and with_grad))

        module_variances = self.fixed_effects['module_variances']
        module_variances = move_data(module_variances, device=device, dtype=self.tensor_scalar_type,
                                             requires_grad=(not self.freeze_sparse_matrix and with_grad))

        return hypertemplate_data, hypertemplate_points, template_data, template_points, control_points, momenta, momenta_t, module_intensities, module_positions, module_variances

    ####################################################################################################################
    ### Writing methods:
    ####################################################################################################################

    def write(self, dataset, population_RER, individual_RER, output_dir, write_residuals=True):

        # Write the model predictions, and compute the residuals at the same time.
        residuals = self._write_model_predictions(dataset, individual_RER, output_dir,
                                                  compute_residuals=write_residuals)

        # Write residuals.
        if write_residuals:
            residuals_list = [[residuals_i_k.data.cpu().numpy() for residuals_i_k in residuals_i]
                              for residuals_i in residuals]
            write_2D_list(residuals_list, output_dir, self.name + "__EstimatedParameters__Residuals.txt")

        # Write the model parameters.
        self._write_model_parameters(output_dir)

    def _write_model_predictions(self, dataset, individual_RER, output_dir, compute_residuals=True):
        device, _ = get_best_device(self.gpu_mode)

        # Initialize.
        hypertemplate_data, hypertemplate_points, template_data, template_points, control_points, momenta, momenta_t, \
        module_intensities, module_positions, module_variances = self._fixed_effects_to_torch_tensors(False, device=device)

        sparse_matrix = self.construct_sparse_matrix(template_points['image_points'], module_positions, module_variances, module_intensities)

        # Deform, write reconstructions and compute residuals.
        self.exponential.set_initial_template_points(template_points)
        self.exponential.set_initial_control_points(control_points)

        residuals = []  # List of torch 1D tensors. Individuals, objects.
        for i, subject_id in enumerate(dataset.subject_ids):
            self.exponential.set_initial_momenta(momenta[i])
            self.exponential.update()

            # Writing the whole flow.
            names = []
            for k, object_name in enumerate(self.objects_name):
                name = self.name + '__flow__' + object_name + '__subject_' + subject_id
                names.append(name)
            #self.exponential.write_flow(names, self.objects_name_extension, self.template, template_data, output_dir)

            deformed_points = self.exponential.get_template_points()
            deformed_data = self.template.get_deformed_data(deformed_points, template_data)
            deformed_data['image_intensities'] += sparse_matrix[i]
            m = torch.max(deformed_data['image_intensities'])
            for k in range(module_positions.shape[1]):
                for j in range(module_positions[i,k].shape[0]):
                    module_positions[i,k,j] = min(99, module_positions[i,k,j])
                    module_positions[i, k, j] = max(0, module_positions[i, k, j])
                deformed_data['image_intensities'][tuple(module_positions[i,k].int())] = 2* m

            if compute_residuals:
                residuals.append(self.multi_object_attachment.compute_distances(
                    deformed_data, self.template, dataset.deformable_objects[i][0]))

            names = []
            for k, (object_name, object_extension) \
                    in enumerate(zip(self.objects_name, self.objects_name_extension)):
                name = self.name + '__Reconstruction__' + object_name + '__subject_' + subject_id + object_extension
                names.append(name)

            self.template.write(output_dir, names,
                                {key: value.detach().cpu().numpy() for key, value in deformed_data.items()})

            deformed_data['image_intensities'] = sparse_matrix[i]
            names = []
            for k, (object_name, object_extension) \
                    in enumerate(zip(self.objects_name, self.objects_name_extension)):
                name = self.name + '__Reconstruction__' + object_name + '__subject_' + subject_id + '_sparsematrix' + object_extension
                names.append(name)

            self.template.write(output_dir, names,
                                {key: value.detach().cpu().numpy() for key, value in deformed_data.items()})

        return residuals

    def _write_model_parameters(self, output_dir):

        # Template.
        device, _ = get_best_device(self.gpu_mode)
        template_names = []
        hypertemplate_data, hypertemplate_points, template_data, template_points, control_points, momenta, momenta_t, module_intensities, module_positions, module_variances \
            = self._fixed_effects_to_torch_tensors(False, device=device)
        self.exponential_t.set_initial_template_points(hypertemplate_points)
        self.exponential_t.set_initial_control_points(control_points)
        self.exponential_t.set_initial_momenta(momenta_t[0])
        self.exponential_t.move_data_to_(device=device)
        self.exponential_t.update()
        template_points = self.exponential_t.get_template_points()
        template_data = self.hypertemplate.get_deformed_data(template_points, hypertemplate_data)
        self.set_template_data({key: value.detach().cpu().numpy() for key, value in template_data.items()})
        for i in range(len(self.objects_name)):
            aux = self.name + "__EstimatedParameters__Template_" + self.objects_name[i] + self.objects_name_extension[i]
            template_names.append(aux)
        self.template.write(output_dir, template_names)

        # Control points.
        write_2D_array(self.get_control_points(), output_dir, self.name + "__EstimatedParameters__ControlPoints.txt")

        # Momenta.
        write_3D_array(self.get_momenta(), output_dir, self.name + "__EstimatedParameters__Momenta.txt")
        write_2D_array(self.get_momenta_t()[0], output_dir, self.name + "__EstimatedParameters__Momenta_t.txt")
        write_3D_array(self.get_module_positions(), output_dir, self.name + "__EstimatedParameters__ModulePositions.txt")
        write_3D_array(self.get_module_intensities(), output_dir, self.name + "__EstimatedParameters__ModuleIntensities.txt")
        write_3D_array(self.get_module_variances(), output_dir, self.name + "__EstimatedParameters__ModuleVariances.txt")