#!/usr/bin/env python3
"""
Creates wallpaper with personal top albums from Last.fm.
"""

import argparse
import configparser
import datetime
import glob
import hashlib
import logging
import math
import os
import re
import shutil
import sys
from functools import lru_cache

import numpy
import pylast
import requests
from PIL import (
    Image,
    ImageChops,
    ImageEnhance,
    ImageFilter,
    ImageMath,
    ImageOps,
    PngImagePlugin,
)
from urllib3.util.retry import Retry

DEFAULT_CONFIG_FILE_PATH = os.path.expanduser("~/.config/lastfm_wallpaper.ini")
DEFAULT_SERVER_NAME = "default"
DEFAULT_ALBUM_COVER_DIR = os.path.expanduser("~/.cache/lastfm_wallpaper")
DEFAULT_MAX_COVER_COUNT = 12
DEFAULT_SPACE = 50

DEFAULT_WIDTH = 1920
DEFAULT_HEIGHT = 1080
DEFAULT_SIZE = "{}x{}".format(DEFAULT_WIDTH, DEFAULT_HEIGHT)

DEFAULT_MAX_TAGS_TO_MATCH = 2

MISSING_CONFIG_ERROR = """\
You have to have your own unique two values for API_KEY and API_SECRET Obtain
yours from https://www.last.fm/api/account/create and save them in following
format in file "{}".

    [{}]
    api_key = xxxxxxxxxxxxxxx
    api_secret = xxxxxxxxxxxxxxx
    user = login_name
"""

SEARCH_PATHS_EXAMPLE = os.path.pathsep.join(
    (
        os.path.join("~", "Music", "{artist} - {album}", "cover.*"),
        os.path.join("~", "Music", "*", "{artist} - {album}", "cover.*"),
    )
)

logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
logger = logging.getLogger(__name__)

@lru_cache(maxsize=1)
def session():
    session = requests.Session()
    retry = Retry(
        total=3,
        read=3,
        connect=3,
        backoff_factor=1,
        status_forcelist=(500, 502, 503, 504),
        method_whitelist=Retry.DEFAULT_METHOD_WHITELIST.union(("POST",)),
    )
    adapter = requests.adapters.HTTPAdapter(max_retries=retry)
    session.mount("http://", adapter)
    session.mount("https://", adapter)
    return session

class DownloadCoverError(RuntimeError):
    pass

class TupleArgument:
    def __init__(self, argument, separator=","):
        self.x, self.y = map(int, argument.split(separator))

class SizeArgument(TupleArgument):
    def __init__(self, size):
        super().__init__(size, "x")

class LayoutArgument:
    def __init__(self, value):
        if value is None:
            self.positions = None
        else:
            self.positions, self.rows, self.columns = parse_layout(value)

class Layout:
    def __init__(
        self,
        positions,
        background,
        rows,
        columns,
        width,
        height,
        space,
        angle_range,
    ):
        self.positions = positions
        self.background = background
        self.rows = rows
        self.columns = columns
        self.space = space

        self.angle_range = angle_range
        self.angles = []

        self.extent = extent = min(
            (height - space * rows) // rows,
            (width - space * columns) // columns,
        )
        self.padding_x = (width - extent * columns) // (columns + 1)
        self.padding_y = (height - extent * rows) // (rows + 1)

    def paste(self, cell, img, offset=(0, 0)):
        img = rotate(img, self.angle(cell))

        row, column = self.position(cell, self.columns)
        x = (
            column * self.extent
            + self.extent // 2
            + (column + 1) * self.padding_x
        )
        y = row * self.extent + self.extent // 2 + (row + 1) * self.padding_y
        paste(img, x + offset[0], y + offset[1], self.background)

    def position(self, cell, columns):
        if self.positions:
            return self.positions[cell]

        return divmod(cell, columns)

    def angle(self, cell):
        if self.angle_range.x >= self.angle_range.y:
            return 0

        count_to_add = cell + 1 - len(self.angles)
        self.angles.extend(
            numpy.random.randint(self.angle_range.x, self.angle_range.y)
            for _ in range(count_to_add)
        )

        return self.angles[cell]

class CoverLoader:
    def __init__(self, album_dir):
        self.album_dir = album_dir
        self.cache = {}

    def cover(self, index, extent):
        path = self.cover_path(index)
        img = self.cache.get(path)
        if not img:
            img = Image.open(path, "r")
            img = img.convert("RGBA")
            self.cache[path] = img
        return img.resize((extent, extent), resample=Image.BICUBIC)

    def cover_path(self, index):
        return image_path(self.album_dir, index + 1)

def parse_config(config_path, server):
    config = configparser.ConfigParser()
    config.read(config_path)
    try:
        return config[server]
    except KeyError:
        raise SystemExit(MISSING_CONFIG_ERROR.format(config_path, server))

def parse_layout(value):
    try:
        positions = [
            [int(x) for x in position.split(",")]
            for position in value.split(" ")
        ]
        min_x = min(positions, key=lambda p: p[0])[0]
        max_x = max(positions, key=lambda p: p[0])[0]
        min_y = min(positions, key=lambda p: p[1])[1]
        max_y = max(positions, key=lambda p: p[1])[1]
        rows = max_x - min_x + 1
        columns = max_y - min_y + 1
        positions = [(x - min_x, y - min_y) for x, y in positions]
        return positions, rows, columns
    except Exception as e:
        logger.exception("Failed to parse positions argument: %s", e)
        raise

def image_info(**args):
    info = PngImagePlugin.PngInfo()
    for k, v in args.items():
        info.add_itxt(k, v)
    return info

def image_path(image_dir, base_name):
    return os.path.join(image_dir, "{}.png".format(base_name))

def fix_name(name):
    """
    Fixes album/artist name.
    """
    name = str(name).strip()
    if name.endswith(")"):
        name = name.rsplit("(", 1)[0].rstrip()
    name = name.replace(".", " ")
    return name

def get_cover_image_from_lastfm(album):
    return album.get_cover_image(pylast.SIZE_MEGA)

def get_cover_image_from_deezer(album):
    url = "https://api.deezer.com/search/autocomplete"
    try:
        resp = session().get(
            url,
            params={
                "q": f"{fix_name(album.artist)} - {fix_name(album.title)}",
            },
        )
        resp.raise_for_status()
        return resp.json()["albums"]["data"][0]["cover_xl"]
    except Exception as e:
        logger.warning("Failed to fetch cover from %r: %s", url, e)

def cover_for_album(album):
    try:
        cover_url = get_cover_image_from_deezer(
            album
        ) or get_cover_image_from_lastfm(album)
        if not cover_url:
            raise DownloadCoverError("Cover URL not available")
    except Exception as e:
        raise DownloadCoverError("Failed to get cover URL: {}".format(e))

    return cover_url

def download_raw(url):
    try:
        r = session().get(url, stream=True)
    except requests.exceptions.RequestException as e:
        raise DownloadCoverError("Failed to download cover: {}".format(e))

    if r.status_code != 200:
        raise DownloadCoverError("Failed to download cover: {}".format(r.text))

    r.raw.decode_content = True
    return r.raw

def cache_path_for_album(album, cache_dir):
    album_id = "{} //// {}".format(album.artist, album.title)
    cache_base_name = hashlib.sha256(album_id.encode("utf-8")).hexdigest()
    return os.path.join(cache_dir, cache_base_name) + ".png"

def save_cover(album, raw_or_path, path):
    img = Image.open(raw_or_path)

    # Workaround for opening 16bit greyscale images.
    # See: https://github.com/python-pillow/Pillow/issues/2574
    if img.mode == "I" and numpy.array(img).max() > 255:
        logger.warning("Fixing 16bit image")
        img = ImageMath.eval("img/256", {"img": img})

    img = img.convert("RGBA")
    info = image_info(artist=album.artist.name, album=album.title)
    img.save(path, pnginfo=info)

def download_cover(album, cache_path):
    cover_url = cover_for_album(album)
    raw = download_raw(cover_url)
    save_cover(album, raw, cache_path)

def lastfm_user(api_key, api_secret, user):
    network = pylast.LastFMNetwork(
        api_key=api_key, api_secret=api_secret, username=user
    )

    return network.get_user(user)

def to_pattern(text):
    case_insensitive = "".join(
        "[{}{}]".format(c.lower(), c.upper()) if c.isalpha() else c
        for c in str(text)
    )
    return "*{}*".format(case_insensitive)

def find_album(album, search):
    for pattern in search:
        artist = to_pattern(album.artist)
        album_title = to_pattern(album.title)
        path = pattern.format(artist=artist, album=album_title)
        paths = glob.iglob(path)
        try:
            return next(paths)
        except StopIteration:
            pass

def get_cover_for_album(album, path, cache_path, search):
    if os.path.isfile(cache_path):
        logger.info('Album "%s": Using cached cover', album)
    else:
        found = find_album(album, search)
        if found:
            logger.info('Album "%s": Getting cover from "%s"', album, found)
            save_cover(album, found, cache_path)
        else:
            try:
                logger.info('Album "%s": Downloading cover', album)
                download_cover(album, cache_path)
            except DownloadCoverError as e:
                logger.warning(e)
                return False

    shutil.copyfile(cache_path, path)
    return True

def download_covers(
    user, album_dir, from_date, to_date, max_count, search, tag_re, max_tags
):
    cache_dir = os.path.join(album_dir, ".cache")
    os.makedirs(cache_dir, exist_ok=True)

    top_items = user.get_weekly_album_charts(
        from_date=from_date.strftime("%s"), to_date=to_date.strftime("%s")
    )

    count = 0

    for top_item in top_items:
        album = top_item.item

        if tag_re:
            tags = album.artist.get_top_tags()
            if not any(
                tag_re.match(tag.item.get_name()) for tag in tags[:max_tags]
            ):
                tag_names = ", ".join(tag.item.get_name() for tag in tags)
                logger.info("No matching tags: %s (%s)", album, tag_names)
                continue

        path = image_path(album_dir, count + 1)
        cache_path = cache_path_for_album(album, cache_dir)
        if get_cover_for_album(
            album, path=path, cache_path=cache_path, search=search
        ):
            count += 1
            if count == max_count:
                break

    return count

def parse_args():
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )

    parser.add_argument(
        "--info",
        action="store_true",
        help="print list of albums in the last wallpaper and exit",
    )

    parser.add_argument(
        "--config", default=DEFAULT_CONFIG_FILE_PATH, help="config file path"
    )
    parser.add_argument(
        "--server",
        default=DEFAULT_SERVER_NAME,
        help="server name (section in config file)",
    )
    parser.add_argument(
        "--dir",
        default=DEFAULT_ALBUM_COVER_DIR,
        help="directory to store album covers",
    )
    parser.add_argument(
        "--size",
        default=DEFAULT_SIZE,
        type=SizeArgument,
        help="wallpaper size",
    )
    parser.add_argument(
        "--count",
        default=DEFAULT_MAX_COVER_COUNT,
        type=int,
        help="maximum cover count",
    )
    parser.add_argument(
        "--rows", default=-1, type=int, help="number of rows; -1 to deduce"
    )
    parser.add_argument(
        "--columns",
        default=-1,
        type=int,
        help="number of columns; -1 to deduce",
    )
    parser.add_argument(
        "--space", default=DEFAULT_SPACE, type=int, help="space between items"
    )
    parser.add_argument(
        "--cached", action="store_true", help="use already downloaded covers"
    )
    parser.add_argument(
        "--search",
        default="",
        help='album search patterns (e.g. "{}")'.format(SEARCH_PATHS_EXAMPLE),
    )
    parser.add_argument(
        "--tags", default="", help="tag search pattern; regular expression"
    )
    parser.add_argument(
        "--max-tags",
        default=DEFAULT_MAX_TAGS_TO_MATCH,
        type=int,
        help="maximum number tags to search",
    )

    parser.add_argument(
        "--angle-range",
        default="0,0",
        type=TupleArgument,
        help="random cover rotation",
    )

    parser.add_argument(
        "--shadow-offset",
        default="1,1",
        type=TupleArgument,
        help="shadow offset",
    )
    parser.add_argument(
        "--shadow-blur", default=4, type=int, help="shadow blur"
    )
    parser.add_argument("--shadow-color", default="black", help="shadow color")

    parser.add_argument(
        "--border-color",
        default="black",
        help='border color; "auto" to auto-detect based on cover',
    )
    parser.add_argument(
        "--border-size", default=10, type=int, help="border size"
    )

    parser.add_argument(
        "--days", default=7, type=int, help="number of days to consider"
    )
    parser.add_argument(
        "--hours",
        default=0,
        type=int,
        help="number of additional hours to consider",
    )
    parser.add_argument(
        "--days-ago",
        default=0,
        type=int,
        help="consider end date X days ago instead of today",
    )

    parser.add_argument(
        "--base",
        default="random",
        help=(
            "base image file"
            '; "random" to pick one of the covers'
            "; <number> to pick the cover at given position"
        ),
    )
    parser.add_argument(
        "--base-blur", default=3, type=int, help="base image blur"
    )
    parser.add_argument(
        "--base-brightness",
        default=80,
        type=int,
        help="base image brightness percentage",
    )
    parser.add_argument(
        "--base-noise",
        default=10,
        type=int,
        help="base image noise percentage",
    )
    parser.add_argument(
        "--base-color",
        default=50,
        type=int,
        help="base image color percentage",
    )

    parser.add_argument(
        "--cover-brightness",
        default=100,
        type=int,
        help="cover image brightness percentage",
    )
    parser.add_argument(
        "--cover-noise",
        default=5,
        type=int,
        help="cover image noise percentage",
    )
    parser.add_argument(
        "--cover-color",
        default=100,
        type=int,
        help="cover image color percentage",
    )
    parser.add_argument(
        "--cover-glow", default=40, type=int, help="cover glow amount"
    )

    parser.add_argument(
        "--random-seed",
        default=-1,
        type=int,
        help=(
            "seed number to initialize random number generator; "
            "random if negative"
        ),
    )

    parser.add_argument(
        "--layout",
        default=None,
        type=LayoutArgument,
        help=(
            "cover positions and layout"
            '; space separated list of "row,column" values'
        ),
    )

    args = parser.parse_args()
    config = parse_config(args.config, args.server)
    config = {
        key.lower().replace("-", "_"): value for key, value in config.items()
    }
    parser.set_defaults(**config)

    return parser.parse_args()

def background_image(path, width, height, blur_radius):
    background = Image.open(path, "r")
    background = background.convert("RGBA")
    extent = math.floor(max(width, height))
    x = (extent - width) // 2
    y = (extent - height) // 2
    background = background.resize((extent, extent), resample=Image.BICUBIC)
    background = background.crop((x, y, x + width, y + height))
    return blur(background, blur_radius)

def add_noise(img, noise_percentage):
    if noise_percentage <= 0:
        return img

    noise = numpy.random.randint(
        0, 255, size=(img.height, img.width, 3), dtype=numpy.uint8
    )
    noise_image = Image.fromarray(noise, mode="RGB").convert("RGBA")
    return ImageChops.blend(img, noise_image, noise_percentage / 100)

def brighter(img, brightness_percentage):
    return ImageEnhance.Brightness(img).enhance(brightness_percentage / 100)

def colorize(img, colorize_percentage):
    return ImageEnhance.Color(img).enhance(colorize_percentage / 100)

def rotate(img, angle):
    if (angle % 360) == 0:
        return img

    # Expand first to have smoother edges.
    img = ImageOps.expand(img, 4, fill=0)
    return img.rotate(
        angle, resample=Image.BICUBIC, expand=True, fillcolor=(0, 0, 0, 0)
    )

def blur(img, radius):
    return img.filter(ImageFilter.GaussianBlur(radius=radius))

def glow(img, amount):
    extent1 = img.width
    img = colorize(img, 200)
    img = brighter(img, 200)
    img = blur(img, amount)

    extent2 = int(img.width * 0.6)
    mask = img.convert("L")
    mask = mask.resize((extent2, extent2))
    d = (extent1 - extent2) // 2
    mask = ImageOps.expand(mask, d, "black")
    mask = mask.resize((img.width, extent1))
    mask = blur(mask, amount)
    img.putalpha(mask)

    return img

def auto_border_color(img):
    img = img.resize((1, 1), resample=Image.BILINEAR)
    img = colorize(img, 200)
    img = brighter(img, 30)
    return img.getpixel((0, 0))

def paste(img, x, y, background):
    dest = [x - img.width // 2, y - img.height // 2]
    src = [0, 0]
    for j in range(2):
        if dest[j] < 0:
            src[j] = -dest[j]
            dest[j] = 0
    background.alpha_composite(img, dest=tuple(dest), source=tuple(src))

def init_random_seed(seed):
    if seed < 0:
        return

    numpy.random.seed(seed)

def print_info(album_dir):
    path = image_path(album_dir, "wallpaper")
    if not os.path.isfile(path):
        raise SystemExit("No wallpaper found")

    img = Image.open(path, "r")

    info = img.info
    albums = info.pop("albums", None)

    info["image"] = path
    info["resolution"] = "{}x{}".format(img.width, img.height)

    for k, v in sorted(info.items()):
        print("# {}: {}".format(k, v))

    if albums:
        print("# albums:\n\n{}".format(albums))

    sys.exit(0)

def main():
    args = parse_args()

    album_dir = args.dir

    if args.info:
        print_info(album_dir)

    init_random_seed(args.random_seed)

    width, height = args.size.x, args.size.y
    max_count = args.count

    rows = args.rows
    columns = args.columns

    if args.layout:
        positions = args.layout.positions
        if len(positions) < max_count:
            raise SystemExit(
                "Expected %s positions but %s specified",
                max_count,
                len(positions),
            )
        if rows < 0:
            rows = args.layout.rows
        if columns < 0:
            columns = args.layout.columns
    else:
        positions = None

    user = lastfm_user(
        api_key=args.api_key, api_secret=args.api_secret, user=args.user
    )

    image_info_dict = {}

    search = [
        os.path.expanduser(path) for path in args.search.split(os.path.pathsep)
    ]

    if args.cached:
        count = max_count
    else:
        to_date = datetime.datetime.now() - datetime.timedelta(
            days=args.days_ago
        )
        from_date = (
            to_date
            - datetime.timedelta(days=args.days)
            - datetime.timedelta(hours=args.hours)
        )
        image_info_dict["dates"] = "{}..{}".format(
            from_date.date(), to_date.date()
        )
        tag_re = re.compile(args.tags, re.IGNORECASE) if args.tags else None
        count = download_covers(
            user=user,
            album_dir=album_dir,
            from_date=from_date,
            to_date=to_date,
            max_count=max_count,
            search=search,
            tag_re=tag_re,
            max_tags=args.max_tags,
        )

    if count <= 0:
        raise SystemExit("No albums in given time range")

    if rows < 0 and columns < 0:
        x = width / height
        _, _, rows = min(
            [count % rows, abs(x - (count / rows) / rows), rows]
            for rows in range(1, count + 1)
        )
    if rows < 0:
        rows = math.ceil(count / columns)
    if columns < 0:
        columns = math.ceil(count / rows)

    scale = max(1, int(width / DEFAULT_WIDTH))
    base_blur = args.base_blur * scale

    loader = CoverLoader(album_dir)

    if args.base == "random":
        i = numpy.random.randint(count)
        path = loader.cover_path(i)
    else:
        try:
            i = int(args.base) - 1
            if i < 0:
                i = count + i + 1
            path = loader.cover_path(i)
        except TypeError:
            path = args.base
    background = background_image(path, width, height, blur_radius=base_blur)

    background = add_noise(background, args.base_noise)
    background = brighter(background, args.base_brightness)
    background = colorize(background, args.base_color)

    if not columns:
        columns = math.ceil(count / rows)
    space = args.space * scale
    layout = Layout(
        positions,
        background,
        rows,
        columns,
        width,
        height,
        space,
        args.angle_range,
    )
    extent = layout.extent

    shadow_size = int(extent * 1.2)
    shadow = Image.new("RGBA", (shadow_size, shadow_size))
    shadow_pos = (shadow_size - extent) // 2
    shadow.paste(
        args.shadow_color,
        (shadow_pos, shadow_pos, extent + shadow_pos, extent + shadow_pos),
    )
    shadow_blur = args.shadow_blur * scale
    shadow = blur(shadow, shadow_blur)
    shadow_offset = (
        args.shadow_offset.x * scale,
        args.shadow_offset.y * scale,
    )

    border = args.border_size * scale

    for i in range(count):
        layout.paste(i, shadow, shadow_offset)

    if args.cover_glow > 0:
        for i in reversed(range(count)):
            extent1 = int(extent * (100 + args.cover_glow) / 100)
            img = loader.cover(i, extent1)
            img = glow(img, extent // 10)
            layout.paste(i, img)

    albums = []

    for i in reversed(range(count)):
        extent1 = extent - 2 * border
        img = loader.cover(i, extent1)
        artist = img.info.get("artist")
        album = img.info.get("album")
        if artist and album:
            albums.insert(0, "{} - {}".format(artist, album))

        border_color = args.border_color
        if border_color == "auto":
            border_color = auto_border_color(img)

        img = ImageOps.expand(img, border, border_color)

        img = add_noise(img, args.cover_noise)
        img = brighter(img, args.cover_brightness)
        img = colorize(img, args.cover_color)

        layout.paste(i, img)

    albums = "\n".join(albums)
    if args.cached:
        logger.info("Using cached covers for albums:\n%s", albums)

    image_info_dict["albums"] = albums
    image_info_dict["url"] = user.get_url()

    path = image_path(album_dir, "wallpaper")
    info = image_info(**image_info_dict)
    background.save(path, pnginfo=info)
    print("Wallpaper saved: {}".format(path))

if __name__ == "__main__":
    main()


import sys
import os
import time
from src.model import *
from src.util import *
import numpy as np
import pandas as pd
from torch.utils.data import DataLoader
from sklearn.model_selection import train_test_split
import torch.nn as nn
import torch.optim as optim
from torch.autograd import Variable

label_path = "/public/malware_dataset/kaggle_microsoft_9_10000/"
train_data_path = label_path + "bytes/"  # Training data
train_label_path = label_path + "kaggle_microsoft_trainlabels.csv"  # Training label
#valid_label_path = label_path + "example-valid-label.csv"  # Validation Label

#name
exp_name = "malconv-classification"

# Parameter
# os.environ["CUDA_VISIBLE_DEVICES"] = "2" # single-GPU

use_gpu = True  #
use_cpu = 32  # Number of cores to use for data loader
display_step = 5  # Std output update rate during training 和 保存训练结果步长
test_step = 50  # Test per n step
learning_rate = 0.0001  #
max_step = 1000  # Number of steps to train
batch_size = 768  #
first_n_byte = (
    100000  # First N bytes of a PE file as the input of MalConv (defualt: 2 million)
)
window_size = 512  # Kernel size & stride for Malconv (defualt : 500)

# output path
log_dir = "/log/"
pred_dir = "/pred/"
checkpoint_dir = "/checkpoint/"
log_file_path = log_dir + exp_name + ".log"
chkpt_acc_path = checkpoint_dir + exp_name + "1000.pt"
pred_path = pred_dir + exp_name + ".pred"

df = pd.read_csv(train_label_path)
train, valid, train_label, valid_label = train_test_split(
    df["Id"],
    df["Class"],
    test_size=0.2,
    stratify=df["Class"],
    random_state=100,
)
"""
# Dataset preparation
class ExeDataset(Dataset):
	def __init__(self, fp_list, data_path, label_list, first_n_byte=2000000):
		self.fp_list = fp_list
		self.data_path = data_path
		self.label_list = label_list
		self.first_n_byte = first_n_byte

	def __len__(self):
		return len(self.fp_list)

	def __getitem__(self, idx):
		try:
			with open(self.data_path + self.fp_list[idx],'rb') as f:
				tmp = [i+1 for i in f.read()[:self.first_n_byte]] # index 0 will be special padding index 每个值加一
				tmp = tmp+[0]*(self.first_n_byte-len(tmp))
		except:
			with open(self.data_path + self.fp_list[idx].lower(),'rb') as f:
				tmp = [i+1 for i in f.read()[:self.first_n_byte]]
				tmp = tmp+[0]*(self.first_n_byte-len(tmp))

		return np.array(tmp), np.array([self.label_list[idx]])
"""

trainset = pd.DataFrame({"id": train, "labels": train_label})
validset = pd.DataFrame({"id": valid, "labels": valid_label})
trainloader = DataLoader(
    ExeDataset(
        list(trainset["id"]), train_data_path, list(trainset["labels"]), first_n_byte
    ),
    batch_size=batch_size,
    shuffle=False,
    num_workers=use_cpu,
    pin_memory=True,
)
validloader = DataLoader(
    ExeDataset(
        list(validset["id"]), train_data_path, list(validset["labels"]), first_n_byte
    ),
    batch_size=batch_size,
    shuffle=False,
    num_workers=use_cpu,
    pin_memory=True,
)
USE_CUDA = torch.cuda.is_available()
device = torch.device("cuda:1" if USE_CUDA else "cpu")

malconv = MalConv(input_length=first_n_byte, window_size=window_size)
malconv = nn.DataParallel(malconv, device_ids=[1,2,3]) # multi-GPU

#malconv = MalConvBase(8, 4096, 128, 32)
bce_loss = nn.BCEWithLogitsLoss()
ce_loss = nn.CrossEntropyLoss()
adam_optim = optim.Adam([{"params": malconv.parameters()}], lr=learning_rate)
sigmoid = nn.Sigmoid()

if use_gpu:
    malconv = malconv.to(device)
    bce_loss = bce_loss.to(device)
    sigmoid = sigmoid.to(device)

step_msg = "step-{}-loss-{:.6f}-acc-{:.4f}-time-{:.2f}s"
valid_msg = "step-{}-tr_loss-{:.6f}-tr_acc-{:.4f}-val_loss-{:.6f}-val_acc-{:.4f}"
log_msg = "{}, {:.6f}, {:.4f}, {:.6f}, {:.4f}, {:.2f}"
history = {}
history["tr_loss"] = []
history["tr_acc"] = []
train_acc = []  # 保存训练结果

valid_best_acc = 0.0
total_step = 0
step_cost_time = 0
valid_idx = list(validset["id"])

while total_step < max_step:

    # Training
    for step, batch_data in enumerate(trainloader):
        start = time.time()

        adam_optim.zero_grad()

        cur_batch_size = batch_data[0].size(0)

        exe_input = batch_data[0].to(device) if use_gpu else batch_data[0]
        exe_input = Variable(exe_input.long(), requires_grad=False)

        label = batch_data[1].to(device) if use_gpu else batch_data[1]
        label = Variable(label, requires_grad=False)
        label = label.squeeze() - 1

        pred = malconv(exe_input)
        loss = ce_loss(pred, label)
        loss.backward()
        adam_optim.step()

        _, predicted = torch.max(pred.data, 1)
        train_Macc = (label.cpu().data.numpy().astype(int) == (predicted.cpu().data.numpy()).astype(int)).sum().item()
        train_Macc = train_Macc / cur_batch_size

        if (step + 1) % display_step == 0:
            print("train：{}".format(train_Macc))

        total_step += 1
        # Interupt for validation
        if total_step % test_step == 0:
            break

    for step, val_batch_data in enumerate(validloader):

        cur_batch_size = val_batch_data[0].size(0)

        exe_input = val_batch_data[0].to(device) if use_gpu else val_batch_data[0]
        exe_input = Variable(exe_input.long(), requires_grad=False)

        label = val_batch_data[1].to(device) if use_gpu else val_batch_data[1]
        label = Variable(label, requires_grad=False)
        label = label.squeeze() - 1

        pred = malconv(exe_input)
        loss = ce_loss(pred, label)
        # loss.backward()
        # adam_optim.step()

        _, predicted = torch.max(pred.data, 1)
        val_Macc = (label.cpu().data.numpy().astype(int) == (predicted.cpu().data.numpy()).astype(int)).sum().item()
        val_Macc = val_Macc / cur_batch_size

        if (step + 1) % display_step == 0:
            print("test：{}".format(val_Macc))


import math
import time
from sklearn import neighbors
import os
import os.path
import pickle
import face_recognition
import cv2
import numpy as np
import random, string
import datetime
import sqlite3

# Create a db
conn = sqlite3.connect('employee.db')
c = conn.cursor()
# Create a table in the db
c.execute('''CREATE TABLE IF NOT EXISTS attendancee (Date text, emp_name text, attendance text, arrival_time text)''')

sql1 = 'DELETE FROM attendancee'
c.execute(sql1)

conn.commit()
conn.close()

def predict(img, knn_clf=None, model_path=None, distance_threshold=0.4):

    if knn_clf is None and model_path is None:
        raise Exception("Must supply knn classifier either thourgh knn_clf or model_path")

    # Load a trained KNN model (if one was passed in)
    if knn_clf is None:
        with open(model_path, 'rb') as f:
            knn_clf = pickle.load(f)

    # Load image file and find face locations
    #X_img = face_recognition.load_image_file(X_img_path)
    X_face_locations = face_recognition.face_locations(img)

    # If no faces are found in the image, return an empty result.
    if len(X_face_locations) == 0:
        return []

    # Find encodings for faces in the test iamge
    faces_encodings = face_recognition.face_encodings(img, known_face_locations=X_face_locations)

    # Use the KNN model to find the best matches for the test face
    closest_distances = knn_clf.kneighbors(faces_encodings, n_neighbors=1)
    are_matches = [closest_distances[0][i][0] <= distance_threshold for i in range(len(X_face_locations))]

    # Predict classes and remove classifications that aren't within the threshold
    return [(pred, loc) if rec else ("unknown", loc) for pred, loc, rec in zip(knn_clf.predict(faces_encodings), X_face_locations, are_matches)]

print("\n Looking for faces via webcam...")

video_capture = cv2.VideoCapture(0)
#time.sleep(2.0)
name_exists = 'unjhuhhb'  #random name string
while True:
    # Using the trained classifier, make predictions for unknown images
        ret, frame = video_capture.read()
        predictions = predict(frame , model_path="trained_model.clf")

        # Print results on the console
        for name, (top, right, bottom, left) in predictions:
            cv2.rectangle(frame,(left, top),(right, bottom),(0,255,0),3)

            # Draw a label with a name below the face
            cv2.rectangle(frame, (left, bottom - 35), (right, bottom), (0, 0, 255), cv2.FILLED)
            font = cv2.FONT_HERSHEY_DUPLEX
            cv2.putText(frame, name, (left + 6, bottom - 6), font, 1.0, (255, 255, 255), 1)

            #if match occurs save employee image and write in database
            if name != "unknown" and name != name_exists:

                face_image = frame[top:bottom, left:right]
                #path where employees' detected faces are stored
                path = " "
                cv2.imwrite(os.path.join(path , name+'.jpg'), face_image)
                print("- Found {} at ({}, {})".format(name, left, top))

                k = datetime.datetime.now()

                tim = str(k.hour)+':'+str(k.minute)+':'+str(k.second)
                dat = str(k.day)+'-'+str(k.month)+'-'+str(k.year)
                ff = 'Present'
                # Connect to the database
                conn = sqlite3.connect('employee.db')
                c = conn.cursor()
                # Insert a row of data
                c.execute("INSERT INTO attendancee VALUES (?,?,?,?)", (dat, name, ff, tim))
                conn.commit()
                conn.close()
                print('Attendance marked')

                name_exists = name

        cv2.imshow('Video', frame)

        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
# Release handle to the webcam
video_capture.release()
cv2.destroyAllWindows()


from .cassiemujoco import pd_in_t, state_out_t, CassieSim, CassieVis

from .trajectory import CassieTrajectory

from math import floor

import numpy as np
import os
import random

import pickle

def getAllTrajectories(speeds):
    trajectories = []

    for i, speed in enumerate(speeds):
        dirname = os.path.dirname(__file__)
        traj_path = os.path.join(dirname, "trajectory", "aslipTrajsTaskSpace", "walkCycle_{}.pkl".format(speed))
        trajectories.append(CassieIKTrajectory(traj_path))

    # print("Got all trajectories")
    return trajectories

class CassieIKTrajectory:
    def __init__(self, filepath):
        with open(filepath, "rb") as f:
            trajectory = pickle.load(f)

        self.qpos = np.copy(trajectory["qpos"])
        self.length = self.qpos.shape[0]
        self.qvel = np.copy(trajectory["qvel"])
        self.rpos = np.copy(trajectory["rpos"])
        self.rvel = np.copy(trajectory["rvel"])
        self.lpos = np.copy(trajectory["lpos"])
        self.lvel = np.copy(trajectory["lvel"])
        self.cpos = np.copy(trajectory["cpos"])
        self.cvel = np.copy(trajectory["cvel"])

# simrate used to be 60
class UnifiedCassieIKEnv:
    def __init__(self, traj="stepping", simrate=60, clock_based=False, state_est=True, training=True, debug=False):
        self.sim = CassieSim("./cassiemujoco/cassie.xml")
        self.vis = None

        # robot state estimation included here
        self.observation_space = np.zeros(46 + 18)

        # motor PD targets
        self.action_space      = np.zeros(10)

        self.speeds = np.array([x / 10 for x in range(0, 21)])
        self.trajectories = getAllTrajectories(self.speeds)
        self.num_speeds = len(self.trajectories)

        self.P = np.array([100,  100,  88,  96,  50])
        self.D = np.array([10.0, 10.0, 8.0, 9.6, 5.0])

        self.u = pd_in_t()

        self.simrate = simrate # simulate X mujoco steps with same pd target
                               # 60 brings simulation from 2000Hz to roughly 30Hz

        self.time    = 0 # number of time steps in current episode
        self.phase   = 0 # portion of the phase the robot is in
        self.counter = 0 # number of phase cycles completed in episode

        # NOTE: each trajectory in trajectories should have the same length
        self.speed = self.speeds[0]
        self.trajectory = self.trajectories[0]

        self.training = training

        # NOTE: a reference trajectory represents ONE phase cycle

        # should be floor(len(traj) / simrate) - 1
        # should be VERY cautious here because wrapping around trajectory
        # badly can cause assymetrical/bad gaits
        # self.phaselen = floor(self.trajectory.length / self.simrate) - 1
        self.phaselen = self.trajectory.length - 1

        # see include/cassiemujoco.h for meaning of these indices
        self.pos_idx = [7, 8, 9, 14, 20, 21, 22, 23, 28, 34]
        self.reward_pos_idx = [7, 8, 9, 14, 20, 21, 22, 23, 28, 34]
        self.vel_idx = [6, 7, 8, 12, 18, 19, 20, 21, 25, 31]

        # for enforcing the correct foot orientation
        self.global_initial_foot_orient = np.array([-0.24135469773826795, -0.24244324494623198, -0.6659363823866352, 0.6629463911006771])
        self.avg_lfoot_quat = np.zeros(4)
        self.avg_rfoot_quat = np.zeros(4)

        ## THIS IS JUST AN IDEA, instead of doing this we are having one ref trajectory per speed
            # params for changing the trajectory
            # self.speed = 2 # how fast (m/s) do we go
            # self.gait = [1,0,0,0]   # one-hot vector of gaits:
                                    # [1, 0, 0, 0] -> walking/running (left single stance, right single stance)
                                    # [0, 1, 0, 0] -> hopping (double stance, flight phase)
                                    # [0, 0, 1, 0] -> skipping (double stance, right single stance, flight phase, right single stance)
                                    # [0, 0, 0, 1] -> galloping (double stance, right single stance, flight, left single stance)

        # maybe make ref traj only send relevant idxs?
        ref_pos, ref_vel = self.get_ref_state(self.phase)
        self.prev_action = ref_pos[self.pos_idx]
        self.phase_add = 1

        # Output of Cassie's state estimation code
        self.cassie_state = state_out_t()

        # for print statements
        self.debug = debug

    def step_simulation(self, action):

        # maybe make ref traj only send relevant idxs?
        if(self.phase == self.phaselen - 1):
            ref_pos, ref_vel = self.get_ref_state(0)
        else:
            ref_pos, ref_vel = self.get_ref_state(self.phase + 1)

        target = action + ref_pos[self.pos_idx]

        ## DO NOT WANT RATE LIMITING
            # h = 0.0005
            # Tf = 1.0 / 300.0
            # alpha = h / (Tf + h)
            # real_action = (1-alpha)*self.prev_action + alpha*target

        # real_action = target

        # diff = real_action - self.prev_action
        # max_diff = np.ones(10)*0.1
        # for i in range(10):
        #     if diff[i] < -max_diff[i]:
        #         target[i] = self.prev_action[i] - max_diff[i]
        #     elif diff[i] > max_diff[i]:
        #         target[i] = self.prev_action[i] + max_diff[i]

        # self.prev_action = real_action
        # real_action = target

        self.u = pd_in_t()
        for i in range(5):
            # TODO: move setting gains out of the loop?
            # maybe write a wrapper for pd_in_t ?
            self.u.leftLeg.motorPd.pGain[i]  = self.P[i]
            self.u.rightLeg.motorPd.pGain[i] = self.P[i]

            self.u.leftLeg.motorPd.dGain[i]  = self.D[i]
            self.u.rightLeg.motorPd.dGain[i] = self.D[i]

            self.u.leftLeg.motorPd.torque[i]  = 0 # Feedforward torque
            self.u.rightLeg.motorPd.torque[i] = 0

            self.u.leftLeg.motorPd.pTarget[i]  = target[i]
            self.u.rightLeg.motorPd.pTarget[i] = target[i + 5]

            self.u.leftLeg.motorPd.dTarget[i]  = 0
            self.u.rightLeg.motorPd.dTarget[i] = 0

        self.cassie_state = self.sim.step_pd(self.u)

    def step(self, action):
        for i in range(self.simrate):
            self.step_simulation(action)
            # calculate running average of foot quaternion
            self.avg_lfoot_quat += self.sim.xquat("left-foot")
            self.avg_rfoot_quat += self.sim.xquat("right-foot")
        self.avg_lfoot_quat /= self.simrate
        self.avg_rfoot_quat /= self.simrate

        height = self.sim.qpos()[2]

        self.time  += 1
        self.phase += 1

        # if self.phase > self.phaselen:
        if self.phase >= self.phaselen:
            self.phase = 0
            self.counter += 1

        # Early termination
        done = not(height > 0.4 and height < 3.0)

        reward = self.compute_reward(action)

        # reset avg foot quaternion
        self.avg_lfoot_quat = np.zeros(4)
        self.avg_rfoot_quat = np.zeros(4)

        # update previous action
        self.prev_action = action

        # TODO: make 0.3 a variable/more transparent
        if reward < 0.3:
            done = True

        return self.get_full_state(), reward, done, {}

    def reset(self):
        random_speed_idx = random.randint(0, self.num_speeds-1)
        self.speed = self.speeds[random_speed_idx]
        # print("current speed: {}".format(self.speed))
        self.trajectory = self.trajectories[random_speed_idx] # switch the current trajectory
        self.phaselen = self.trajectory.length - 1
        self.phase = random.randint(0, self.phaselen)
        self.time = 0
        self.counter = 0

        qpos, qvel = self.get_ref_state(self.phase)
        # qpos[2] -= .1

        self.sim.set_qpos(qpos)
        # self.sim.set_qvel(qvel)
        self.sim.set_qvel(np.zeros(qvel.shape))

        # Need to reset u? Or better way to reset cassie_state than taking step
        self.cassie_state = self.sim.step_pd(self.u)

        # maybe make ref traj only send relevant idxs?
        ref_pos, ref_vel = self.get_ref_state(self.phase)
        self.prev_action = ref_pos[self.pos_idx]

        return self.get_full_state()

    # used for plotting against the reference trajectory
    def reset_for_test(self):
        random_speed_idx = random.randint(0, self.num_speeds)
        self.speed = 0
        self.trajectory = self.trajectories[0] # switch the current trajectory
        self.phaselen = self.trajectory.length - 1
        self.phase = 0
        self.time = 0
        self.counter = 0

        qpos, qvel = self.get_ref_state(self.phase)

        self.sim.set_qpos(qpos)
        self.sim.set_qvel(qvel)

        # maybe make ref traj only send relevant idxs?
        ref_pos, ref_vel = self.get_ref_state(self.phase)
        self.prev_action = ref_pos[self.pos_idx]

        # Need to reset u? Or better way to reset cassie_state than taking step
        self.cassie_state = self.sim.step_pd(self.u)

        return self.get_full_state()

    def update_selected_trajectory(self, new_speed):
        self.speed = new_speed
        self.trajectory = self.trajectories[(np.abs(self.speeds - self.speed)).argmin()]
        self.phaselen = self.trajectory.length - 1

    # NOTE: this reward is slightly different from the one in Xie et al
    # see notes for details
    def compute_reward(self, action):
        qpos = np.copy(self.sim.qpos())
        qvel = np.copy(self.sim.qvel())

        ref_pos, ref_vel = self.get_ref_state(self.phase)

        # TODO: should be variable; where do these come from?
        # TODO: see magnitude of state variables to gauge contribution to reward
        # weight = [0.15, 0.15, 0.1, 0.05, 0.05, 0.15, 0.15, 0.1, 0.05, 0.05]

        # weight = [0.05, 0.05, 0.25, 0.25, 0.05,
        #           0.05, 0.05, 0.25, 0.25, 0.05]

        #weight = [.1] * 10

        footpos_error     = 0
        com_vel_error     = 0
        action_penalty    = 0
        foot_orient_penalty = 0
        straight_diff = 0

        # enforce distance between feet and com
        ref_rfoot, ref_lfoot  = self.get_ref_footdist(self.phase + 1)

        # left foot
        lfoot = self.cassie_state.leftFoot.position[:]
        rfoot = self.cassie_state.rightFoot.position[:]
        for j in [0, 1, 2]:
            footpos_error += np.linalg.norm(lfoot[j] - ref_lfoot[j]) +  np.linalg.norm(rfoot[j] - ref_rfoot[j])

        if self.debug:
            print("ref_rfoot: {}  rfoot: {}".format(ref_rfoot, rfoot))
            print("ref_lfoot: {}  lfoot: {}".format(ref_lfoot, lfoot))
            print(footpos_error)

        # try to match com velocity
        ref_cvel = self.get_ref_com_vel(self.phase + 1)

        # center of mass vel: x, y, z
        cvel = self.cassie_state.pelvis.translationalVelocity
        for j in [0, 1, 2]:
            com_vel_error += np.linalg.norm(cvel[j] - ref_cvel[j])

        # # each joint pos, skipping feet
        # for i, j in enumerate(self.reward_pos_idx):
        #     target = ref_pos[j]
        #     actual = qpos[j]

        #     if j == 20 or j == 34:
        #         joint_error += 0
        #     else:
        #         joint_error += (target - actual) ** 2

        # action penalty
        action_penalty = np.linalg.norm(action - self.prev_action)

        # foot orientation penalty
        foot_orient_penalty = np.linalg.norm(self.avg_rfoot_quat - self.global_initial_foot_orient) + np.linalg.norm(self.avg_lfoot_quat - self.global_initial_foot_orient)

        # straight difference penalty
        straight_diff = np.abs(qpos[1])
        if straight_diff < 0.05:
            straight_diff = 0

        reward = 0.3 * np.exp(-footpos_error) +    \
                 0.3 * np.exp(-com_vel_error) +    \
                 0.1 * np.exp(-action_penalty) +     \
                 0.2 * np.exp(-foot_orient_penalty) + \
                 0.1 * np.exp(-straight_diff)

        if self.debug:
            print("reward: {10}\nfoot:\t{0:.2f}, % = {1:.2f}\ncom_vel:\t{2:.2f}, % = {3:.2f}\naction_penalty:\t{4:.2f}, % = {5:.2f}\nfoot_orient_penalty:\t{6:.2f}, % = {7:.2f}\nstraight_diff:\t{8:.2f}, % = {9:.2f}\n\n".format(
            0.3 * np.exp(-footpos_error),          0.3 * np.exp(-footpos_error) / reward * 100,
            0.3 * np.exp(-com_vel_error),          0.3 * np.exp(-com_vel_error) / reward * 100,
            0.1 * np.exp(-action_penalty),         0.1 * np.exp(-action_penalty) / reward * 100,
            0.2 * np.exp(-foot_orient_penalty),    0.2 * np.exp(-foot_orient_penalty) / reward * 100,
            0.1  * np.exp(-straight_diff),         0.1  * np.exp(-straight_diff) / reward * 100,
            reward
            )
            )
            print("actual speed: {}\tdesired_speed: {}".format(qvel[0], self.speed))
        return reward

    # get the corresponding state from the reference trajectory for the current phase
    def get_ref_state(self, phase=None):

        # print("phase: {}\t (phaselen = {})".format(phase, self.phaselen))

        if phase is None:
            phase = self.phase

        if phase > self.phaselen:
            phase = 0

        # print("looped phase: {}".format(phase))
        # print()

        # pos = np.copy(self.trajectory.qpos[phase * self.simrate])
        pos = np.copy(self.trajectory.qpos[phase])

        # this is just setting the x to where it "should" be given the number
        # of cycles
        #pos[0] += (self.trajectory.qpos[-1, 0] - self.trajectory.qpos[0, 0]) * self.counter
        pos[0] += (self.trajectory.qpos[-1, 0] - self.trajectory.qpos[0, 0]) * self.counter

        # ^ should only matter for COM error calculation,
        # gets dropped out of state variable for input reasons

        # setting lateral distance target to 0?
        # regardless of reference trajectory?
        pos[1] = 0

        # vel = np.copy(self.trajectory.qvel[phase * self.simrate])
        vel = np.copy(self.trajectory.qvel[phase])

        return pos, vel

    # get the corresponding state from the reference trajectory for the current phase
    def get_ref_footdist(self, phase=None):

        if phase is None:
            phase = self.phase

        if phase > self.phaselen:
            phase = 0

        rpos = np.copy(self.trajectory.rpos[phase])
        lpos = np.copy(self.trajectory.lpos[phase])

        return rpos, lpos

    def get_ref_com_vel(self, phase=None):

        if phase is None:
            phase = self.phase

        if phase > self.phaselen:
            phase = 0

        cvel = np.copy(self.trajectory.cvel[phase])

        return cvel

    def get_ref_ext_state(self, phase=None):

        if phase is None:
            phase = self.phase

        if phase > self.phaselen:
            phase = 0

        rpos = np.copy(self.trajectory.rpos[phase])
        rvel = np.copy(self.trajectory.rvel[phase])
        lpos = np.copy(self.trajectory.lpos[phase])
        lvel = np.copy(self.trajectory.lvel[phase])
        cpos = np.copy(self.trajectory.cpos[phase])
        cvel = np.copy(self.trajectory.cvel[phase])

        return rpos, rvel, lpos, lvel, cpos, cvel

    def get_full_state(self):
        qpos = np.copy(self.sim.qpos())
        qvel = np.copy(self.sim.qvel())

        # TODO: maybe convert to set subtraction for clarity
        # {i for i in range(35)} -
        # {0, 10, 11, 12, 13, 17, 18, 19, 24, 25, 26, 27, 31, 32, 33}

        # this is everything except pelvis x and qw, achilles rod quaternions,
        # and heel spring/foot crank/plantar rod angles
        # note: x is forward dist, y is lateral dist, z is height

        # makes sense to always exclude x because it is in global coordinates and
        # irrelevant to phase-based control. Z is inherently invariant to
        # trajectory despite being global coord. Y is only invariant to straight
        # line trajectories.

        # [ 0] Pelvis y
        # [ 1] Pelvis z
        # [ 2] Pelvis orientation qw
        # [ 3] Pelvis orientation qx
        # [ 4] Pelvis orientation qy
        # [ 5] Pelvis orientation qz
        # [ 6] Left hip roll         (Motor [0])
        # [ 7] Left hip yaw          (Motor [1])
        # [ 8] Left hip pitch        (Motor [2])
        # [ 9] Left knee             (Motor [3])
        # [10] Left shin                        (Joint [0])
        # [11] Left tarsus                      (Joint [1])
        # [12] Left foot             (Motor [4], Joint [2])
        # [13] Right hip roll        (Motor [5])
        # [14] Right hip yaw         (Motor [6])
        # [15] Right hip pitch       (Motor [7])
        # [16] Right knee            (Motor [8])
        # [17] Right shin                       (Joint [3])
        # [18] Right tarsus                     (Joint [4])
        # [19] Right foot            (Motor [9], Joint [5])
        pos_index = np.array([1,2,3,4,5,6,7,8,9,14,15,16,20,21,22,23,28,29,30,34])

        # [ 0] Pelvis x
        # [ 1] Pelvis y
        # [ 2] Pelvis z
        # [ 3] Pelvis orientation wx
        # [ 4] Pelvis orientation wy
        # [ 5] Pelvis orientation wz
        # [ 6] Left hip roll         (Motor [0])
        # [ 7] Left hip yaw          (Motor [1])
        # [ 8] Left hip pitch        (Motor [2])
        # [ 9] Left knee             (Motor [3])
        # [10] Left shin                        (Joint [0])
        # [11] Left tarsus                      (Joint [1])
        # [12] Left foot             (Motor [4], Joint [2])
        # [13] Right hip roll        (Motor [5])
        # [14] Right hip yaw         (Motor [6])
        # [15] Right hip pitch       (Motor [7])
        # [16] Right knee            (Motor [8])
        # [17] Right shin                       (Joint [3])
        # [18] Right tarsus                     (Joint [4])
        # [19] Right foot            (Motor [9], Joint [5])
        vel_index = np.array([0,1,2,3,4,5,6,7,8,12,13,14,18,19,20,21,25,26,27,31])

        if(self.phase == 0):
            ext_state = np.concatenate(self.get_ref_ext_state(self.phaselen - 1))
        else:
            ext_state = np.concatenate(self.get_ref_ext_state(self.phase))

        robot_state = np.concatenate([
            [self.cassie_state.pelvis.position[2] - self.cassie_state.terrain.height], # pelvis height
            self.cassie_state.pelvis.orientation[:],                                 # pelvis orientation
            self.cassie_state.motor.position[:],                                     # actuated joint positions

            self.cassie_state.pelvis.translationalVelocity[:],                       # pelvis translational velocity
            self.cassie_state.pelvis.rotationalVelocity[:],                          # pelvis rotational velocity
            self.cassie_state.motor.velocity[:],                                     # actuated joint velocities

            self.cassie_state.pelvis.translationalAcceleration[:],                   # pelvis translational acceleration

            self.cassie_state.joint.position[:],                                     # unactuated joint positions
            self.cassie_state.joint.velocity[:]                                      # unactuated joint velocities
        ])

        return np.concatenate([robot_state, ext_state])

    def reset_for_normalization(self):
        return self.reset()

    def render(self):
        if self.vis is None:
            self.vis = CassieVis(self.sim, "./cassiemujoco/cassie.xml")

        return self.vis.draw(self.sim)


import numpy as np

from abr_control.interfaces.interface import Interface as BaseInterface
from . import jaco2_rs485

class Interface(BaseInterface):
    """ Interface class for the Jaco2 Kinova arm.

    Handles the overhead of interacting with the Cython code,
    conforms the functions to the ABR_Control interface API.

    Parameters
    ----------
    robot_config : class instance
        contains all relevant information about the arm such as:
        number of joints, number of links, mass information etc.
    """

    def __init__(self, robot_config, display_error_level=2, use_redis=False):
        """ Constructor

        Parameters
        ----------
        display_error_level: int, optional (Default: 3)
            set the level of messages that will be printed to screen
            1. errors, warnings, info, and debugging
            2. errors, warnings, and info
            3: errors and warnings
            4: errors only
            any other int to not display any messages
            NOTE: it is highly recommended to set display_error_level
            to at least 4 to see important error messages
        use_redis: boolean, optional (Default: False)
            True: send joint info to redis server during position mode which is
                  otherwise unavailable during movement
            False: do no import or use redis
        """

        super(Interface, self).__init__(robot_config)
        self.jaco2 = jaco2_rs485.pyJaco2(display_error_level, use_redis)

    def connect(self):
        """ All initial setup, establish RS485 connection
        """

        self.jaco2.Connect()

    def disconnect(self):
        """ Any socket closing etc that must be done
        """

        self.jaco2.Disconnect()

    def get_feedback(self):
        """ Returns a dictionary of relevant information

        Returns the current joint and joint velocity information
        to the controller [radians] [radians/second] respectively
        """

        # convert from degrees from the Jaco into radians
        # Jaco API uses degrees
        feedback = self.jaco2.GetFeedback()
        feedback['q'] = (np.array(feedback['q']) * np.pi / 180.0) % (2 * np.pi)
        feedback['dq'] = np.array(feedback['dq']) * np.pi / 180.0
        return feedback

    def get_torque_load(self):
        """ Returns the torque at each joint in Nm
        """

        return self.jaco2.GetTorqueLoad()['torque_load']

    def init_force_mode(self):
        """ Changes the arm to torque control mode

        To switch to torque mode the arm must be in a position where the
        torque load at the joints is minimized. Suggested to use the
        Config.INIT_TORQUE_POSITION
        """

        self.jaco2.InitForceMode()

    def init_position_mode(self):
        """ Changes the arm into position control mode

        Changes the arm into position control mode, where target joint
        angles can be provided, and the motors switch into servo mode.
        """

        self.jaco2.InitPositionMode()

    def open_hand(self, hand_open):
        """ Incrementally open or close the hand

        Parameters
        ----------
        hand_open : boolean
            True to open hand, False to close hand
        """

        self.jaco2.SendTargetAnglesHand(hand_open)

    def send_forces(self, u):
        """ Applies the set of torques u to the arm.

        NOTE: As part of the Kinova RS485 API, if a torque is not
        applied every 200ms then the arm reverts back to position
        control and the InitForceMode function must be called again.

        Parameters
        ----------
        u : numpy.array
            float value of torques to apply to each joint [Nm]
        """
        self.jaco2.SendForces(u)

    def send_target_angles(self, q):
        """ Moves the arm to the specified joint angles

        Moves the arm to the specified joint angles using
        the on-board PD controller.

        q : numpy.array
            the target joint angles [radians]
        """
        # TODO: need to account for negative degrees
        # convert from radians into degrees the Jaco expects
        q = np.array(q) * 180.0 / np.pi
        self.jaco2.SendTargetAngles(q)


import numpy as np
import os
import random
from PIL import Image

def load(path):
    return PIL2Chainer(Image.open(path))

def save(path, data, rescale=True):
    Chainer2PIL(data, rescale).save(path)

def PIL2Chainer(img, scale=True):
    img = np.array(img)
    if len(img.shape) == 2:
        img = img.astype(np.float32).reshape((1, img.shape[0], img.shape[1]))
    else:
        buf = np.zeros((img.shape[2], img.shape[0], img.shape[1]), dtype=np.uint8)
        for i in range(3):
            buf[i,::] = img[:,:,i]
        img = buf.astype(np.float32)
    if scale:
        # img -= 128
        img /= 256
    return img

def Chainer2PIL(data, rescale=True):
    data = np.array(data)
    if rescale:
        data *= 256
        # data += 128
    if data.dtype != np.uint8:
        data = np.clip(data, 0, 255)
        data = data.astype(np.uint8)
    if data.shape[0] == 1:
        buf = data.astype(np.uint8).reshape((data.shape[1], data.shape[2]))
    else:
        buf = np.zeros((data.shape[1], data.shape[2], data.shape[0]), dtype=np.uint8)
        for i in range(3):
            a = data[i,:,:]
            buf[:,:,i] = a
    img = Image.fromarray(buf)
    return img


#!/usr/bin/env python3
"""
Process microphone input stream endlessly.
The streaming of the input by the underlying sounddevice library
might take ~7% of a single i7 processor's utilization, before
any processing of ours.

With the vad model employed, it's around ~40% of a single cpu's time.
Luckily we have multi-core machines these days ...
"""

import sys
import time
import sounddevice as sd
# required to avoid crashing in assigning the callback input which is a numpy object
import numpy as np
import webrtcvad
import audio2numpy as a2n

EMERGENCY_THRESHOLD = 30
emergencyDetectionDuration = 5.5  # seconds
voiceDetectionTimes = 0

channels = [1]
# translate channel numbers to be 0-indexed
mapping = [c - 1 for c in channels]

# get the default audio input device and its sample rate
device_info = sd.query_devices(None, 'input')
sample_rate = int(device_info['default_samplerate'])

interval_size = 30  # audio interval size in ms
downsample = 1

block_size = sample_rate * interval_size / 1000

# get an instance of webrtc's voice activity detection
vad = webrtcvad.Vad()
vad.set_mode(3)

print("reading audio stream from default audio input device:\n" +
      str(sd.query_devices()) + '\n')
print(F"audio input channels to process: {channels}")
print(F"sample_rate: {sample_rate}")
print(F"window size: {interval_size} ms")
print(F"datums per window: {block_size}")
print()

def emergencyMessage(audioLocation):
    print("Playing emergency audio")
    status = None
    try:
        data, fs = a2n.audio_from_file(audioLocation)
        sd.play(data, fs)
        status = sd.wait()
    except KeyboardInterrupt:
        print('\nInterrupted by user')
    except Exception as e:
        print(type(e).__name__ + ': ' + str(e))
    if status:
        print('Error during playback: ' + str(status))

def voice_activity_detection(audio_data):
    return vad.is_speech(audio_data, sample_rate)

def audio_callback(indata, frames, time, status):
    global voiceDetectionTimes
    # if voiceDetection == True:
    #     print("Audio already detected")
    #     return
    """This is called (from a separate thread) for each audio block."""
    if status:
        print(F"underlying audio stack warning:{status}", file=sys.stderr)

    assert frames == block_size
    # possibly downsample, in a naive way
    audio_data = indata[::downsample, mapping]
    # normalize from [-1,+1] to [0,1], you might not need it with different microphones/drivers
    audio_data = map(lambda x: (x+1)/2, audio_data)
    # adapt to expected float type
    audio_data = np.fromiter(audio_data, np.float16)

    # uncomment to debug the audio input, or run sounddevice's mic input visualization for that
    #print(f'{sum(audio_data)} \r', end="")
    #print(f'min: {min(audio_data)}, max: {max(audio_data)}, sum: {sum(audio_data)}')

    audio_data = audio_data.tobytes()

    voiceDetection = voice_activity_detection(audio_data)
    # use just one line to show the detection status (speech / not-speech)
    print(f'Voz detectada?: {voiceDetection} \r', end="")
    if voiceDetection:
        voiceDetectionTimes += 1
    #     print("Audio detected redirecting to real agent -> :)")

def detectEmergencyFromStream():
    with sd.InputStream(
            device=None,  # the default input device
            channels=max(channels),
            samplerate=sample_rate,
            blocksize=int(block_size),
            callback=audio_callback):
        sd.sleep(int(emergencyDetectionDuration * 1000))
        # avoid shutting down for endless processing of input stream audio
        # while True:
        #     time.sleep(0.1)  # intermittently wake up

for i in range(2):
    emergencyMessage("./playback/emergencia.mp3")
    detectEmergencyFromStream()

print("voiceDetectionTimes: ", voiceDetectionTimes)
if voiceDetectionTimes > EMERGENCY_THRESHOLD:
    print("Emergency call")
else:
    print("Joke call")


import numpy as np
import cv2
import os
import time

from Detector import YoloTrt

class FYDetector():
    '''
    FYDetector - Fast YOLO Detector
    Applies TensorRT and heavily use of PyTorch GPU calculations for best inference time.
    '''

    def __init__(self, config_path, weights_path, names_path, obj_thresh=0.3, filter_classes=None, batch_size=1):
        '''
        config_path    - str       - path to darknet cfg file (yolov3.cfg)
        weights_path   - str       - path to darknet weights file (yolov3.weights)
        names_path     - str       - path to classes names (coco.names)
        obj_thresh     - float     - confidence threshold for objects
        filter_classes - list[int] - list of classes to output
        batch_size     - int       - size of batch for batch inference, constant in time
        '''

        if not os.path.exists(config_path):
            raise ValueError("Invalid config path `" +
                            os.path.abspath(config_path)+"`")
        if not os.path.exists(weights_path):
            raise ValueError("Invalid weight path `" +
                            os.path.abspath(weights_path)+"`")

        self.yolo = YoloTrt.YOLO_TRT(config_path, weights_path, names_path, batch_size=batch_size, obj_thresh=obj_thresh)
        if type(filter_classes) is list:
            self.acceptable_classes = filter_classes

    def detect(self, data):
        if type(data) is list:
            boxes, classes, confs = self.yolo.detect(data)
        else:
            boxes, classes, confs = self.yolo.detect(data)
        return boxes


import tensorflow as tf
from tensorflow import keras
import time

from ThinkGear import ThinkGear
import pygame
from pygame.time import Clock
from pygame.locals import *
import random
from collections import deque
import numpy as np

my_device = ThinkGear("COM5")

# Recreate the exact same model, including its weights and the optimizer
model = tf.keras.models.load_model('saved_models/boosting/cnn40')

# Show the model architecture
model.summary()

pygame.init()

gameDisplay=pygame.display.set_mode((600,600))
pygame.display.set_caption("Platypus")

white=255,255,255
black=0,0,0
red=255,0,0

images = [[], [], [], []]

i = 1
images[0] += [pygame.image.load("data/images/up/up" + str(i) + ".png")]
images[1] += [pygame.image.load("data/images/down/down" + str(i) + ".png")]
images[2] += [pygame.image.load("data/images/left/left" + str(i) + ".png")]
images[3] += [pygame.image.load("data/images/right/right" + str(i) + ".png")]

reads_per_data = 256

raw_data = deque(maxlen=int(reads_per_data * 1)) #192
lastdata = 0

data_size = 180
time_per_data = 1

# Set up the drawing window
screen = pygame.display.set_mode([500, 500])
font = pygame.font.Font(pygame.font.get_default_font(), 36)

# Run until the user asks to quit
running = True
pos = (250,250)
dir = (0,0)
clock = Clock()
speed = 20

score = 0
foodpos = (250, 100)

for reads in range(reads_per_data):
        delay = time.time()
        my_device.fetch_data()
        data = my_device.data
        if 'eeg_raw' in data:
            raw_data.append(data['eeg_raw'])
            lastdata = data['eeg_raw']
        else:
            raw_data.append(lastdata)

        time.sleep(time_per_data/reads_per_data - (time.time() - delay))

running = True
evaluation = [0,0,0,0]
input = np.array(256)

while running:
    input = np.array([list(raw_data)])
    input = input - np.apply_along_axis(lambda a: np.convolve(a, np.ones(11)/11, 'same'), axis=1, arr=input) + 60
    input = input[..., np.newaxis]
    new_evaluation = model.predict(input)[0]
    # print(new_evaluation)
    for i in range(4):
        evaluation[i] = (0.99 * evaluation[i]) + (0.01 * new_evaluation[i])
    direction = max(list(range(4)), key = lambda j: evaluation[j])
    delta  = 1 / float(clock.tick(reads_per_data))
    # Did the user click the window close button?
    if direction == 0:
        print("Up")
        dir = (0, -speed)
    if direction == 1:
        print("Down")
        dir = (0, speed)
    if direction == 2:
        print("Left")
        dir = (-speed, 0)
    if direction == 3:
        print("Right")
        dir = (speed, 0)

    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        if event.type == pygame.KEYDOWN:
            if event.key == K_ESCAPE:
                running = False

    # Fill the background with white
    screen.fill((50, 50, 50))
    if(len(raw_data) == reads_per_data):
        pygame.draw.circle(screen, (0, 150, 0), foodpos, 5)

        # Draw a solid blue circle in the center
        pygame.draw.circle(screen, (200, 200, 0), pos, 10)
    x = (pos[0] + (dir[0] * delta)) % 500
    y = (pos[1] + (dir[1] * delta)) % 500
    pos = (x,y)

    text_surface = font.render("Score: " + str(score), True, (250, 250, 250))
    screen.blit(text_surface, dest=(10,10))

    if ((pos[0]-foodpos[0])**2 + (pos[1] - foodpos[1]) ** 2) < 225:
        score += 1
        print("Score")

        while ((pos[0]-foodpos[0])**2 + (pos[1] - foodpos[1]) ** 2) < 225:
            foodpos = (random.randint(10, 490), random.randint(40, 490))

    # gameDisplay.blit(images[dir][random.randrange(len(images[dir]))], (0,0))
    # pygame.display.update()
    time_save = time.time()
    pygame.display.flip()

    my_device.fetch_data()
    data = my_device.data
    if 'eeg_raw' in data:
        raw_data.append(data['eeg_raw'])
        lastdata = data['eeg_raw']
    else:
        raw_data.append(lastdata)

    for event in pygame.event.get():
            if event.type == QUIT:
                running = False
            elif event.type == KEYDOWN:
                if event.key == K_ESCAPE:
                    running = False

    time.sleep(time_per_data/reads_per_data - (time.time() - time_save))

pygame.quit()
quit()


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

# # Desafio 1
#
# Para esse desafio, vamos trabalhar com o data set [Black Friday](https://www.kaggle.com/mehdidag/black-friday), que reúne dados sobre transações de compras em uma loja de varejo.
#
# Vamos utilizá-lo para praticar a exploração de data sets utilizando pandas. Você pode fazer toda análise neste mesmo notebook, mas as resposta devem estar nos locais indicados.

# ## _Set up_ da análise

# In[1]:

import pandas as pd
import numpy as np
from sklearn import preprocessing

# In[2]:

black_friday = pd.read_csv("black_friday.csv")

# ## Inicie sua análise a partir daqui

# O data set black_friday é uma disponibilização da Analytics Vidhya acessível através do Kaggle. Segundo a plataforma da Codenation, o data set traz algumas variáveis relativas à transações comerciais realizadas durante a Black Friday em uma determinada loja de varejo. Cada observação é relativa a um determinado item comprado por um usuário. O objetivo deste documento não é explorar os dados do dataframe, mas sim responder as questões propostas pelo desafio 1 do ACELERADEV Data Science. Portanto, não houve qualquer alteração no documento, exceto a complementação com as respostas necessárias para cada questão.

# In[12]:

black_friday.head(3)

# ## Questão 1
#
# Quantas observações e quantas colunas há no dataset? Responda no formato de uma tuple `(n_observacoes, n_colunas)`.

# Este exercício poderia ser resolvido de duas formas, a primeira utilizando apenas o atributo shape que retornaria a quantidade de linhas e a quantidade de colunas do dataframe. Contudo, crei duas variáveis para que possamos ver que é possível mostrar separadamente a quantidade de linhas, através do indicador [0] e a quantidade de colunas com o indicador [1]. Por fim, foi reservado no atributo tupla o valor de ambos.

# In[3]:

def q1():
    n_observacoes = black_friday.shape[0]
    n_colunas = black_friday.shape[1]
    tupla = (n_observacoes, n_colunas)
    return (tupla)

# ## Questão 2
#
# Há quantas mulheres com idade entre 26 e 35 anos no dataset? Responda como um único escalar.

# Novamente usamos o atributo shape para retornar a contagem de linhas, mas dessa vez precisamos de um filtro. Para isso, é usada a forma abaixo, que faz um "corte" dentro do dataframe, criando uma lista com True e False para todas as linhas que atendem as condições descritas, ou seja, mulheres entre 26-35 anos.

# In[4]:

def q2():
    return black_friday[(black_friday['Gender']=='F') & (black_friday['Age']=='26-35')].shape[0]

# ## Questão 3
#
# Quantos usuários únicos há no dataset? Responda como um único escalar.

# O método nunique provém da biblioteca do numpy (importada no início do código). Seu objetivo é contar quantos valores únicos há em um array. Usando como parâmetro a coluna "User_ID" é fácil descobrir quantos valores usuários únicos há.

# In[5]:

def q3():
    return int(black_friday["User_ID"].nunique())

# ## Questão 4
#
# Quantos tipos de dados diferentes existem no dataset? Responda como um único escalar.

# Mais uma vez o método unique pode nos salvar. Dessa vez, é utilizado em conjunto com um recurso de outra biblioteca (pandas). O dtypes retorna quais são os tipos de dados existentes em um determinado conjunto de dados. Abaixo então vemos que, é selecionado os tipos (dtypes) únicos(nunique) do dataframe black_friday e para garantir que será um valor único escalar, utiliza-se o int no início.

# In[6]:

def q4():
    return int(black_friday.dtypes.nunique())

# ## Questão 5
#
# Qual porcentagem dos registros possui ao menos um valor null (`None`, `ǸaN` etc)? Responda como um único escalar entre 0 e 1.

# O valor percentual de uma determinada ocorrência é calculado a partir de x/total de casos, onde x é o parâmetro utilizado (proposto pelo exercício). Vejamos, o shape, como já visto, faz uma contagem de linhas, enquanto que o dropna remove os valores faltantes. Podemos então utilizar a técnica da primeira questão, fazer um "corte" no nosso dataframe a partir deste filtro. Você deve estar se questionando "mas isso não trará então o percentual dos valores preenchidos?", por isso utiliza-se o 1- no início da equação, ou seja, o valor complementar.

# In[7]:

def q5():
    percentual = 1-(black_friday.dropna().shape[0]/black_friday.shape[0])
    return percentual

# ## Questão 6
#
# Quantos valores null existem na variável (coluna) com o maior número de null? Responda como um único escalar.

#
Na fórmula abaixo temos uma série de métodos para resolver este problema. Se, por curiosidade, fizer o teste executando cada um dos métodos, um de cada vez, é possível perceber o filtro que cada método faz. Inicialmente o isnull identifica quais são os dados são nulos de forma booleana, então sum mostra a somatória dos nulls em cada variável e o max finalmente retorna qual é coluna possui mais dados nulos. Lembrando que essa forma é para alcançar o resultado pedido no exercício, porém, esses métodos podem ser executados de forma separa para outras finalidades.

# In[12]:

def q6():
    return black_friday.isnull().sum().max()

# ## Questão 7
#
# Qual o valor mais frequente (sem contar nulls) em `Product_Category_3`? Responda como um único escalar.

# Este exercício pode ser realizado a partir do que foi aplicado na questão 2, com o método dropna, recorta-se os valores faltantes. O que foi feito abaixo então:
#
Atribuiu-se um novo dataframe bf para não perder o dataframe anterior;
#
Utilizou-se o método mode que retorna qual é o valor mais frequente;
#
Mudou-se este retorno para inteiro, afim de retornar um único valor escalar.

# Assim então, tem-se que o valor mais frequente é: 16

# In[9]:

def q7():
    bf = black_friday.dropna()
    return int(bf['Product_Category_3'].mode())

# ## Questão 8
#
# Qual a nova média da variável (coluna) `Purchase` após sua normalização? Responda como um único escalar.

# A normalização ajusta os valores que estão em diferentes medidas para uma escala fictícia comum, entre -1 e 1, onde resultados entre -1 e 0, representam que o valor puro era negativo.

# Nas importações, é possível ver que além do pandas e numpy, também foi o preprocessing da biblioteca sklearn. Através dele é possível executar a normalização de forma mais prática.
#
Primeiro foi retornado como um array em x a coluna "Purchase"
Em seguida, o método MinMaxScaler transforma cada elemento, de modo que esteja no intervalo especificado no conjunto. Como não está estabelecido, por default ele considera entre 0 e 1.
Em x_scaled os dados são ajustados e transformados em um array
#
E por fim, é retornado a média deste array

# In[10]:

def q8():
    x = black_friday['Purchase'].values #retorna um array
    min_max_scaler = preprocessing.MinMaxScaler()
    x_scaled = min_max_scaler.fit_transform(x.reshape(-1,1))
    return x_scaled.mean()

# ## Questão 9
#
# Quantas ocorrências entre -1 e 1 inclusive existem da variável `Purchase` após sua padronização? Responda como um único escalar.

# A padronização consiste em subtrair de um valor de um conjunto de dados a sua média e dividir o resultado pelo desvio padrão do conjunto ou variável. Novamente, será utilizado recursos do preprocessing, mas agora quem irá nos auxiliar é o método StandardScaler() que faz a padronização de forma simples. Os passos são os seguintes:
#
Primeiro foi retornado como um array em x a coluna "Purchase" em y
Em seguida, o método StandardScaler transforma cada elemento
Em y_scaled os dados são ajustados e transformados em um array
Em ocorrencias é filtrado quais elementos estão entre -1 e 1
E então, é retornado o tamanho (pelo len) este valor único.

# In[1]:

def q9():
    from sklearn.preprocessing import StandardScaler
    y = black_friday['Purchase'].values
    standard = StandardScaler()
    y_standard = standard.fit_transform(y.reshape(-1, 1))
    ocorrencias = y_standard[(y_standard > -1) & (y_standard < 1)]
    return len (ocorrencias)

# ## Questão 10
#
# Podemos afirmar que se uma observação é null em `Product_Category_2` ela também é em `Product_Category_3`? Responda com um bool (`True`, `False`).

# Existem N formas para fazer essa comparação, uma delas é criando dois atributos auxiliares (aux1 e aux2) que retorna valores booleanos, sendo True caso o elemento seja nulo e caso o contrário, o retorno é False. Depois basta comparar as linhas pelo shape[0] de aux1 e aux 2 para saber se são iguais, ou seja, se uma observação é null em `Product_Category_2` ela também é em `Product_Category_3`.

# In[10]:

def q10():
    aux1 = black_friday['Product_Category_2'].isnull()
    aux2 = black_friday['Product_Category_3'].isnull()
    return bool (aux1.shape[0] == aux2.shape[0])

# In[11]:

q10()


from numpy import mean
from wordstats import Word, WordInfo
from zeeguu.core import model
from zeeguu.core.language.difficulty_estimator_strategy import DifficultyEstimatorStrategy
from zeeguu.core.util.text import split_words_from_text
from zeeguu.core.model.word_knowledge.word_interaction_history import WordInteractionHistory
from zeeguu.core.model import UserWord, Language
from nltk.stem import SnowballStemmer
from collections import defaultdict, Counter

from zeeguu.core.constants import WIH_READ_CLICKED, WIH_READ_NOT_CLICKED_IN_SENTENCE, WIH_READ_NOT_CLICKED_OUT_SENTENCE, WIH_WRONG_EX_RECOGNIZE,\
WIH_WRONG_EX_TRANSLATE,WIH_WRONG_EX_CHOICE,WIH_WRONG_EX_MATCH

class WordHistoryDifficultyEstimator(DifficultyEstimatorStrategy):

    CUSTOM_NAMES = ["history"]

    def __init__(self, language: 'model.Language', user: 'model.User'):
        self.user = user
        self.language = language
        self.score_map = dict()

        # determines word scores
        #words_history = WordInteractionHistory.find_all_word_histories_for_user_language(user, language)

        #words_found = [w.word.word.lower() for w in words_history]
        #event_history = [w.interaction_history for w in words_history]

        #words_score = [max(1 - len(e) / 1, 0) for e in event_history]

    # creates estimator that determines the ratio of new words for a given text
    @classmethod
    def recurrence(cls, language: 'model.Language', user: 'model.User'):
        """
                This estimator computes the ratio of new words for a given user and language
                :param language: language of the text that needs to be estimated
                :param user: the user for which the difficulty estimation needs to be done
                :rtype: WordHistoryDifficultyEstimator
                :return: WordHistoryDifficultyEstimator with initialized user, language and word => score map
                        which can be used for determining scores for multiple articles for the same user and language
        """

        estimator = cls(language, user)

        # determine word scores
        words_history = WordInteractionHistory.find_all_word_histories_for_user_language(user, language)
        words_found = [w.word.word for w in words_history]
        event_history = [w.interaction_history for w in words_history]

        words_score = [0 for e in event_history]

        estimator.score_map = dict(zip(words_found, words_score))

        return estimator

    @classmethod
    def recurrence_until_timestamp(cls, language: 'model.Language', user: 'model.User',max_timestamp):
        """
                This estimator computes the ratio of new words for a given user and language
                :param language: language of the text that needs to be estimated
                :param user: the user for which the difficulty estimation needs to be done
                :rtype: WordHistoryDifficultyEstimator
                :return: WordHistoryDifficultyEstimator with initialized user, language and word => score map
                        which can be used for determining scores for multiple articles for the same user and language
        """

        estimator = cls(language, user)

        # determine word scores
        words_history = WordInteractionHistory.find_all_word_histories_for_user_language(user, language)

        words_found = []
        event_history = []
        for wh in words_history:
            history = [event for event in wh.interaction_history if event.seconds_since_epoch <= max_timestamp]

            if len(history) > 0:
                words_found.append(wh.word.word)
                event_history.append(history)

        words_score = [0 for e in event_history]

        estimator.score_map = dict(zip(words_found, words_score))

        return estimator

    @classmethod
    def difficulty(cls, language: 'model.Language', user: 'model.User', mode = 1):
        """
                This estimator computes the ratio of new words for a given user and language
                :param language: language of the text that needs to be estimated
                :param user: the user for which the difficulty estimation needs to be done
                :rtype: WordHistoryDifficultyEstimator
                :return: WordHistoryDifficultyEstimator with initialized user, language and word => score map
                        which can be used for determining scores for multiple articles for the same user and language
        """

        estimator = cls(language, user)

        # determine word scores
        words_history = WordInteractionHistory.find_all_word_histories_for_user_language(user, language)

        words_found = [w.word.word for w in words_history]
        event_history = [w.interaction_history for w in words_history]

        words_score = []

        if mode == 1:

            for e in event_history:
                if len(e) > 10:
                    words_score.append(0)
                else:
                    N_seen_context = sum([event == WIH_READ_NOT_CLICKED_IN_SENTENCE for event in event_history])
                    N_seen = sum([event == WIH_READ_NOT_CLICKED_OUT_SENTENCE for event in event_history])
                    # N_clicked = sum([event == WIH_READ_CLICKED for event in event_history])

                    if N_seen_context > 3 or N_seen > 6:
                        words_score.append(0)
                    else:
                        words_score.append(min(1 - N_seen_context / 4, 1 - N_seen / 7))
        elif mode == 2:
            for e in event_history:
                words_score.append(max(1 - len(e)/10,0))

        elif mode == 3:
            length = 0
            for e in event_history[::-1]:
                if e is not WIH_READ_CLICKED or e is not WIH_WRONG_EX_CHOICE or e is not WIH_WRONG_EX_MATCH or e is not WIH_WRONG_EX_RECOGNIZE or\
                    e is not WIH_WRONG_EX_TRANSLATE:
                    length += 1
                words_score.append(max(1 - len(e)/10,0))

        estimator.score_map = dict(zip(words_found, words_score))

        return estimator

    @classmethod
    def difficulty_until_timestamp(cls, language: 'model.Language', user: 'model.User', max_timestamp, mode = 1, scaling = 10.0, scaling2 = 20.0):
        """
                This estimator computes the ratio of new words for a given user and language
                :param language: language of the text that needs to be estimated
                :param user: the user for which the difficulty estimation needs to be done
                :rtype: WordHistoryDifficultyEstimator
                :return: WordHistoryDifficultyEstimator with initialized user, language and word => score map
                        which can be used for determining scores for multiple articles for the same user and language
        """

        estimator = cls(language, user)

        # determine word scores
        words_history = WordInteractionHistory.find_all_word_histories_for_user_language(user, language)

        words_found = []
        event_history = []
        for wh in words_history:
            history = [event for event in wh.interaction_history if event.seconds_since_epoch <= max_timestamp]

            if history:
                words_found.append(wh.word.word)
                event_history.append(history)

        words_score = []

        if mode == 1:

            for e in event_history:
                N_seen_context = sum([event.event_type == WIH_READ_NOT_CLICKED_IN_SENTENCE for event in e])
                N_seen = sum([event.event_type == WIH_READ_NOT_CLICKED_OUT_SENTENCE for event in e])
                # N_clicked = sum([event == WIH_READ_CLICKED for event in event_history])

                words_score.append(max(0,1 - N_seen_context / scaling - N_seen / scaling2))
        elif mode == 2:
            for e in event_history:
                words_score.append(max(1 - len(e)/scaling,0))

        elif mode == 3:
            length = 0
            for e in event_history:
                events= e[::-1]
                for event in events:
                    if event.event_type is not WIH_READ_CLICKED or event.event_type is not WIH_WRONG_EX_CHOICE or event.event_type is not WIH_WRONG_EX_MATCH or event.event_type is not WIH_WRONG_EX_RECOGNIZE or \
                            event.event_type is not WIH_WRONG_EX_TRANSLATE:
                        length += 1
                words_score.append(max(1 - len(e)/scaling,0))

        estimator.score_map = dict(zip(words_found, words_score))

        return estimator

    def estimate_difficulty(self, text: str):
        """
        This estimator computes the difficulty based on the scoring map
        :param text: See DifficultyEstimatorStrategy
        :rtype: dict
        :return: The dictionary contains the keys and return types
                    normalized: float (0<=normalized<=1)
                    discrete: string [EASY, MEDIUM, HARD]
        """

        # Calculate difficulty for each word
        words = split_words_from_text(text)
        stemmer = SnowballStemmer(self.language.name.lower())
        words = [stemmer.stem(w.lower()) for w in words]

        # frequency and length
        word_frequency = Counter(words)
        total_words = len(words)

        # score per word
        word_scores = {w: self.word_difficulty(self.score_map, True, w) for w in word_frequency}

        # If we can't compute the text difficulty, we estimate hard
        if (len(word_scores)) == 0:
            difficulty_scores = dict(
                median=1.0,
                median_unique=1.0,
                normalized=1.0,
                discrete= "HARD",
                unique_ratio=1.0
            )
        else:

            # words above median score
            word_scores_sorted= sorted(word_scores.items(), key=lambda item: item[1])
            center = int(round(len(word_scores_sorted) / 2, 0))
            word_scores_median = dict(word_scores_sorted[center:])
            total_words_median = sum([word_frequency[w] for w in word_scores_median])

            # If we can't compute the text difficulty, we estimate hard

            difficulty_scores = dict(
                median=sum([ s*word_frequency[w] for w, s in word_scores_median.items()])/total_words_median,
                median_unique=mean(list(word_scores_median.values())),
                normalized=sum([ s*word_frequency[w] for w, s in word_scores.items()])/total_words,  # Originally called 'score_average'
                #todo: discrete=discrete_difficulty,               # Originally called 'estimated_difficulty'
                unique_ratio=mean(list(word_scores.values()))      # Ratio of unknown
            )

        return difficulty_scores

    @classmethod
    def discrete_text_difficulty(cls, median_difficulty: float):
        """
        :param median_difficulty:
        :return: a symbolic representation of the estimated difficulty
         the values are between "EASY", "MEDIUM", and "HARD"
        """
        if median_difficulty < 0.3:
            return "EASY"
        if median_difficulty < 0.4:
            return "MEDIUM"
        return "HARD"

    @classmethod
    # TODO: must test this thing
    def word_difficulty(cls, known_probabilities: dict, personalized: bool, word: Word):
        """
        # estimate the difficulty of a word, given:
            :param word_info:
            :param known_probabilities:
            :param personalized:
            :param word:

        :return: a normalized value where 0 is (easy) and 1 is (hard)
        """

        # Assume word is difficult and unknown
        estimated_difficulty = 1.0

        # Check if the user knows the word
        try:
            known_probability = known_probabilities[word]  # Value between 0 (unknown) and 1 (known)
        except KeyError:
            known_probability = None

        if personalized and known_probability is not None:
            estimated_difficulty = float(known_probability)

        return estimated_difficulty