import json
import uuid
import random
import string
from re import search
from time import time
from hashlib import md5
from pathlib import Path
from urllib import parse
from asyncio import sleep as asyncsleep
from typing import Dict, List, Tuple, Union, Literal

from httpx import AsyncClient
from nonebot.log import logger

from .__meta__ import (
    POOL_API,
    ROLE_API,
    ROOT_URL,
    LOCAL_DIR,
    TOKEN_API,
    EXPIRE_SEC,
    GACHA_TYPE,
    AUTHKEY_API,
    CLIENT_SALT,
    CLIENT_TYPE,
    CLIENT_VERSION,
    ROOT_OVERSEA_URL,
    datetime_with_tz,
)

def formatInput(input: str, find: Literal["url", "cookie"]) -> Dict:
    """
    输入内容格式化，可根据输入内容提取抽卡记录链接或米游社 Cookie

    * ``param input: str`` 输入内容
    * ``param find: Literal["url", "cookie"]`` 提取内容类型
    - ``return: Dict`` 提取结果，格式为 ``{"url": "链接"}`` 或 ``{"login_ticket": "xx", "stoken": "xx", ...}``，出错时返回 ``{}``
    """  # noqa: E501

    if find == "cookie":
        if not input:
            return {}
        cookieDict = {}
        params = [
            "account_id",
            "cookie_token",
            "login_ticket",
            "login_uid",
            "ltoken",
            "ltuid",
            "mid",
            "stoken",
            "stuid",
        ]
        for param in params:
            match = search(rf"(^| ){param}=([^;]*)", input)
            if match:
                # 去除结尾分号、空格、引号
                cookieDict[param] = match.group(2).rstrip("; '\"")
        # account_id, login_uid, stuid, ltuid 同为米游社 ID
        for param in ["account_id", "login_uid", "stuid", "ltuid"]:
            if cookieDict.get(param):
                for p in ["account_id", "login_uid", "stuid", "ltuid"]:
                    cookieDict[p] = cookieDict[param]
                break
        # return "; ".join([f"{k}={v}" for k, v in cookieDict.items()])
        return cookieDict
    elif find == "url":
        # ref: https://ihateregex.io/expr/url
        urlReg = r"https?:\/\/(www\.)?[-a-zA-Z0-9@:%._\+~#=]{1,256}\.[a-zA-Z0-9()]{1,6}\b([-a-zA-Z0-9()!@:%_\+.~#?&\/\/=]*)"  # noqa: E501
        matchUrl = search(urlReg, input.replace("&", "&"))
        if matchUrl:
            # 找到链接后替换 Host
            root, query = matchUrl[0].split("?")
            urlRoot = (
                ROOT_OVERSEA_URL
                if any(x in root for x in ["webstatic-sea", "hk4e-api-os"])
                else ROOT_URL
            )
            return {"url": urlRoot + "?" + query}
        return {}

async def configHelper(qq: str, data: Dict = {}) -> Dict:
    """
    配置缓存助手，既可根据 QQ 读取配置，也可根据数据写入/删除配置缓存

    * ``param qq: str`` 目标 QQ，为 ``"0"`` 时返回全部配置
    * ``param data: Dict = {}`` 配置数据，根据是否传入决定更新或读取。配置中有链接、Cookie 或抽卡记录缓存 或 配置中包含 force/delete 参数时才写入
    - ``return: Dict`` 目标 QQ 的配置数据，删除时返回 ``{}``，出错时返回 ``{"error": "错误信息"}``
    """  # noqa: E501

    cfgFile = LOCAL_DIR / "config.json"
    # 尝试读取本地配置文件
    if not cfgFile.exists():
        cfgFile.write_text("{}", encoding="utf-8")
        cfg = {}
    else:
        cfg = json.loads(cfgFile.read_text(encoding="utf-8"))
    assert isinstance(cfg, Dict)
    # 根据是否传入配置数据决定写入/删除或读取
    if data:
        # 配置中有链接、Cookie 或抽卡记录缓存 或 配置中包含 force/delete 参数时视为有效输入
        if data.get("delete"):
            validInput, delMode = True, data["delete"]
            modeStr = "配置已删除" if delMode == "全部配置" else "记录已删除"
            data.pop("delete")
        elif data.get("force"):
            validInput = bool(data["force"])
            delMode, modeStr = "", ("已更新" if validInput else "未更新")
            data.pop("force")
        else:
            validInput = bool(data.get("url") or data.get("cookie") or data.get("logs"))
            delMode, modeStr = "", ("已更新" if validInput else "未更新")
        try:
            if validInput:
                if delMode == "全部配置":
                    cfg.pop(qq)
                else:
                    cfg[qq] = data
                cfgFile.write_text(
                    json.dumps(cfg, ensure_ascii=False, indent=2), encoding="utf-8"
                )
            else:
                cfg[qq] = {}
                logger.info(f"QQ{qq} 的配置写入被跳过")
            logger.info(f"QQ{qq} 的配置{modeStr}\n{'{}' if delMode else cfg[qq]}")
            return {} if delMode else cfg[qq]
        except Exception as e:
            logger.opt(exception=e).error(f"QQ{qq} 的配置缓存{modeStr}失败")
            return {"error": f"QQ{qq} 的配置缓存{modeStr}失败！"}
    else:
        return cfg if qq == "0" else cfg.get(qq, {"error": f"暂无 QQ{qq} 的抽卡记录配置！"})

async def logsHelper(file: Union[Path, str], data: Dict = {}) -> Tuple[str, Dict]:
    """
    抽卡记录缓存助手，既可根据 ``file`` 路径读取抽卡记录，也可根据 ``data`` 数据写入/删除抽卡记录缓存

    * ``param file: Union[Path, str]`` 抽卡记录缓存文件路径
    * ``param data: Dict = {}`` 抽卡记录数据，根据是否传入决定写入/删除或读取
    - ``return: Tuple[str, Dict]`` 抽卡记录所属 UID（出错时返回错误信息）、抽卡记录数据（写入/删除时固定返回 ``{}``）
    """

    uid = search(r"gachalogs-([0-9]{9}).json", str(file))
    if not uid:
        raise ValueError(f"记录所属 UID 不存在！{str(file)}")
    uid, logsFile = uid.group(1), Path(file)
    # 根据是否传入抽卡记录数据决定写入/删除或读取
    if data:
        delMode = bool(data.get("delete"))
        modeStr = "删除" if delMode else "更新"
        try:
            if delMode:
                logsFile.unlink(missing_ok=True)
            else:
                logsFile.write_text(
                    json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8"
                )
            logger.info(f"UID{uid} 的抽卡记录已{modeStr}")
            return uid, {}
        except Exception as e:
            logger.opt(exception=e).error(f"UID{uid} 的抽卡记录缓存{modeStr}失败")
            return f"UID{uid} 的抽卡记录缓存{modeStr}失败！", {}
    elif logsFile.exists():
        logs = json.loads(logsFile.read_text(encoding="utf-8"))
        assert isinstance(logs, Dict)
        return uid, logs
    else:
        raise ValueError(f"UID{uid} 的本地抽卡记录不存在！")

async def queryMihoyo(
    cookie: str, aType: Literal["获取令牌", "获取角色", "获取卡池", "生成密钥"], data: Dict = {}
) -> Dict:
    """
    米哈游接口请求，支持以下功能：
    - 获取 ``stoken`` 令牌
    - 获取指定 ``cookie`` 名下原神游戏角色数据
    - 查询最新祈愿活动 ``gacha_id`` 卡池数据
    - 生成用于抽卡记录查询的 ``authkey`` 密钥

    * ``param cookie: str`` 米哈游通行证 Cookie，获取令牌和获取卡池时可传入空
    * ``param aType: Literal["获取令牌", "获取角色", "获取卡池", "生成密钥"]`` 请求类型
    * ``param data: Dict = {}`` 请求数据，根据需要被赋值为 ``params`` ``content`` 等
    - ``return: Dict`` 请求结果，出错时返回 ``{"error": "错误信息"}``
    """

    if aType == "获取卡池":
        # 可能不需要每次都向米游社发起请求，此处将返回固定为常驻祈愿活动数据
        return {"type": "200", "pool": "fecafa7b6560db5f3182222395d88aaa6aaac1bc"}
    t = str(int(time()))
    r = "".join(random.sample(string.ascii_lowercase + string.digits, 6))
    m = md5(f"salt={CLIENT_SALT}&t={t}&r={r}".encode()).hexdigest()
    headers = {
        "cookie": cookie,
        "ds": f"{t},{r},{m}",
        "host": "api-takumi.mihoyo.com",
        "referer": "https://app.mihoyo.com",
        "user-agent": "okhttp/4.8.0",
        "x-rpc-app_version": CLIENT_VERSION,
        "x-rpc-channel": "mihoyo",
        "x-rpc-client_type": CLIENT_TYPE,
        "x-rpc-device_id": str(uuid.uuid3(uuid.NAMESPACE_URL, cookie)),
        "x-rpc-device_model": "SM-977N",
        "x-rpc-device_name": "Samsung SM-G977N",
        "x-rpc-sys_version": "12",
    }
    if aType == "获取令牌":
        headers.update(
            {
                "origin": "https://webstatic.mihoyo.com",
                "referer": "https://webstatic.mihoyo.com/",
                "user-agent": (
                    "Mozilla/5.0 (Linux; Android 12; SM-G977N Build/SP1A.210812.016"
                    "; wv) AppleWebKit/537.36 (KHTML, like Gecko) Version/4.0 Chrome"
                    f"/104.0.5112.69 Mobile Safari/537.36 miHoYoBBS/{CLIENT_VERSION}"
                ),
                "x-requested-with": "com.mihoyo.hyperion",
            }
        )
    elif aType == "生成密钥":
        headers["content-type"] = "application/json; charset=UTF-8"
    async with AsyncClient() as client:
        try:
            if aType == "获取令牌":
                res = await client.get(TOKEN_API, headers=headers, params=data)
                rt = {"stoken": res.json()["data"]["list"][0]["token"]}
            elif aType == "获取角色":
                res = await client.get(ROLE_API, headers=headers)
                rt = [
                    char
                    for char in res.json()["data"]["list"]
                    if char["game_biz"] == "hk4e_cn"
                ][0]
            elif aType == "获取卡池":
                res = await client.get(POOL_API, params={"ts": str(time())[:8]})
                rt = {
                    "type": "200",
                    "pool": [
                        p["gacha_id"]
                        for p in res.json()["data"]["list"]
                        if p["gacha_type"] == 200
                    ][0],
                }
            elif aType == "生成密钥":
                res = await client.post(
                    AUTHKEY_API,
                    headers=headers,
                    content=json.dumps(data, ensure_ascii=False),
                )
                rt = {"authkey": res.json()["data"]["authkey"]}
            else:
                raise ValueError(f"未知的请求类型：{aType}")
        except Exception as e:
            logger.opt(exception=e).error(f"米游社 {aType} 请求失败")
            return {"error": f"[{e.__class__.__name__}] 未能成功{aType}！"}
    return rt

async def checkAuthKey(url: str, skipFmt: bool = True) -> str:
    """
    抽卡记录链接验证，检查传入抽卡记录链接 `url` 是否有效（AuthKey 有效期 24h）

    * ``param url: str`` 抽卡记录链接
    * ``param skipFmt: bool = True`` 是否跳过格式检查，传入 ``False`` 将根据格式检查结果修正抽卡记录链接
    - ``return: str`` 抽卡记录链接，出错时返回错误信息
    """

    if not skipFmt:
        urlRes = formatInput(url, find="url")
        if not urlRes:
            return "未找到有效的抽卡记录链接！"
        else:
            url = urlRes["url"]
    logger.debug(f"验证抽卡记录链接 {url}")
    async with AsyncClient() as client:
        try:
            res = await client.get(url)
            resJson = res.json()
            # checkFile = LOCAL_DIR / f"checkAuthKey-{int(time())}.json"
            # with open(checkFile, "w", encoding="utf-8") as f:
            #     resJson["url"] = url
            #     json.dump(resJson, f, ensure_ascii=False, indent=2)
        except json.decoder.JSONDecodeError:
            return "链接返回无法解析！"
        except Exception as e:
            logger.opt(exception=e).error("抽卡记录链接验证出错")
            return f"[{e.__class__.__name__}] 链接验证出错！"
    if not resJson["data"]:
        if resJson.get("message", "") == "authkey timeout":
            return "链接 AuthKey 失效！"
        if resJson.get("message", "") == "authkey error":
            return "链接 AuthKey 错误！"
        # ref: https://webstatic.mihoyo.com/admin/mi18n/hk4e_cn/20190926_5d8c80193de82/20190926_5d8c80193de82-zh-cn.json  # noqa: E501
        logger.error(
            f"抽卡记录链接有问题 {resJson.get('retcode', 777)} {resJson.get('message', '')}"
        )
        return f"[{resJson.get('retcode', 777)}]链接有问题！"
    return url

async def updateLogsUrl(url: str, cookie: str) -> Tuple[str, Dict]:
    """
    抽卡记录链接更新，可根据 `cookie` 初始化或更新抽卡记录链接

    * ``param url: str`` 抽卡记录链接
    * ``param cookie: str`` 含有 `stoken` 字段的米游社 Cookie
    - ``return: Tuple[str, Dict]`` 抽卡记录链接（出错时返回 ``"错误信息"``）、角色及 Cookie 字典数据（出错或旧链接未过期时返回 ``{}``）
    """  # noqa: E501

    # 检查传入链接是否仍然有效
    url = await checkAuthKey(url)
    if url.startswith("https://"):
        return url, {}
    # 提取 cookie 中有效字段字典
    usefulCk = formatInput(cookie, find="cookie")
    # Cookie 验证及补全
    if not usefulCk:
        return "无法自动更新链接！", {}
    elif not usefulCk.get("account_id"):
        return "Cookie 缺少米游社 ID 数据！", {}
    elif not usefulCk.get("stoken"):
        mysId, loginTicket = usefulCk.get("account_id"), usefulCk.get("login_ticket")
        if not loginTicket:
            return "Cookie 缺少 login_ticket 数据！", {}
        data = {"login_ticket": loginTicket, "token_types": "3", "uid": mysId}
        stokenRes = await queryMihoyo("", "获取令牌", data=data)
        if not stokenRes.get("stoken"):
            return "获取令牌失败！", {}
        usefulCk["stoken"] = stokenRes["stoken"]
    ckStr = "; ".join([f"{k}={v}" for k, v in usefulCk.items()])
    # 获取 Cookie 名下角色数据
    role = await queryMihoyo(ckStr, "获取角色")
    if role.get("error"):
        return role["error"], {"cookie": ckStr}
    # 更新抽卡记录链接中的 AuthKey
    data = {
        "auth_appid": "webview_gacha",
        "game_biz": role["game_biz"],
        "game_uid": role["game_uid"],
        "region": role["region"],
    }
    authkeyRes = await queryMihoyo(ckStr, "生成密钥", data=data)
    if not authkeyRes.get("authkey"):
        return "生成密钥失败！", {"cookie": ckStr}
    # 更新抽卡记录链接中的卡池 ID
    poolRes = await queryMihoyo("", "获取卡池")
    if not poolRes.get("pool"):
        poolRes = {
            "type": "200",  # init_type
            "pool": "fecafa7b6560db5f3182222395d88aaa6aaac1bc",  # gacha_id
        }
    # 提取传入 url 中请求参数字典
    parsed = parse.urlparse(url)
    querys = dict(parse.parse_qsl(str(parsed.query)))
    # 更新抽卡记录链接
    if not querys:
        querys = {
            "authkey_ver": "1",
            "sign_type": "2",
            "auth_appid": "webview_gacha",
            "init_type": poolRes["type"],
            "gacha_id": poolRes["pool"],
            "timestamp": str(int(time())),
            "lang": "zh-cn",
            "device_type": "mobile",
            # "ext": {"loc":{"x":-672.817138671875,"y":122.54130554199219,"z":-87.4539794921875},"platform":"IOS"},  # noqa: E501
            # "game_version": "CNRELiOS2.8.0_R9182063_S9401797_D9464149",
            "plat_type": "android",  # ios
            "region": role["region"],
            "authkey": authkeyRes["authkey"],
            "game_biz": role["game_biz"],
            "gacha_type": "301",
            "page": "1",
            "size": "6",
            "end_id": 0,
        }
    else:
        querys.update(
            {
                "init_type": poolRes["type"],
                "gacha_id": poolRes["pool"],
                "timestamp": str(int(time())),
                "region": role["region"],
                "authkey": authkeyRes["authkey"],
                "game_biz": role["game_biz"],
            }
        )
    urlRoot = (
        ROOT_OVERSEA_URL
        if any(x in url for x in ["webstatic-sea", "hk4e-api-os"])
        else ROOT_URL
    )
    url = urlRoot + "?" + parse.urlencode(querys)
    return await checkAuthKey(url), {"role": role, "cookie": ckStr}

def getGachaLogsApi(logUrl: str, gachaType: str, page: int, endId: int) -> str:
    """
    指定类型抽卡记录接口生成

    * ``param logUrl: str`` 抽卡记录链接
    * ``param gachaType: str`` 祈愿类型（请求参数 ``gacha_type``）
    * ``param page: int`` 页码（请求参数 ``page``）
    * ``param endId: int`` 结束 ID（请求参数 ``end_id``）
    - ``return: str`` 指定类型抽卡记录接口
    """

    parsed = parse.urlparse(logUrl)
    querys = dict(parse.parse_qsl(str(parsed.query)))
    querys.update(
        {
            "lang": "zh-cn",
            "gacha_type": gachaType,
            "page": str(page),
            "size": "20",
            "end_id": str(endId),
        }
    )
    return logUrl.split("?")[0] + "?" + parse.urlencode(querys)

async def getSingleTypeLogs(logUrl: str, gachaType: str) -> List:
    """
    指定类型抽卡记录获取

    * ``param logUrl: str`` 抽卡记录链接
    * ``param gachaType: str`` 祈愿类型，实际类型应为 ``Literal["100", "200", "301", "302"]``
    - ``return: List`` 指定类型抽卡记录
    """

    logsList, endId, page = [], 0, 1
    async with AsyncClient() as client:
        while True:
            logger.debug(f"正在获取 {GACHA_TYPE[gachaType]} 第 {page} 页")
            api = getGachaLogsApi(logUrl, gachaType, page, endId)
            res = await client.get(api)
            try:
                resJson = res.json()
                if resJson.get("message") == "visit too frequently":
                    logger.info("访问过于频繁，等待 3 秒后重试")
                    await asyncsleep(3)
                elif resJson.get("data", {}).get("list", []):
                    # 成功解析记录数据，更新请求参数
                    logsList.extend(resJson["data"]["list"])
                    endId = resJson["data"]["list"][-1]["id"]
                    page += 1
                    await asyncsleep(0.6)
                else:
                    # 未解析到记录数据，跳出循环
                    break
            except json.decoder.JSONDecodeError:
                logger.error(f"{GACHA_TYPE[gachaType]} 第 {page} 页解析失败！")
                await asyncsleep(2)
                continue
    return logsList

async def getAllTypeLogs(logUrl: str) -> Dict:
    """
    全部抽卡记录获取，可根据输入的抽卡记录链接获取 6 个月内全部抽卡记录

    * ``param logUrl: str`` 抽卡记录链接
    - ``return: Dict`` 全部抽卡记录，格式为``{"msg": "uid 或错误信息", "logs": {}]``
    """

    start, uidGot, newLogs = time(), "", {}
    # 获取最新抽卡记录
    for banner in GACHA_TYPE:
        remoteLogs = await getSingleTypeLogs(logUrl, banner)
        if remoteLogs:
            newLogs[banner] = remoteLogs
            if not uidGot:
                # 从记录中获取一次 UID
                uidGot = newLogs[banner][0]["uid"]
    msg = str(uidGot) if uidGot else "获取最新抽卡记录失败！"
    logger.info(
        "刷新 UID{} 的抽卡记录耗时 {}s".format(
            (uidGot if uidGot else "未知"), round(time() - start, 3)  # type: ignore
        )
    )
    return {"msg": msg, "logs": newLogs}

async def mergeLogs(locLogs: Dict, newLogs: Dict, config: Dict, qq: str) -> Dict:
    """
    本地记录数据与新增记录数据合并

    * ``param locLogs: Dict`` 本地记录数据
    * ``param newLogs: Dict`` 新增记录数据
    * ``param config: Dict`` 配置文件
    * ``param qq: str`` 目标 QQ
    - ``return: Dict`` 合并后的记录数据，格式为``{"uid": "123456789", "msg": "文字消息", "logs": {}]``
    """

    msgList, logs = [], {}  # 消息列表、待写入记录
    for banner in GACHA_TYPE:
        locItems, newItems = locLogs.get(banner, []), newLogs.get(banner, [])
        # 本地记录同步至待写入记录
        if locItems:
            logs[banner] = locItems
        # 本地记录与最新记录相同，跳过
        if locItems == newItems:
            continue
        # UID 不同，立即单独展示最新记录
        if len(locItems) and len(newItems) and locItems[0]["uid"] != newItems[0]["uid"]:
            logger.info(
                "{}({}) 中发现 UID 不同（{}!={}），跳过合并".format(
                    GACHA_TYPE[banner], banner, locItems[0]["uid"], newItems[0]["uid"]
                )
            )
            return {
                "uid": newItems[0]["uid"],
                "msg": "新增与缓存数据 UID 不一致\n跳过合并，单独展示新增记录..",
                "logs": newLogs,
            }
        # 本地记录与最新记录对比合并
        tempList, locItemsSimp = [], [[i["time"], i["name"]] for i in locItems]
        for i in range(len(newItems)):
            newItem = [newItems[i]["time"], newItems[i]["name"]]
            if newItem not in locItemsSimp:
                # 新增记录暂存
                tempList.append(newItems[i])
        if len(tempList):
            # 新增记录同步至待写入记录，保证新增数据在最前
            msgList.append(f"新增 {len(tempList)} 条{GACHA_TYPE[banner]}记录..")
            tempList.extend(locItems)
            logs[banner] = tempList
    # 无记录跳过后续缓存更改
    if not logs:
        return {"msg": "未发现任何记录.."}
    # 保存合并后的记录数据
    uid = logs[list(logs.keys())[0]][0]["uid"]
    cache = LOCAL_DIR / f"gachalogs-{uid}.json"
    config.update(
        {
            # "url": initUrl,
            # "cookie": args,
            "logs": str(cache),
            "time": int(time()),
            # "game_biz": role["game_biz"],
            "game_uid": uid,
            # "region": role["region"],
        }
    )
    res, _ = await logsHelper(cache, logs)
    if not res.isdigit():
        msgList.append(res)
    res = await configHelper(qq, config)
    if res.get("error"):
        msgList.append(res["error"])
    return {"uid": uid, "msg": "\n".join(msgList), "logs": logs}

async def getFullGachaLogs(config: Dict, qq: str, force: bool) -> Dict:
    """
    抽卡记录数据获取入口，可根据配置获取最新完整抽卡记录

    * ``param config: Dict`` 配置数据
    * ``param qq: str`` 目标 QQ
    * ``param force: bool`` 是否强制更新抽卡记录
    - ``return: Dict`` 抽卡记录数据，格式为``{"uid": "123456789", "msg": "文字消息", "logs": {}]``
    """

    # 读取抽卡记录缓存
    uid, locLogs = await logsHelper(config["logs"]) if config["logs"] else ("无记录", {})
    # msg = uid if not uid.isdigit() else ""
    if (not config["url"]) or (  # 缺少有效链接
        (not force) and (int(time()) - config["time"] < EXPIRE_SEC)  # 有缓存、未要求强制更新
    ):
        timeStr = datetime_with_tz(config["time"]).strftime("%m-%d %H:%M")
        timeTip = f"这是 {timeStr} 创建的缓存.." if config["time"] else "暂无本地记录！"
        logger.info(f"返回 QQ{qq} 于 {timeStr} 生成的抽卡记录缓存")
        return {"uid": uid, "msg": timeTip, "logs": locLogs}
    # 获取最新抽卡记录
    newLogsRes = await getAllTypeLogs(config["url"])
    if not str(newLogsRes["msg"]).isdigit():
        return {"msg": newLogsRes["msg"]}
    # 合并数据
    return await mergeLogs(locLogs, newLogsRes["logs"], config, qq)


from email import header
from email.policy import default
from textwrap import indent

#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File    :   evaluate_realtions.py
@Time    :   2022/03/08 10:51:08
@Author  :   Xixin Hu
@Version :   1.0
@Contact :   xixinhu97@foxmail.com
@Desc    :   None
'''

# here put the import lib

import json
from typing import DefaultDict
from pkg_resources import run_script

from tqdm import tqdm
import csv
import sys
import os
# print(sys.path)
sys.path.append(os.getcwd())
# sys.path.append("..") # CWQ dir

from components.utils import dump_json, load_json
import torch
from tqdm import tqdm
import pandas as pd

def gen_rel_match_data_set(split, filter_base=True):
    """make sentence-relation pair for cross-encoder training"""
    res_lines = []

    if filter_base:
        relation_data_file_name = f'data/rel_match/CWQ_{split}_relation_data_old.json'
    else:
        relation_data_file_name = f'data/rel_match/CWQ_{split}_relation_data.json'

    train_relation_map = load_json(f'data/label_maps/CWQ_{split}_label_maps.json')

    print(f'Loading {split}...')
    with open(relation_data_file_name,'r') as f:
        databank = json.load(f)
        print(len(databank))
        idx =0
        for qid,data in tqdm(databank.items(),desc=f'Processing {split}'):

            if split=='train': # for training
                for rel in data['positive_rels']:
                    if filter_base: # filter base relation
                        if rel.startswith('base.'):
                            continue

                    if '\t' in rel: # process illegal relation
                        rel = rel.split('\t')[0]
                    line = str(idx)+"\t"+data['question'] + '\t' + rel+'\t'+'1'
                    res_lines.append(line+'\n')
                    idx+=1
                for rel in data['negative_rels']:
                    if '\t' in rel:
                        rel = rel.split('\t')[0]
                    line = str(idx)+"\t"+data['question'] + '\t' + rel+'\t'+'0'
                    res_lines.append(line+'\n')
                    idx+=1
            else:
                added_rels = set()
                for rel in data['cand_rels']: # take cand_rels
                    if filter_base: # filter base relation
                        if rel.startswith('base.'):
                            continue

                    if '\t' in rel: # process illegal relation
                        rel = rel.split('\t')[0]
                    added_rels.add(rel)
                    if rel in data['positive_rels']:
                        line = str(idx)+"\t"+data['question'] + '\t' + rel+'\t'+'1'
                    else:
                        line = str(idx)+"\t"+data['question'] + '\t' + rel+'\t'+'0'
                    res_lines.append(line+'\n')
                    idx+=1

                for rel in train_relation_map.keys(): # add rels in train set
                    if rel not in added_rels:
                        added_rels.add(rel)
                        if '\t' in rel: # process illegal relation
                            rel = rel.split('\t')[0]
                        if rel in data['positive_rels']:
                            line = str(idx)+"\t"+data['question'] + '\t' + rel+'\t'+'1'
                        else:
                            line = str(idx)+"\t"+data['question'] + '\t' + rel+'\t'+'0'
                        res_lines.append(line+'\n')
                        idx+=1

    with open(f'data/rel_match/CWQ_{split}_rel_match.tsv','w') as f:
        f.writelines(res_lines)

def make_pseudo_sorted_dataset(split):
    pseudo_sorted_data = {}
    rel_data_bank = load_json(f'data/rel_match/CWQ_{split}_relation_data.json')
    for qid,data in tqdm(rel_data_bank.items(),desc='Building pseudo cand rel set'):
        pos_rels = data['positive_rels']
        neg_rels = data['negative_rels']
        pseudo_sorted_data[qid] = pos_rels+neg_rels

    dump_json(pseudo_sorted_data,f'data/rel_match/CWQ_{split}_cand_rels_sorted.json',indent=4)

def make_sorted_dataset(split, from_cand=False):
    # with open('logits.pt','r') as f:

    logits = torch.load(f'data/rel_match/predict/{split}/logits.pt',map_location=torch.device('cpu'))

    # print(logits)
    # print(len(logits))
    logits_list = list(logits.squeeze().numpy())
    print('Logits len:',len(logits_list))

    split_tsv = f'CWQ_{split}_rel_match_old.tsv'

    # with open(split_tsv,'r') as f:
    #     tsv_lines = f.readlines()
    # print('TSV lines:', len(tsv_lines))

    split_df = pd.read_csv(split_tsv, header=None, delimiter='\t'
                            ,names=['id','question','rel','label']
                            ,error_bad_lines=True
                            ,quoting=csv.QUOTE_NONE)
                            #,dtype={"id":int, "question":str, "rel":str, 'label':int}

    print('Dataframe len:',len(split_df))

    # length check
    assert len(logits_list) == len(split_df)

    enumerate_cand_rel_bank  = load_json(f'data/rel_match/CWQ_{split}_relation_data_old.json')

    # with open('logits.txt','w') as f:
    #     for logit in tqdm(logits_list):
    #         f.write(str(logit)+'\n')

    split_dataset = f'data/origin/ComplexWebQuestions_{split}.json'
    if split=='test_new':
        split_dataset= 'data/origin/ComplexWebQuestions_test.json'
    split_dataset = load_json(split_dataset)
    question2id = {x['question']:x['ID'] for x in split_dataset}

    # cand_rel_bank = {} # Dict[Question, List[Relation:logit]]
    cand_rel_bank = DefaultDict(dict)
    cand_rel_bank_from_enumerate = {}
    for idx,logit in tqdm(enumerate(logits_list),total=len(logits_list),desc=f'Reading logits {split}'):
        logit = float(logit)
        question = split_df.loc[idx]['question']
        rel = split_df.loc[idx]['rel']
        cwq_id = question2id.get(question,None)

        if not cwq_id:
            continue
        else:
            if not from_cand:
                cand_rel_bank[cwq_id][rel]=logit
            else:
                # only retain the relations in cand rels
                if cwq_id in cand_rel_bank_from_enumerate:
                    cand_rel_enumreate = cand_rel_bank_from_enumerate[cwq_id]
                else:
                    cand_rel_enumreate = set(enumerate_cand_rel_bank[cwq_id]['cand_rels'])
                    cand_rel_bank_from_enumerate[cwq_id] = cand_rel_enumreate

                if rel in cand_rel_enumreate: # only retain the relations in cand rels
                    cand_rel_bank[cwq_id][rel]=logit

    cand_rel_file_name = f'data/rel_match/CWQ_{split}_cand_rels_from_cand.json' if from_cand else f'data/rel_match/CWQ_{split}_cand_rels.json'
    dump_json(cand_rel_bank,cand_rel_file_name,indent=4)

    final_candRel_map = {} # Dict[Question,List[Rel]]   sorted by logits

    for ori_data in tqdm(split_dataset,total=len(split_dataset),desc=f'Sort by logits {split}'):
        qid = ori_data['ID']
        cand_rel_map = cand_rel_bank.get(qid,None)
        if not cand_rel_map:
            final_candRel_map[qid]=[]
        else:
            cand_rel_list = list(cand_rel_map.keys())
            cand_rel_list.sort(key=lambda x:float(cand_rel_map[x]),reverse=True)
            final_candRel_map[qid]=cand_rel_list

    sorted_cand_rel_name = f'data/rel_match/CWQ_{split}_cand_rels_sorted.json' if not from_cand else f'data/rel_match/CWQ_{split}_cand_rels_sorted_from_cand.json'
    dump_json(final_candRel_map,sorted_cand_rel_name,indent=4)

def check_sorted_dataset_coverage(split,sorted_rel_file):

    sorted_rel_bank = load_json(sorted_rel_file)
    label_map_file = f'data/label_maps/CWQ_{split}_label_maps.json'
    if split=='test_new':
        label_map_file = f'data/label_maps/CWQ_test_label_maps.json'
    label_maps_bank = load_json(label_map_file)

    gt_0_pre = 0
    gt_0_recall = 0
    top1_pre = 0
    top1_recall = 0
    top3_pre = 0
    top3_recall = 0
    top4_pre = 0
    top4_recall = 0
    top5_pre = 0
    top5_recall = 0
    top10_pre = 0
    top10_recall = 0
    total = len(sorted_rel_bank)
    for qid in sorted_rel_bank:
        cand_rel_list = sorted_rel_bank[qid]
        gold_rel_set = label_maps_bank[qid]['rel_label_map'].keys()

        top1_cand = set([x[0] for x in cand_rel_list[:1]])
        top3_cand = set([x[0] for x in cand_rel_list[:3]])
        top4_cand = set([x[0] for x in cand_rel_list[:4]])
        top5_cand = set([x[0] for x in cand_rel_list[:5]])
        top10_cand = set([x[0] for x in cand_rel_list[:10]])

        gt_0_cand = set([x[0] for x in cand_rel_list if x[1]>=0])
        if not gt_0_cand:
            gt_0_cand = set([x[0] for x in cand_rel_list[:1]])

        gt_0_pre += len(gt_0_cand & gold_rel_set) / len(gt_0_cand) if gt_0_cand else 0
        gt_0_recall += len(gt_0_cand & gold_rel_set) / len(gold_rel_set) if gold_rel_set else 0
        gt_0_f1 = 2*gt_0_pre*gt_0_recall/(gt_0_pre+gt_0_recall) if (gt_0_pre+gt_0_recall)>0 else 0

        top1_pre += len(top1_cand & gold_rel_set) / len(top1_cand) if top1_cand else 0
        top1_recall += len(top1_cand & gold_rel_set) / len(gold_rel_set) if gold_rel_set else 0
        top1_f1 = 2*top1_pre*top1_recall/(top1_pre+top1_recall) if (top1_pre+top1_recall)>0 else 0

        top3_pre += len(top3_cand & gold_rel_set) / len(top3_cand) if top3_cand else 0
        top3_recall += len(top3_cand & gold_rel_set) / len(gold_rel_set) if gold_rel_set else 0
        top3_f1 = 2*top3_pre*top3_recall/(top3_pre+top3_recall) if (top3_pre+top3_recall)>0 else 0

        top4_pre += len(top4_cand & gold_rel_set) / len(top4_cand) if top4_cand else 0
        top4_recall += len(top4_cand & gold_rel_set) / len(gold_rel_set) if gold_rel_set else 0
        top4_f1 = 2*top4_pre*top4_recall/(top4_pre+top4_recall) if (top4_pre+top4_recall)>0 else 0

        top5_pre += len(top5_cand & gold_rel_set) / len(top5_cand) if top5_cand else 0
        top5_recall += len(top5_cand & gold_rel_set) / len(gold_rel_set) if gold_rel_set else 0
        top5_f1 = 2*top5_pre*top5_recall/(top5_pre+top5_recall) if (top5_pre+top5_recall)>0 else 0

        top10_pre += len(top10_cand & gold_rel_set) / len(top10_cand) if top10_cand else 0
        top10_recall += len(top10_cand & gold_rel_set) / len(gold_rel_set) if gold_rel_set else 0
        top10_f1 = 2*top10_pre*top10_recall/(top10_pre+top10_recall) if (top10_pre+top10_recall)>0 else 0

    print(split,'TOTAL: ',total)
    print(f'GT 0 pre: {gt_0_pre/total}, GT 0 recall: {gt_0_recall/total}, GT 0 f1: {gt_0_f1/total}')
    print(f'TOP 1 pre: {top1_pre/total}, TOP 1 recall: {top1_recall/total}, TOP 1 f1: {top1_f1/total}')
    print(f'TOP 3 pre: {top3_pre/total}, TOP 3 recall: {top3_recall/total}, TOP 3 f1: {top3_f1/total}')
    print(f'TOP 4 pre: {top4_pre/total}, TOP 4 recall: {top4_recall/total}, TOP 4 f1: {top4_f1/total}')
    print(f'TOP 5 pre: {top5_pre/total}, TOP 5 recall: {top5_recall/total}, TOP 5 f1: {top5_f1/total}')
    print(f'TOP 10 pre: {top10_pre/total}, TOP 10 recall: {top10_recall/total}, TOP 10 f1: {top10_f1/total}')

def check_cand_dataset_coverage(split):
    """Check the recall of candidate relations"""
    data_bank = load_json(f'data/rel_match/CWQ_{split}_relation_data_old.json')
    cand_recall = 0
    for qid,data in data_bank.items():
        positive_rels = data['positive_rels']
        negative_rels = data['negative_rels']
        cand_rels = data['cand_rels']

        cand_recall += len(set(cand_rels)& set(positive_rels))/len(set(positive_rels))

    print(f'{split} cand recall: {cand_recall/len(data_bank)}')

def make_sorted_dataset_from_logits(split, logits_file, tsv_file, output_dir):

    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    logits = torch.load(logits_file,map_location=torch.device('cpu'))

    # print(logits)
    # print(len(logits))
    logits_list = list(logits.squeeze().numpy())
    print('Logits len:',len(logits_list))

    tsv_df = pd.read_csv(tsv_file, delimiter='\t',dtype={"id":int, "question":str, "relation":str, 'label':int})
                            #quoting=csv.QUOTE_NONE,

    print('Tsv len:', len(tsv_df))
    # print(tsv_df.head())
    print('Question Num:',len(tsv_df['question'].unique()))

    assert(len(logits_list)==len(tsv_df))

    split_dataset = load_json(f'data/origin/ComplexWebQuestions_{split}.json')
    question2id = {x['question']:x['ID'] for x in split_dataset}

    # cand_rel_bank = {} # Dict[Question, List[Relation:logit]]
    cand_rel_bank = DefaultDict(dict)
    for idx,logit in tqdm(enumerate(logits_list),total=len(logits_list),desc=f'Reading logits {split}'):
        logit = float(logit)
        question = tsv_df.loc[idx]['question']
        rel = tsv_df.loc[idx]['relation'].split("|")[0]
        cwq_id = question2id.get(question,None)

        if not cwq_id:
            print(question)
            cand_rel_bank[cwq_id]= {}
        else:
            cand_rel_bank[cwq_id][rel]=logit

    cand_rel_logit_map = {}
    for qid in tqdm(cand_rel_bank,total=len(cand_rel_bank),desc='Sorting rels...'):
        cand_rel_maps = cand_rel_bank[qid]
        cand_rel_list = [(rel,logit) for rel,logit in cand_rel_maps.items()]
        cand_rel_list.sort(key=lambda x:x[1],reverse=True)

        cand_rel_logit_map[qid]=cand_rel_list

    dump_json(cand_rel_logit_map,os.path.join(output_dir,f'CWQ_{split}_cand_rel_logits.json'),indent=4)

    final_candRel_map = DefaultDict(list) # Dict[Question,List[Rel]]   sorted by logits

    for ori_data in tqdm(split_dataset,total=len(split_dataset),desc=f'{split} Dumping... '):
        qid = ori_data['ID']
        # cand_rel_map = cand_rel_bank.get(qid,None)
        cand_rel_list = cand_rel_logit_map.get(qid,None)
        if not cand_rel_list:
            final_candRel_map[qid]=[]
        else:
            # cand_rel_list = list(cand_rel_map.keys())
            # cand_rel_list.sort(key=lambda x:float(cand_rel_map[x]),reverse=True)
            final_candRel_map[qid]=[x[0] for x in cand_rel_list]

    sorted_cand_rel_name = os.path.join(output_dir,f'CWQ_{split}_cand_rels_sorted.json')
    dump_json(final_candRel_map,sorted_cand_rel_name,indent=4)

if __name__=='__main__':

    output_dir = 'data/rel_match/relations/sorted_results_new'
    # split = 'dev'
    for split in ['dev']:
        logits_file = f'data/rel_match/relations/predictions_new/{split}/logits.pt'
        tsv_file = f'data/rel_match/relations/biEncoder_top200_new/CWQ.{split}.biEncoder.train_all.richRelation.top200.tsv'
        make_sorted_dataset_from_logits(split, logits_file, tsv_file, output_dir)

    # split = 'test'
    for split in ['dev']:
        sorted_logit_file_name = os.path.join(output_dir,f'CWQ_{split}_cand_rel_logits.json')
        check_sorted_dataset_coverage(split,sorted_logit_file_name)

    # for split in ['train','test','dev']:
    #     # gen_rel_match_data_set(split)
    #     check_cand_dataset_coverage(split)

    # gen_rel_match_data_set('test')

    # make_pseudo_sorted_dataset('dev')
    # make_pseudo_sorted_dataset('train')
    # for split in ['test','dev','train']:
    #     make_sorted_dataset(split)

    # 根据logits重建排序后的候选关系
    # make_sorted_dataset('test_new')
    # make_sorted_dataset('test',from_cand=True)
    # make_sorted_dataset('dev',from_cand=True)
    # make_sorted_dataset('train',from_cand=True)

    # 检查排序后的候选关系的覆盖率
    # check_sorted_dataset_coverage('test')
    # check_sorted_dataset_coverage('dev')
    # check_sorted_dataset_coverage('train')
    # check_sorted_dataset_coverage('test_new')

    # check_sorted_dataset_coverage('dev')
    # check_sorted_dataset_coverage('train')


#1
from sklearn.datasets import load_iris

dataset = load_iris()
x = dataset.data
y = dataset.target
#print(dataset.DESCR)

#2
import numpy as np
#axis=0 压缩行，求各列的平均数，返回1*m； axis=1则相反；默认求所有值的平均数
attribute_means = x.mean(axis=0)
x_d = np.array(x>=attribute_means,dtype='int')

#3 OneR-1
from collections import defaultdict
from operator import  itemgetter
 # params依次：数据集；分类类别（数组）（诸如男，女）；选好的特征值索引（特征值列表排序后的索引）；给定特征值（如，长度宽度）
 #  训练特征值目的：[一个给定特征+值情况下]1.得到最可能（最高频）的类别以及2.得到错误率（次数）
def train_feature_value(x,classifications,feature_index,feature_value):
    #分类类别记数
    class_counts = defaultdict(int)
    #整个数据集中取每个【个体】以及【个体的类别】，若【个体的特征值】符合【给定的特征值】,【这个类别】符合条件记数+1-->即是找出所有【符合给定特征值】的类别记数情况
    for sample,classification in zip(x,classifications):
        if sample[feature_index] == feature_value:
            class_counts[classification] += 1

    #对【符合给定特征值】类别记数情况进行[降序]排序并取[最大记数]的一种情况下的【类别】
    sorted_class_counts = sorted(class_counts.items(),key=itemgetter(1),reverse = True)
    most_frequent_class = sorted_class_counts[0][0]

    #计算这条规则下的错误率 表示的是分类不适用的个体数量 这里class_value 相当于classification
    # 如果类别不是最高出现频率类别则 取出来这个类别的记数，组成列表
    incorrect_predictions = [class_count for class_value,class_count
                             in class_counts.items()
                             if class_value != most_frequent_class]
    error = sum(incorrect_predictions)
    return most_frequent_class,error

#4 OneR-2
 #对于某项特征(feature_index)，遍历其每一个特征值(feature_value)，使用train_feature_value函数(得到最高频类别，和这个特征值出现错误的次数)
 #  定义函数把所有错误次数加起来，得到的是该特征下的总错误次数
def train_on_feature(x,classifications,feature_index):
    #set 是一个只有key的集合,同数学意义
    #返回 feature_index 所指的列，如feature_index == 1，则返回第一列所组成的数组,最后转化为集合(去重)
    feature_values = set(x[:,feature_index])
    #预测器{'特征值':类别}；errors列表
    predictors = {}
    errors = []
    for current_value in feature_values:
        most_frequent_class,error = train_feature_value(x,classifications,feature_index,current_value)
        predictors[current_value] = most_frequent_class
        errors.append(error)
    total_error = sum(errors)
    return predictors,total_error

#5 OneR-3 -test
from sklearn.cross_validation import train_test_split
 #将数据集分成两部分训练及测试集，防止过拟合，random_state指定切分随机状况；每次切分相同随机状况，切分结果相同，值为none每次切分结果真正随机。
Xd_train, Xd_test, y_train, y_test = train_test_split(x_d,y,random_state=14)
all_predictors = {}
errors = {}
# shape 返回数组维度（降序），shape[0]最高维;矩阵则对应行[0]，列[1]
for feature_index in range(Xd_train.shape[1]):
    predictors, total_error = train_on_feature(Xd_train,y_train,feature_index)
    print()
    all_predictors[feature_index] = predictors
    errors[feature_index] = total_error
#从所有特征中 找出 [错误率]最低的特征，以此为依据作为分类唯一规则-->得到最佳特征
best_feature,best_error = sorted(errors.items(),key=itemgetter(1))[0]
#最佳特征-->最佳特征值,创建model模型
model = {'feature':best_feature,
         'predictor':all_predictors[best_feature]}
#?????????
print(model)
def predict(x_test,model):
    variable = model['feature']
    predictor = model['predictor']
    y_predicted = np.array([predictor[int(sample[variable])] for sample in x_test])
    return y_predicted

y_predicted = predict(Xd_test,model)
accuracy = np.mean(y_predicted == y_test) *100
print('The test accuracy is {:.1f}%'.format(accuracy))

#示例上
#print(all_predictors)
#{0: ({0: 0, 1: 2}, 41), 1: ({0: 1, 1: 0}, 58), 2: ({0: 0, 1: 2}, 37), 3: ({0: 0, 1: 2}, 37)}
#示例上的 all_predictors 中 得到的字典值为train_on_feature的两个值，所以在建立模型中取第一个[0]第二个是总错误率
#自己 -->缺少了错误个数
#print(all_predictors)
#{feature(long?wide?):{feature_value(how long?how wide?):class(what kind?)},(mayInclude mistakes)}
#{0: {0: 0, 1: 2}, 1: {0: 1, 1: 0}, 2: {0: 0, 1: 2}, 3: {0: 0, 1: 2}}
#https://blog.csdn.net/leafage_m/article/details/79046435


from numpyro import sample, plate
import numpyro.distributions as dists
from numpyro.infer.autoguide import AutoDiagonalNormal
import jax.numpy as jnp
import numpy as np
import pandas as pd
from twinify.mixture_model import MixtureModel
from twinify.model_loading import TModelFunction, TGuideFunction
import argparse
from typing import Iterable, Tuple
from collections import OrderedDict
from functools import partial

#ordinal_features = ["dAge", "iFertil", "iEnglish", "iImmigr", "iYearsch"]
categorical_features = ["dAncstry1", "dAncstry2", "iMarital"] + ["dAge", "iFertil", "iEnglish", "iImmigr", "iYearsch"]
binary_features = ["iSex", "iVietnam", "iKorean", "iMobillim"]
target = ["dPoverty"]

#all_features = ordinal_features + categorical_features + binary_features
all_features = categorical_features + binary_features

reference_groups = OrderedDict(
        dAncstry1  = 1,
        dAncstry2  = 1,
        iMarital   = 0,
        dAge       = 2,
        iFertil    = 0,
        iEnglish   = 0,
        iImmigr    = 0,
        iYearsch   = 1

    )

def preprocess(ori_df):
    df = pd.DataFrame()
    for feature in all_features:
        df[feature] = ori_df[feature].astype('category').cat.codes
    df["dPoverty"] = (ori_df.dPoverty.astype(int) < 2).map({False: 0, True: 1})
    y = df.pop('dPoverty')
    X = df[all_features]
    num_data = len(X)
    return (X, y), num_data

def postprocess(posterior_samples, ori_df):
    syn_X = posterior_samples['X']
    syn_y = posterior_samples['y']
    syn_data = jnp.hstack([syn_X, syn_y[:, np.newaxis]])
    columns = all_features.copy()
    columns.remove('dPoverty')
    columns.append('dPoverty')
    syn_df = pd.DataFrame(syn_data, columns = columns)

    encoded_syn_df = syn_df.copy()
    for feature in all_features:
        if feature != 'dPoverty':
            category_map = ori_df[feature].astype('category').cat.categories
            encoded_syn_df[feature] = pd.Categorical.from_codes(encoded_syn_df[feature], category_map)
    encoded_syn_df['dPoverty'] = encoded_syn_df['dPoverty'].astype(bool)

    return syn_df, encoded_syn_df

def onehot_encode(x, num_values):
    lut = jnp.eye(num_values)
    return lut[x]

def adjusted_onehot_encode(x, num_values, reference_idx):
    encoded = onehot_encode(x, num_values)
    return jnp.hstack((encoded[:, :reference_idx], encoded[:, (reference_idx+1):]))

def model_factory(twinify_args: argparse.Namespace, unparsed_args: Iterable[str], ori_df: pd.DataFrame) -> Tuple[TModelFunction, TGuideFunction]:
    model_args_parser = argparse.ArgumentParser("Model 1")
    model_args_parser.add_argument('--init_scale', type=float, default=0.1, help='Initial value for scales in variational AutoDiagonalNormal.')
    args = model_args_parser.parse_args(unparsed_args, twinify_args)

    k = args.k

    # collect number of categories and index of reference group/category for each feature
    d_sizes = OrderedDict()
    reference_indices = OrderedDict()
    for feature, reference_category in reference_groups.items():
        idx = list(ori_df[feature].astype('category').cat.categories).index(reference_category)
        reference_indices[feature] = idx
        d_sizes[feature] = len(ori_df[feature].astype('category').cat.categories)

    #for feature in ordinal_features:
    #    d_sizes[feature] = len(ori_df[feature].astype('category').cat.categories)

    def model(X=None, y=None, num_obs_total=None):
        d = len(all_features)
        batch_size = 1

        if y is not None:
            assert X is not None
            batch_size = X.shape[0]
            assert y.shape[0] == batch_size
            assert d == X.shape[1]
        else:
            assert X is None

        if num_obs_total is None:
            num_obs_total = batch_size

        ############## PRIORs ##################

        # Prior for mixture model of regressors
        prior_samples = []
        probs_for_mm = []
        for feature in all_features:
            if feature not in binary_features:
                prior_samples.append(sample(f'{feature}_probs', dists.Dirichlet(jnp.ones((k, d_sizes[feature])))))
                probs_for_mm.append(dists.Categorical(probs=prior_samples[-1]))
            else:
                prior_samples.append(sample(f'{feature}_probs', dists.Beta(jnp.ones((k,)), jnp.ones((k,)))))
                probs_for_mm.append(dists.Bernoulli(probs=prior_samples[-1]))

        pis = sample('pis_probs', dists.Dirichlet(jnp.ones(k)))

        mixture_dist = MixtureModel(probs_for_mm, pis)

        # Set up Poisson regression: this sets the prior for the regression weights
        # and also prepares functions that convert each categorical into (adjusted) one-hot vectors
        coefs = []
        feature_encoders = []
        for feature in all_features:
            if feature in categorical_features:
                adjusted_d_size = d_sizes[feature]- 1
                feature_coefs = sample(f"weight_{feature}", dists.Normal(jnp.zeros(adjusted_d_size), jnp.ones(adjusted_d_size)))
                coefs.append(feature_coefs)
                feature_encoders.append(partial(adjusted_onehot_encode, num_values=d_sizes[feature], reference_idx=reference_indices[feature]))

            else:
                feature_coefs = sample(f"weight_{feature}", dists.Normal(0.0, 1.0), sample_shape=(1,))
                feature_encoders.append(lambda x: x[:, jnp.newaxis])
                coefs.append(feature_coefs)

        weight = jnp.concatenate(coefs)

        intercept = sample("intercept", dists.Normal(0., 1.))

        ############# LIKELIHOOD ####################

        with plate('batch', num_obs_total, batch_size):
            X = sample('X', mixture_dist, obs=X)

            encoded_X = jnp.hstack([
                encode(X[:, i]) for i, encode in enumerate(feature_encoders)
            ])

            logit = encoded_X @ weight + intercept
            y = sample('y', dists.Poisson(rate=jnp.exp(logit)), obs=y)

    guide = AutoDiagonalNormal(model, init_scale=args.init_scale)

    return model, guide


import asyncio
from asyncio.locks import Event
from contextlib import contextmanager

from replayserver.send.stream import DelayedReplayStream
from replayserver.errors import MalformedDataError, \
    CannotAcceptConnectionError
from replayserver.collections import AsyncCounter

class Sender:
    def __init__(self, delayed_stream):
        self._stream = delayed_stream
        self._conn_count = AsyncCounter()
        self._ended = Event()
        asyncio.ensure_future(self._lifetime())

    @classmethod
    def build(cls, stream, **kwargs):
        delayed_stream = DelayedReplayStream.build(stream, **kwargs)
        return cls(delayed_stream)

    @contextmanager
    def _connection_count(self):
        self._conn_count.inc()
        try:
            yield
        finally:
            self._conn_count.dec()

    async def handle_connection(self, connection):
        if self._stream.ended():
            raise CannotAcceptConnectionError(
                "Reader connection arrived after replay ended")
        with self._connection_count():
            await self._write_header(connection)
            await self._write_replay(connection)

    async def _write_header(self, connection):
        header = await self._stream.wait_for_header()
        if header is None:
            raise MalformedDataError("Malformed replay header")
        await connection.write(header.data)

    async def _write_replay(self, connection):
        position = 0
        while True:
            data = await self._stream.wait_for_data(position)
            if not data:
                break
            position += len(data)
            conn_open = await connection.write(data)
            if not conn_open:
                break

    def close(self):
        pass

    async def _lifetime(self):
        await self._stream.wait_for_ended()
        await self._conn_count.wait_until_empty()
        self._ended.set()

    async def wait_for_ended(self):
        await self._ended.wait()


import argparse
import os
import copy
import gradio as gr

import numpy as np
import torch
from PIL import Image, ImageDraw, ImageFont

# Grounding DINO
import GroundingDINO.groundingdino.datasets.transforms as T
from GroundingDINO.groundingdino.models import build_model
from GroundingDINO.groundingdino.util import box_ops
from GroundingDINO.groundingdino.util.slconfig import SLConfig
from GroundingDINO.groundingdino.util.utils import clean_state_dict, get_phrases_from_posmap

# segment anything
from segment_anything import build_sam, SamPredictor
import cv2
import numpy as np
import matplotlib.pyplot as plt

# diffusers
import PIL
import requests
import torch
from io import BytesIO
from stable_diffusion_masked_diffedit import StableDiffusionMaskedDiffeditPipeline
from diffusers import DDIMScheduler

# chatgpt
from chatgpt import call_chatgpt

# blip2
from transformers import AutoProcessor, Blip2ForConditionalGeneration

def load_image(image_path):
    # load image
    image_pil = Image.open(image_path).convert("RGB")  # load image

    transform = T.Compose(
        [
            T.RandomResize([800], max_size=1333),
            T.ToTensor(),
            T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        ]
    )
    image, _ = transform(image_pil, None)  # 3, h, w
    return image_pil, image

def load_model(model_config_path, model_checkpoint_path, device):
    args = SLConfig.fromfile(model_config_path)
    args.device = device
    model = build_model(args)
    checkpoint = torch.load(model_checkpoint_path, map_location="cpu")
    load_res = model.load_state_dict(clean_state_dict(checkpoint["model"]), strict=False)
    print(load_res)
    _ = model.eval()
    return model

def get_grounding_output(model, image, caption, box_threshold, text_threshold, with_logits=True, device="cpu"):
    caption = caption.lower()
    caption = caption.strip()
    if not caption.endswith("."):
        caption = caption + "."
    model = model.to(device)
    image = image.to(device)
    with torch.no_grad():
        outputs = model(image[None], captions=[caption])
    logits = outputs["pred_logits"].cpu().sigmoid()[0]  # (nq, 256)
    boxes = outputs["pred_boxes"].cpu()[0]  # (nq, 4)
    logits.shape[0]

    # filter output
    logits_filt = logits.clone()
    boxes_filt = boxes.clone()
    filt_mask = logits_filt.max(dim=1)[0] > box_threshold
    logits_filt = logits_filt[filt_mask]  # num_filt, 256
    boxes_filt = boxes_filt[filt_mask]  # num_filt, 4
    logits_filt.shape[0]

    # get phrase
    tokenlizer = model.tokenizer
    tokenized = tokenlizer(caption)
    # build pred
    pred_phrases = []
    for logit, box in zip(logits_filt, boxes_filt):
        pred_phrase = get_phrases_from_posmap(logit > text_threshold, tokenized, tokenlizer)
        if with_logits:
            pred_phrases.append(pred_phrase + f"({str(logit.max().item())[:4]})")
        else:
            pred_phrases.append(pred_phrase)

    return boxes_filt, pred_phrases

def show_mask(mask, ax, random_color=False):
    if random_color:
        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
    else:
        color = np.array([30/255, 144/255, 255/255, 0.6])
    h, w = mask.shape[-2:]
    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
    ax.imshow(mask_image)

def show_box(box, ax, label):
    x0, y0 = box[0], box[1]
    w, h = box[2] - box[0], box[3] - box[1]
    ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0,0,0,0), lw=2))
    ax.text(x0, y0, label)

def adjust_encoding_ratio(instructions, feedbacks, cur_encoding_ratio):
    full_instructions = []
    for instruction in instructions:
        full_instruction = '{}. The current encoding ratio is {}.'.format(instruction, cur_encoding_ratio)
        full_instructions.append(full_instruction)

    encoding_ratio = call_chatgpt(full_instructions, feedbacks=feedbacks)

    return encoding_ratio

def maskedit(image_path, user_instructions, is_blip2_description, encoding_ratio, mask_mode):

    # cfg
    config_file = 'GroundingDINO/groundingdino/config/GroundingDINO_SwinT_OGC.py'  # change the path of the model config file
    grounded_checkpoint = 'GroundingDINO/weights/groundingdino_swint_ogc.pth'  # change the path of the model
    sam_checkpoint = 'sam_vit_h_4b8939.pth'
    # image_path = args.input_image
    # user_instructions = args.user_instructions
    # is_blip2_description = args.is_blip2_description
    box_threshold = 0.3
    text_threshold = 0.25
    # encoding_ratio = args.encoding_ratio
    device = 'cuda' if torch.cuda.is_available() else 'cpu'

    # load image
    image_pil, image = load_image(image_path)

    if is_blip2_description:
        blip_processor = AutoProcessor.from_pretrained('Salesforce/blip2-flan-t5-xl')
        blip_model = Blip2ForConditionalGeneration.from_pretrained('Salesforce/blip2-flan-t5-xl', torch_dtype=torch.float16).to(device)

        blip_prompt_1 = "Is this a photo, a painting, a drawing, or other kind of arts?"
        blip_inputs = blip_processor(image_pil, text=blip_prompt_1, return_tensors="pt").to(device, torch.float16)
        generated_ids = blip_model.generate(**blip_inputs, max_new_tokens=20)
        generated_text = blip_processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
        blip_prompt_2 = "{} of".format(generated_text.capitalize())
        blip_inputs = blip_processor(image_pil, text=blip_prompt_2, return_tensors="pt").to(device, torch.float16)
        generated_ids = blip_model.generate(**blip_inputs, max_new_tokens=20)
        description = blip_processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
        description = "{} {}.".format(blip_prompt_2, description)
        descriptions = [description]
        print('Blip2 description: {}'.format(description))
    else:
        descriptions = None
    instructions = [user_instructions]
    det_prompt, prompt_inversion, prompt, _ = call_chatgpt(instructions, descriptions=descriptions)
    print('Segmentation prompt: {}'.format(det_prompt))
    print('Inverted prompt: {}'.format(prompt_inversion))
    print('Editing prompt: {}'.format(prompt))

    # load model
    model = load_model(config_file, grounded_checkpoint, device=device)

    # run grounding dino model
    boxes_filt, pred_phrases = get_grounding_output(
        model, image, det_prompt, box_threshold, text_threshold, device=device
    )

    # initialize SAM
    predictor = SamPredictor(build_sam(checkpoint=sam_checkpoint).to(device))
    image = cv2.imread(image_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    predictor.set_image(image)

    size = image_pil.size
    H, W = size[1], size[0]
    for i in range(boxes_filt.size(0)):
        boxes_filt[i] = boxes_filt[i] * torch.Tensor([W, H, W, H])
        boxes_filt[i][:2] -= boxes_filt[i][2:] / 2
        boxes_filt[i][2:] += boxes_filt[i][:2]

    boxes_filt = boxes_filt.cpu()
    transformed_boxes = predictor.transform.apply_boxes_torch(boxes_filt, image.shape[:2]).to(device)

    masks, _, _ = predictor.predict_torch(
        point_coords = None,
        point_labels = None,
        boxes = transformed_boxes.to(device),
        multimask_output = False,
    )
    # masks: [1, 1, 512, 512]

    # Masked diffedit pipeline
    if mask_mode == 'merge':
        masks = torch.sum(masks, dim=0).unsqueeze(0)
        masks = torch.where(masks > 0, True, False)

    mask = masks[0][0].cpu().numpy() # simply choose the first mask, which will be refine in the future release
    mask_pil = Image.fromarray(mask)
    image_pil = Image.fromarray(image)

    scheduler =DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False,
                                        set_alpha_to_one=False)
    pipe = StableDiffusionMaskedDiffeditPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", scheduler=scheduler)
    pipe.to(device)

    image_pil = image_pil.resize((512, 512))
    mask_pil = mask_pil.resize((512, 512))
    image = pipe(prompt, image=image_pil, mask_image=mask_pil, prompt_inversion=prompt_inversion, encoding_ratio=encoding_ratio).images[0]
    image = image.resize(size)

    return image

with gr.Blocks() as demo:
    with gr.Row():
        gr.Markdown("""# Instruct Diffedit: to edit an image based on a user instruction.
                No mask or prompting is needed.""")
    with gr.Row():
        with gr.Column():
            image = gr.Image(type='filepath', label='Input image', value='./example_input/cat_dog.jpg')
            user_instructions = gr.Textbox(lines=3, label='User instructions', value='Change the cat to a fox')
            is_blip2_description = gr.Checkbox(label='Use Blip2 description')
            encoding_ratio = gr.Slider(label='Encoding ratio', minimum=0.0, maximum=1.0, step=0.05, value=0.5)
            mask_mode = gr.Dropdown(label='Mask mode', choices=['merge', 'max'], value='max')

        with gr.Column():
            output_image = gr.Image(type='pil', label='Output image')
            edit_btn = gr.Button('Edit')
            edit_btn.click(fn=maskedit, inputs=[image, user_instructions, is_blip2_description, encoding_ratio, mask_mode], outputs=[output_image], api_name='maskedit')

demo.queue().launch(share=True)


import unittest
import os.path as op
import numpy as np
import six
from scipy.signal import argrelmax, argrelmin, resample
from scipy.io import loadmat
from numpy.testing import assert_allclose
from pyhht.emd import EMD

class TestEMD(unittest.TestCase):

    def setUp(self):
        self.trend = self.ts = np.linspace(0, 1, 10000)
        self.mode1 = np.sin(2 * np.pi * 5 * self.ts)
        self.mode2 = np.sin(2 * np.pi * 10 * self.ts)

    def test_zeromean_decomposition(self):
        """Test if the EMD decomposes a signal properly if it has local mean
        close to zero."""
        t = 2 * np.pi * self.ts
        signal = np.cos(80 * t) + 0.8 * np.sin(50 * t) + \
            0.6 * np.sin(25 * t) + 0.4 * np.sin(10 * t) + 0.3 * np.cos(3 * t)
        emd = EMD(signal)
        imfs = emd.decompose()
        self.assertEqual(imfs.shape[0], 6)

    def test_emd_multidimensional_signal_error(self):
        """Check if EMD raises an error for multidimensional signals."""
        signal = self.trend + self.mode1 + self.mode2
        signal = np.c_[signal, signal]
        self.assertRaises(ValueError, EMD, signal)

    def test_imfs_total_no_error(self):
        """
        Check if the sum of the IMFs is sufficiently close to the input signal.
        """
        signal = np.sum([self.trend, self.mode1, self.mode2], axis=0)
        emd = EMD(signal)
        imfs = emd.decompose()
        assert_allclose(imfs.sum(0), signal)

    def test_decomposition(self):
        """Test the decompose method of the emd class."""
        signal = np.sum([self.trend, self.mode1, self.mode2], axis=0)
        decomposer = EMD(signal, t=self.ts)
        imfs = decomposer.decompose()
        if six.PY3:
            self.assertTupleEqual(imfs.shape, (3, signal.shape[0]))
        else:
            self.assertItemsEqual(imfs.shape, (3, signal.shape[0]))

    def test_noisy_signal(self):
        """Test if decomposing a noisy signal works."""
        fpath = op.join(op.abspath(op.dirname(__file__)), "testdata",
                        "gabor.mat")
        signal = loadmat(fpath)['gabor'].ravel()
        signal = resample(signal, signal.shape[0] * 100)
        signal += np.random.normal(size=signal.shape)
        engine = EMD(signal)
        engine.decompose()

    def test_maxiter(self):
        """Check if the maxiter parameter is respected."""
        fpath = op.join(op.abspath(op.dirname(__file__)), "testdata",
                        "gabor.mat")
        signal = loadmat(fpath)['gabor'].ravel()
        signal = resample(signal, signal.shape[0] * 100)
        signal += np.random.normal(size=signal.shape)
        engine = EMD(signal, maxiter=200)
        engine.decompose()
        self.assertLessEqual(engine.nbit, 200)

    def test_fixe(self):
        """Check if the fixe parameter is respected."""
        fpath = op.join(op.abspath(op.dirname(__file__)), "testdata",
                        "gabor.mat")
        signal = loadmat(fpath)['gabor'].ravel()
        signal = resample(signal, signal.shape[0] * 100)
        signal += np.random.normal(size=signal.shape)
        engine = EMD(signal, fixe=100, n_imfs=3)
        engine.decompose()
        self.assertListEqual(engine.nbits, [100] * 3)

    def test_residue(self):
        """Test the residue of the emd output."""
        signal = np.sum([self.trend, self.mode1, self.mode2], axis=0)
        decomposer = EMD(signal, t=self.ts)
        imfs = decomposer.decompose()
        n_imfs = imfs.shape[0]
        n_maxima = argrelmax(imfs[n_imfs - 1, :])[0].shape[0]
        n_minima = argrelmin(imfs[n_imfs - 1, :])[0].shape[0]
        self.assertTrue(max(n_maxima, n_minima) <= 2)

    def test_optional_nbsym(self):
        """Test if the optional nbsym parameter works when set to False. See
        jaidevd/pyhht#29."""
        fpath = op.join(op.abspath(op.dirname(__file__)), "testdata",
                        "issue29.npy")
        x = np.load(fpath)
        emd = EMD(x, nbsym=False)
        imfs = emd.decompose()
        self.assertEqual(imfs.shape[0], 4)
        np.testing.assert_allclose(x, imfs.sum(0))

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


from django.shortcuts import get_object_or_404
from django.shortcuts import render, redirect
from django.template.response import TemplateResponse
from django.views.generic import TemplateView
from django_tables2 import RequestConfig
from guardian.decorators import permission_required

from .forms import AddressForm, PersonForm, PhoneForm, EmailForm
from .models import Address, Person, Email, Phone
from .tables import ListBookTable

#
# class ListBookView(TemplateView):
#
#     def get(self, request, *args, **kwargs):
#         table = ListBookTable(Person.objects.filter(user=request.user), order_by='-name')
#         RequestConfig(request, paginate={'per_page': 10}).configure(table)
#
#         return render(request, 'template_list.html', {
#             'table': table, 'title': 'List Persons', 'link_add': 'person_new'
#         })
#
# list_book = permission_required('address_book.view_person', return_403=True)(ListBookView.as_view())

# second method to list book view
@permission_required('address_book.view_person', return_403=True)
def list_book(request):
    table = ListBookTable(Person.objects.filter(user=request.user), order_by='-name')
    RequestConfig(request, paginate={'per_page': 10}).configure(table)

    return render(request, 'template_list.html', {
        'table': table, 'title': 'List Persons', 'link_add': 'person_new'
    })

class PersonsView(TemplateView):
    def get(self, request):
        try:
            persons = Person.objects.filter(user=request.user).order_by('surname')
        except:
            persons = None
        column_name = ['id', 'Name', 'Groups', 'Detail']
        column_data = []
        for person in persons:
            group_list = []
            for group in person.group_set.all():
                group_link = '/group/group_detail/{}'.format(group.pk)
                group_list.append(
                    {'col': group.name,
                     'link': group_link
                     }
                )

            edit_list = [

                {'col': 'View', 'link': '/address_book/person/{}'.format(person.pk)},
                {'col': 'Edit', 'link': '/address_book/person_edit/{}'.format(person.pk)},
                {'col': 'Delete', 'link': '/address_book/person_del/{}'.format(person.pk)},

            ]

            row = [
                {'col': person.id},
                {'col': person.name + '' + person.surname},
                {'many_link': group_list},
                {'many_link': edit_list},
            ]
            column_data.append(row)

        ctx = {'title': 'Persons List ({})'.format(request.user),
               'link_add': '/address_book/person_new',
               'column_name': column_name,
               'column_data': column_data
               }
        return render(request=request, template_name="forms.html",
                      context=ctx)

class PersonsAllView(TemplateView):

    def get(self, request):
        column_name = []
        column_data = []
        try:
            persons = Person.objects.all()
        except:
            persons = None

        if persons:
            column_name = [
                'id',
                'Name',
                # 'View detail'
            ]
            for person in persons:
                row = [
                    {'col': person.id},
                    {'col': person.name + '' + person.surname},
                    # {'col': 'View', 'link': '/address_book/person/{}'.format(person.pk)},
                ]
                column_data.append(row)

        ctx = {'title': 'Persons all List',
               # 'link_add': '/address_book/person_new',
               'column_name': column_name,
               'column_data': column_data
               }
        return render(request=request, template_name="forms.html",
                      context=ctx)

        # return render(request=request, template_name="persons_all.html",
        #               context={'persons': persons})

class PersonView(TemplateView):

    def get(self, request, id):
        ctx = {}
        if id:
            person = get_object_or_404(Person, id=id, user=request.user)
            # check persmission if user doesnt have permission return "no permission"
            ctx['instance'] = person

        return TemplateResponse(request, 'person.html', ctx)

    def post(self, request): #TODO ?
        id_person = request.POST.get("id_person")
        id_address = request.POST.get("id_address")
        id_phone = request.POST.get("id_phone")
        id_email = request.POST.get("id_email")
        if id:
            instance_person = get_object_or_404(Person, id=id_person)
            instance_address = get_object_or_404(Address, id=id_address)
            instance_phone = get_object_or_404(Phone, id=id_phone)
            instance_email = get_object_or_404(Email, id=id_email)

            person = PersonForm(request.POST or None, instance=instance_person)
            address = AddressForm(request.POST or None, instance=instance_address)
            phone = PhoneForm(request.POST or None, instance=instance_phone)
            email = EmailForm(request.POST or None, instance=instance_email)

            ctx = {"form_person": person, "form_address": address,
                   "form_phone": phone, "form_email": email,
                   "title": "Person"}

            return TemplateResponse(request, 'person_new.html', ctx)

class PersonNewView(TemplateView):

    def get(self, request):
        person = PersonForm()
        address = AddressForm()

        ctx = {"form_person": person, "form_address": address,
               "title": "Person"}
        return TemplateResponse(request, 'person_new.html', ctx)

    def post(self, request):
        person_form = PersonForm(request.POST)
        address_form = AddressForm(request.POST)
        if person_form.is_valid() and address_form.is_valid():
            address = address_form.save(commit=True)
            person = person_form.save(commit=False)
            person.address_id = address.id
            person.user = request.user
            person.save()
            return redirect('person', id=person.id)

class PersonEditView(TemplateView):

    def get(self, request, id):
        if id:
            instance_person = get_object_or_404(Person, id=id)
            instance_address = get_object_or_404(Address, id=instance_person.address_id)

            person = PersonForm(request.POST or None, instance=instance_person)
            address = AddressForm(request.POST or None, instance=instance_address)

            ctx = {"form_person": person, "form_address": address,
                   "title": "Person", "id": id}

            return TemplateResponse(request, 'person_edit.html', ctx)

    def post(self, request, id):
        instance_person = get_object_or_404(Person, id=id)
        instance_address = get_object_or_404(Address, id=instance_person.address_id)

        person_form = PersonForm(request.POST or None, instance=instance_person)
        address_form = AddressForm(request.POST or None, instance=instance_address)

        if address_form.is_valid():
            address = address_form.save(commit=True)
        if person_form.is_valid():
            person = person_form.save(commit=False)
            person.address_id = address.id
            person.save()
            return redirect('person', id=person.id)

class PersonDelView(TemplateView):

    def get(self, request, id):
        """Method get delete person data in table address phone email and person"""
        if id:
            person = get_object_or_404(Person, id=id)
            # address = person.address
            #
            # address = get_object_or_404(Address, id=person.address_id)
            # phone = get_object_or_404(Phone, id=person.phone_id)
            # email = get_object_or_404(Email, id=person.email_id)
            #
            # if address:
            #     address.delete()
            # if phone:
            #     phone.delete()
            # if email:
            #     email.delete()
            if person:
                person.delete()

        return redirect('persons_list')


import os
from werkzeug.utils import secure_filename
import requests
from flask import Flask, request
from flask_cors import CORS

app = Flask(__name__)
CORS(app)
# uploads_dir = '/home/vest1/Desktop/multi-01-starting-setup/backend/data'
uploads_dir='/app/data'

# post both files at once
# files saved into the data directory
@app.route('/upload', methods=['POST'])
def upload_files():
    if request.method == 'POST':
        f1 = request.files['file1']
        f2 = request.files['file2']
        f1.save(os.path.join(uploads_dir, secure_filename(f1.filename)))
        f2.save(os.path.join(uploads_dir, secure_filename(f2.filename)))
        return 'file uploaded successfully\n'

@app.route('/checksum', methods=['GET'])
def determine_file_similarity():
    if request.method == 'GET':
        f1 = request.args.get('f1')
        f2 = request.args.get('f2')
        payload = {'f1': f1, 'f2': f2}
        response = requests.get("http://checksum:5001", params=payload)
        return response.json()

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


#!/usr/bin/python
# -*- coding: utf-8 -*-
#
# TalkTalkTalk
#
# is an easy-installable small chat room, with chat history.
#
# author:  Joseph Ernest (twitter: @JosephErnest)
# url:     http://github.com/josephernest/talktalktalk
# license: MIT license

import sys, json, bleach, time, threading, dumbdbm, random, re
import daemon
from bottle import route, run, view, request, post, ServerAdapter, get, static_file
from gevent import pywsgi
from geventwebsocket.handler import WebSocketHandler
from geventwebsocket.exceptions import WebSocketError
from collections import deque
from config import PORT, HOST, ADMINNAME, ADMINHIDDENNAME, ALLOWEDTAGS

idx = 0

def websocket(callback):
    def wrapper(*args, **kwargs):
        callback(request.environ.get('wsgi.websocket'), *args, **kwargs)
    return wrapper

class GeventWebSocketServer(ServerAdapter):
    def run(self, handler):
        server = pywsgi.WSGIServer((self.host, self.port), handler, handler_class=WebSocketHandler)
        server.serve_forever()

def main():
    global idx
    db = dumbdbm.open('talktalktalk.db', 'c')
    idx = len(db)

    users = {}
    pings = {}
    usermessagetimes = {}

    def send_userlist():
        for u in users.keys():
            if not u.closed:
                u.send(json.dumps({'type' : 'userlist', 'connected': users.values()}))

    def clean_username(usr, ws):
        username = bleach.clean(usr, tags=ALLOWEDTAGS, strip=True)
        #username = re.sub('[ :]', '', username)      # removes " ", ":", and the evil char "" http://unicode-table.com/fr/200D/
        username = re.sub(r'\W+', '', username)       # because of spam and usage of malicious utf8 characters, let's use alphanumeric usernames only for now
        username = username[:16]
        if username.lower() == ADMINNAME or username == '':
            username = 'user' + str(random.randint(0, 1000))
            ws.send(json.dumps({'type' : 'usernameunavailable', 'username' : username}))
        elif username.lower() == ADMINHIDDENNAME:
            username = ADMINNAME
            ws.send(json.dumps({'type' : 'displayeduser', 'username' : username}))
        return username

    def dbworker():        # when a user disappears during more than 30 seconds (+/- 10), remove him/her from the userlist
        while True:
            userlistchanged = False
            t = time.time()
            for ws in users.copy():
                if t - pings[ws] > 30:
                    del users[ws]
                    del pings[ws]
                    userlistchanged = True
            if userlistchanged:
                send_userlist()
            time.sleep(10)

    dbworkerThread = threading.Thread(target=dbworker)
    dbworkerThread.daemon = True
    dbworkerThread.start()

    @get('/ws', apply=[websocket])
    def chat(ws):
        global idx
        usermessagetimes[ws] = deque(maxlen=10)
        while True:
            try:
                receivedmsg = ws.receive()
                if receivedmsg is not None:

                    receivedmsg = receivedmsg.decode('utf8')
                    if len(receivedmsg) > 4096:      # this user is probably a spammer
                        ws.send(json.dumps({'type' : 'flood'}))
                        break

                    pings[ws] = time.time()

                    if receivedmsg == 'ping':         # ping/pong packet to make sure connection is still alive
                        ws.send('id' + str(idx-1))    # send the latest message id in return
                        if ws not in users:           # was deleted by dbworker
                            ws.send(json.dumps({'type' : 'username'}))
                    else:
                        usermessagetimes[ws].append(time.time())                           # flood control
                        if len(usermessagetimes[ws]) == usermessagetimes[ws].maxlen:
                            if usermessagetimes[ws][-1] - usermessagetimes[ws][0] < 5:     # if more than 10 messages in 5 seconds (including ping messages)
                                ws.send(json.dumps({'type' : 'flood'}))                    # disconnect the spammer
                                break

                        msg = json.loads(receivedmsg)

                        if msg['type'] == 'message':
                            message = (bleach.clean(msg['message'], tags=ALLOWEDTAGS, strip=True)).strip()

                            if ws not in users:         # is this really mandatory ?
                                username = clean_username(msg['username'], ws)
                                users[ws] = username
                                send_userlist()

                            if message:
                                if len(message) > 1000:
                                    message = message[:1000] + '...'
                                s = json.dumps({'type' : 'message', 'message': message, 'username': users[ws], 'id': idx, 'datetime': int(time.time())})
                                db[str(idx)] = s                # Neither dumbdbm nor shelve module allow integer as key... I'm still looking for a better solution!
                                idx += 1
                                for u in users.keys():
                                    u.send(s)

                        elif msg['type'] == 'messagesbefore':
                            idbefore = msg['id']
                            ws.send(json.dumps({'type' : 'messages', 'before': 1, 'messages': [db[str(i)] for i in range(max(0,idbefore - 100),idbefore)]}))

                        elif msg['type'] == 'messagesafter':
                            idafter = msg['id']
                            ws.send(json.dumps({'type' : 'messages', 'before': 0, 'messages': [db[str(i)] for i in range(idafter,idx)]}))

                        elif msg['type'] == 'username':
                            username = clean_username(msg['username'], ws)
                            if ws not in users:          # welcome new user
                                ws.send(json.dumps({'type' : 'messages', 'before': 0, 'messages': [db[str(i)] for i in range(max(0,idx - 100),idx)]}))
                            users[ws] = username
                            send_userlist()
                else:
                    break
            except (WebSocketError, ValueError, UnicodeDecodeError):      # ValueError happens for example when "No JSON object could be decoded", would be interesting to log it
                break

        if ws in users:
            del users[ws]
            del pings[ws]
            send_userlist()

    @route('/')
    @route('/index.html')
    @view('talktalktalk.html')
    def index():
        context = {'request': request}
        return (context)

    @route('/popsound.mp3')
    def popsound():
        return static_file('popsound.mp3', root='.')

    run(host=HOST, port=PORT, debug=True, server=GeventWebSocketServer)

class talktalktalk(daemon.Daemon):
    def run(self):
        main()

if len(sys.argv) == 1:           # command line interactive mode
    main()

elif len(sys.argv) == 2:         # daemon mode
    daemon = talktalktalk(pidfile='_.pid', stdout='log.txt', stderr='log.txt')

    if 'start' == sys.argv[1]:
        daemon.start()
    elif 'stop' == sys.argv[1]:
        daemon.stop()
    elif 'restart' == sys.argv[1]:
        daemon.restart()


from flask import render_template, redirect, url_for, flash, request, make_response, jsonify, send_from_directory, send_file
from app import db
from app.models import User, Teacher_Class, Class, Student_Class, Activity_Results, Student_Activity, Misspelled_Word, Class_Configs, Student_Class_Level, Find_Me_Configurations
from flask_login import login_user, logout_user, current_user, login_required
from app.main import bp
from datetime import datetime
import unidecode
import urllib.parse as urlparse
import random
from app.utils import checkForIncorrectCharacters
import requests
from app.doodle.routes import findWordForDoodle
from app.everyday_object_classification.routes import findWord, findCustomWord
from app.image_classification.routes import getCategoryWord, findCategoryWord
from app.create_classifier.routes import findCustomGameWord, getCustomGameWord
import os
import app
import json

languageIndices = {"French": 0, "Spanish": 1}
sports = {"golf":["golf", "golf"],"tennis":["tennis", "tenis"],"football":["foot", "fútbol"],"boxing":["boxe", "boxeo"],"skiing":["ski", "esquí"],"horseriding":["équitation", "equitación"],"basketball":["basket", "baloncesto"],"swimming":["nager", "nadando"],"volleyball":["volley-ball", "voleibol"],"rugby":["rugby", "rugby"]}
emotionsEmojis = {"angry":["fâché", "enojado"],"awkward":["gênant", "torpe"],"annoyed":["agacé", "irritado"],"cool":["impressionnant", "increíble"],"sad":["triste", "triste"],"shock":["choc", "choque"],"love":["amour", "amor"],"tired":["fatigué", "cansado"],"sick":["malade", "enfermos"], "scared":["effrayé", "asustado"], "funny":["marrant", "gracioso"]}
animals = {"bird":["oiseau", "pájaro"],"cat":["chat", "gato"],"cow":["vache", "vaca"],"dog":["chien", "perro"],"frog":["grenouille", "rana"],"hippo":["hipopotame", "hipopótamo"],"horse":["cheval", "caballo"],"lion":["leon", "león"],"monkey":["singe", "mono"], "bear": ["ours", "oso"], "fish":["poisson", "pez"], "pig":["porc", "cerdo"], "sheep": ["mouton", "oveja"], "tiger": ["tigre", "tigre"]}
vehicles = {"aeroplane":["avion", "avión"],"bike":["vélo", "bicicleta"],"bus":["autobus", "autobús"],"car":["voiture", "coche"],"train":["train", "entrenar"],"ferry":["traversier", "transportar"],"helicopter":["hélicoptère", "helicóptero"],"motorbike":["moto", "moto"],"taxi":["taxi", "taxi"], "truck":["camion", "camión"]}
classroom = {"monitor":["moniteur"],"glasses":["lunettes"],"computer":["ordinateur"],"mouse":["souris"],"CD Player":["lecteur CD"],"pencil case":["trousse"],"pencil sharpener":["taille-crayon"],"pen":["stylo"],"bottle":["bouteille"], "chair":["chaise"], "desk":["bureau"], "eraser":["gomme"], "printer":["imprimante"], "projector":["projecteur"], "backpack":["sac à dos"], "clock":["l'horloge"], "ruler":["règle"]}
allVocab = dict(sports)
allVocab.update(emotionsEmojis)
allVocab.update(animals)
allVocab.update(vehicles)
numbersDict = {1: ['un', "uno"], 2: ['deux', "dos"],3: ['trois', "tres"],4: ['quatre', "cuatro"],5: ['cinq', "cinco"],6: ['six', "seis"],7: ['sept', "siete"],8: ['huit', "ocho"],9: ['neuf', "nueve"],10: ['dix', "diez"],11: ['onze', "once"]}
voicesForApplication = {0: "French Female", 1: "Spanish Male"}
translationLanguage = {0: "fr", 1:"es"}

def get_class_vocab(dict):
    config = Class_Configs.query.filter_by(class_id=current_user.current_class).first()
    if config is None:
        return dict
    vocab_data = json.loads(config.vocab_data)
    vocab_data = vocab_data["Scramble"]
    vocab_dict = {}
    userClass = Class.query.filter_by(class_id=current_user.current_class).first_or_404()
    classLanguage = userClass.get_language()
    if classLanguage == "French":
        for eng in vocab_data:
            vocab_dict[eng] = [vocab_data[eng], ""]
    else:
        for eng in vocab_data:
            vocab_dict[eng] = ["" , vocab_data[eng]]

    print(vocab_dict)
    dict.update(vocab_dict)
    return dict

@bp.before_request
def before_request():
    if current_user.is_authenticated:
        current_user.last_seen = datetime.utcnow()
        db.session.commit()

@bp.route('/')
@login_required
def index():
    # print(current_user.current_class)
    if current_user.current_class == None:
        return redirect(url_for("user_management.createPracticeArea"))
    userClass = Class.query.filter_by(class_id=current_user.current_class).first_or_404()
    classLanguage = userClass.get_language()
    languageIndex = getLanguageIndex(classLanguage)
    get_class_vocab(allVocab)
    print(classLanguage)
    resp = make_response(render_template('index.html'))
    resp.set_cookie('word', " ")
    resp.set_cookie('correct', "0")
    resp.set_cookie('numGames', "0")
    resp.set_cookie('activity_id' " ")
    resp.set_cookie('languageIndex', str(languageIndex))
    return resp

def getLanguageIndex(lanaguage):
    return languageIndices[lanaguage]

@bp.route('/difficulty', methods=['GET', 'POST'])
@login_required
def diff():
    if request.method == 'POST':
        selectedClass = Class.query.filter_by(class_id=current_user.current_class).first()
        customGames = selectedClass.get_custom_games()
        config = Class_Configs.query.filter_by(class_id=current_user.current_class).first()
        if config is None:
            resp = make_response(render_template('level.html', config={}, current_level=0, customGames=customGames))
        else:
            currentLevel = current_user.get_current_level()
            configData = {'layout': json.loads(config.layout_data), 'data': json.loads(config.vocab_data),
                          'unlock_data': json.loads(config.unlock_data)}
            resp = make_response(render_template('level.html', config=configData, current_level=currentLevel,customGames=customGames))
        resp.set_cookie('correct', "0")
        resp.set_cookie('numGames', "0")
        resp.set_cookie('activity_id' " ")
        resp.set_cookie('word', " ")
        return resp
    else:
        selectedClass = Class.query.filter_by(class_id=current_user.current_class).first()
        customGames = selectedClass.get_custom_games()
        config = Class_Configs.query.filter_by(class_id=current_user.current_class).first()
        if config is None:
            resp = make_response(render_template('level.html', config={}, current_level=0, customGames=customGames))
        else:
            currentLevel = current_user.get_current_level()
            print(currentLevel)
            configData = {'layout': json.loads(config.layout_data), 'data': json.loads(config.vocab_data),
                          'unlock_data': json.loads(config.unlock_data)}
            resp = make_response(render_template('level.html', config=configData, current_level=currentLevel, customGames=customGames))
        resp.set_cookie('correct', "0")
        resp.set_cookie('numGames', "0")
        resp.set_cookie('activity_id' " ")
        resp.set_cookie('word', " ")
        return resp

@bp.route('/chooseGame', methods=['GET', 'POST'])
@login_required
def chooseGame():
    if request.method == 'POST':
        resp = make_response(render_template('level.html', config={}, current_level=0))
        resp.set_cookie('correct', "0")
        resp.set_cookie('numGames', "0")
        resp.set_cookie('activity_id' " ")
        resp.set_cookie('word', " ")
        return resp
    else:
        resp = make_response(render_template('level.html', config={}, current_level=0))
        resp.set_cookie('correct', "0")
        resp.set_cookie('numGames', "0")
        resp.set_cookie('activity_id' " ")
        resp.set_cookie('word', " ")
        return resp

@bp.route('/checkguess',methods=['POST'])
@login_required
def check():
    print("IN check guess")
    numGames = int(request.cookies.get("numGames"))
    languageIndex = int(request.cookies.get("languageIndex"))
    numGames += 1
    numCorrect = int(request.cookies.get("correct"))
    url = request.referrer
    path = urlparse.urlparse(url).path
    print(path)
    guess = request.form['guess'].strip()
    print(guess + "in check guess")
    word = request.cookies.get('word')
    print(word, "from swipe")
    if path == "/numbersOrSpelling" or path == "/swipeWords":
        word = allVocab[word][languageIndex]
        word = unidecode.unidecode(word)
    print(word + "------------")
    if "/swipe1" in path:
        print("in swipe1")
        category = path.split("/")[2]
        if languageIndex == 0:
            language = "fr"
        else:
            language = "es"
        word = findCategoryWord(word, category, language)
        print(word + " from swipe1")
        print(guess + " from swwipe1")
    if "/customSwipeGame" in path:
        category = path.split("/")[3]
        print(category)
        word = findCustomGameWord(word, category, current_user.current_class)
        print(word, guess)
    guess = guess.lower()
    word = word.lower()

    if guess == word or guess==(word+" "):

        numCorrect += 1
        resp = make_response(render_template('correct.html', path=path, score=[numCorrect, numGames]))
        resp.set_cookie("correct", str(numCorrect))
        resp.set_cookie("numGames", str(numGames))

        return resp
    else:
        resp = make_response(render_template('incorrect.html', path=path, score=[numCorrect, numGames]))
        resp.set_cookie("numGames", str(numGames))
        if path == "/numbersOrSpelling":
            checkForIncorrectCharacters(guess, word)
        return resp

@bp.route('/scramble', methods=['POST', "GET"])
@login_required
def scramble():
    word = request.cookies.get('word')
    active_activity = request.cookies.get("activity_id")
    languageIndex = int(request.cookies.get("languageIndex"))
    if active_activity != "":
        activity_id = active_activity
    else:
        activity_id = request.form["activity_id"]
    # print(word)
    if word == " ":
        key = random.choice(list(allVocab.keys()))
        print(key)

        wordForGame = allVocab[key][languageIndex]
        word = wordForGame
        print(word)

    # else:
    #     wordForGame = allVocab[word][0]

    wordlist = list(word)
    scrwordlist = random.sample(wordlist, len(wordlist))

    scrword = " "

    for i in range(0, len(scrwordlist)):
        scrword = (scrword + scrwordlist[i] + " ")
    print(scrword)
    scramble = {"Your word is....": scrword}

    resp = make_response(render_template('guess.html', scramble=scramble))
    resp.set_cookie('word', word)
    resp.set_cookie('activity_id', activity_id)
    return resp

@bp.route('/tts', methods=['POST', "GET"])
@login_required
def tts():
    word = request.cookies.get('word')
    active_activity = request.cookies.get("activity_id")
    languageIndex = int(request.cookies.get("languageIndex"))
    if active_activity != "":
        activity_id = active_activity
    else:
        print("from form")
        activity_id = request.form["activity_id"]
    print(active_activity, activity_id)
    # print(word)
    if word == " ":
        key = random.choice(list(allVocab.keys()))
        print(key)
        word = key
        print(word)
        wordForGame = allVocab[key][languageIndex]
    else:
        wordForGame = allVocab[word][languageIndex]

    voice = getVoice(languageIndex)

    scramble = {"Your word is....": [wordForGame, voice]}

    resp = make_response(render_template('tts.html', scramble=scramble))
    resp.set_cookie('word', word)
    resp.set_cookie('activity_id', activity_id)
    return resp

@bp.route('/about')
def about():
    return render_template('about.html')

@bp.route('/download')
def downloadUserManual():
    if current_user.role == "Student":
        return send_file('static/manuals/StudentManual.pdf', as_attachment=True)
    return send_file('static/manuals/TeacherManual.pdf', as_attachment=True)

@bp.route('/playAgain', methods=['POST'])
def playAgain():
    print(request.values)
    status = request.form["status"]
    path = request.form["path"]
    print(path)
    if(status == "correct"):
        resp = make_response(redirect(path))
        resp.set_cookie('word', " ")
        return resp
    resp = make_response(redirect(path))
    return resp

@bp.route('/answer',methods=['POST'])
def answer():
    #print(current_user.role)
    word = request.cookies.get('word')
    print(word)
    path = request.form["path"]
    languageIndex = int(request.cookies.get("languageIndex"))
    voice = ""
    english = word

    if path == "/doodle":
        translated, model = findWordForDoodle(word, languageIndex)
    elif path == "/find":
        translated = findWord(word, languageIndex)
    elif path == "/customFind":
        findMeConfig = Find_Me_Configurations.query.filter_by(class_id=current_user.current_class).first()
        translated = findCustomWord(word, findMeConfig.for_loc, findMeConfig.eng_loc)

    elif "findCategory" in path or "swipe1" in path:
        if languageIndex == 0:
            language = "fr"
        else:
            language = "es"
        category = path.split("/")[2]
        translated = findCategoryWord(word, category, language)

    elif "customGame" in path or "customSwipeGame" in path:
        print("customGame")
        translated = findCustomGameWord(word, path.split("/")[3], path.split("/")[2])

    elif path == "/tilesGame":
        print("tiles")
        newVocab = current_user.get_students_new_vocab(current_user.current_class).all()
        englishWords = [vocab.english for vocab in newVocab]
        translations =  [vocab.translation for vocab in newVocab]
        index = englishWords.index(word)
        english = word
        word = translations[index]
        translated = word

    elif path == "/numbersOrSpelling" or path == "/swipeWords":
        english = word
        word = allVocab[word][languageIndex]
        translated = word

    elif path == "/numbers":
        english = word
        translated = numbersDict[int(english)][languageIndex]

    elif word in allVocab:
        translated = allVocab[word][languageIndex]

    else:
        print(word)
        translated = word
        #word = allVocab[word][languageIndex]
        translated = word

    if path == "/speech":
        voice = getVoice(languageIndex)

    if path != "/speech" and path != "/scramble":
        misspelled_word = Misspelled_Word(class_id=current_user.current_class, user_id=current_user.user_id, word=english, translated_word=translated)
        db.session.add(misspelled_word)
        db.session.commit()
    dword = { "Your word was....":word}
    resp = make_response(render_template('answer.html',dword=dword, path=path, voice=voice))
    resp.set_cookie("word", " ")
    return resp

@bp.route("/endgame")
def endGame():
    #Update DB with score and currect user ID
    if(int(request.cookies.get("numGames")) == 0):
        score = 0
    else:
        score = int(request.cookies.get("correct")) / int(request.cookies.get("numGames")) * 100
    activity_id = int(request.cookies.get("activity_id"))
    print(activity_id)
    activity_result = Activity_Results(activity_id=activity_id, score=score, class_id=current_user.current_class)
    db.session.add(activity_result)
    db.session.commit()
    student_activity = Student_Activity(activity_instance_id=activity_result.activity_instance_id, student_id=current_user.user_id)

    db.session.add(student_activity)
    db.session.commit()

    config = Class_Configs.query.filter_by(class_id=current_user.current_class).first()
    if config is not None:
        current_level = current_user.get_current_level()
        active_level = current_user.get_active_level()
        print(current_level, active_level)
        if active_level < current_level:
            ##do update here
            current_user.freeze_stats()
            current_user.update_level(current_level)
            print("Current Level: " + str(current_level))
            if active_level != 3:
                flash("Congratulations, you have completed Course " + str(active_level) + ", you can now see your stats for this course!")
            else:
                flash("Congratulations you have completed all courses in your class! More statistics have been unlocked as well as a Certificate of Completion!")
            return redirect(url_for("user_management.user", username=current_user.username))

    return redirect(url_for("main.index"))

@bp.route("/swipeWords", methods=["POST", "GET"])
@login_required
def swipeWords():
    word = request.cookies.get('word')
    print(word, "empty")
    get_class_vocab(allVocab)
    verbs.update(allVocab)
    active_activity = request.cookies.get("activity_id")
    languageIndex = int(request.cookies.get("languageIndex"))
    if active_activity != "":
        activity_id = active_activity
    else:
        print("from form")
        activity_id = request.form["activity_id"]
    print(active_activity, activity_id)
    print(word)
    if word == " ":
        key = random.choice(list(verbs.keys()))
        print(key, "word")
        word = verbs[key][languageIndex]
        print(word)
        wordForGame = key
    else:
        wordForGame = word
    scramble = {"Your word is....": wordForGame}

    resp = make_response(render_template('swipeWords.html', scramble=scramble))
    print(word, wordForGame)
    resp.set_cookie('word', wordForGame)
    resp.set_cookie('activity_id', activity_id)
    return resp

verbs = {"to go": ["aller", "ir"], "to think": ["penser", "pensar"], "to talk": ["parler", "hablar"]}
from random import shuffle

@bp.route("/getWords/<word>")
def getWords(word):
    get_class_vocab(allVocab)
    verbs.update(allVocab)
    languageIndex = int(request.cookies.get("languageIndex"))
    print(languageIndex)
    word = " ".join(word.split("_"))
    print(word)
    keyWord = verbs[word][languageIndex]
    print(keyWord)
    data = [keyWord]
    for i in range(9):
        print(i)
        data.append(verbs[random.choice(list(verbs.keys()))][languageIndex])
    shuffle(data)
    print(data)
    response = {'data': data}
    return jsonify(response)

@bp.route("/phraseTranslator", methods=['POST', 'GET'])
@login_required
def phraseTranslator():
    languageIndex = int(request.cookies.get("languageIndex"))
    language = getTranslationLanguage(languageIndex)
    return render_template("tagger.html", language=language)

@bp.route("/info", methods=['POST', 'GET'])
def info():
    return render_template("landing_page.html")

def getTranslationLanguage(languageIndex):
    return translationLanguage[languageIndex]

def getVoice(languageIndex):
    return voicesForApplication[languageIndex]


# Some magic for Python3
try:
  import SocketServer as socketserver
except ImportError:
  import socketserver

import logging
import sys
import threading

class ATCommandUDPHandler(socketserver.BaseRequestHandler):
  def command(self, cmd):
    logging.info('Got command: %s' % (cmd,))

  def handle(self):
    commands = self.request[0].strip().split('\n')
    socket = self.request[1]

    logging.info('%s -> %s' % (self.client_address, repr(commands)))
    for command in commands:
      self.command(command)

class ATCommandListener(threading.Thread):
  def __init__(self):
    super(ATCommandListener, self).__init__()

    self.should_quit = False
    self.cv = threading.Condition()

    (bind_host, bind_port) = '127.0.0.1', 5556
    self._server = socketserver.UDPServer((bind_host, bind_port),
        ATCommandUDPHandler)

  def run(self):
    self._server.serve_forever()

def serve():
  # Create an AT command listener
  at_listener = ATCommandListener()
  at_listener.daemon = True
  at_listener.start()

  # Wait for the quit command
  at_listener.cv.acquire()
  while not at_listener.should_quit:
    at_listener.cv.wait(0.5)
  at_listener.cv.release()

if __name__ == '__main__':
  logging.basicConfig(level=logging.INFO)
  serve()


from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException
from pydantic import HttpUrl
from starlette.requests import Request
from tweepy import API

from app.api.deps import get_backend, get_twitter_api
from app.core import TwitterSentiment
from app.decorator import async_log_response
from app.feature_store.backends import Backend
from app.feature_store.core import Keys, get_cache, set_cache
from app.schemas import SentimentResponse

router = APIRouter()

@router.get("/inference", response_model=SentimentResponse)
@async_log_response
async def inference(
    request: Request,
    tweetUrl: HttpUrl,
    background_tasks: BackgroundTasks,
    feature_store: Backend = Depends(get_backend),
    twitter_api: API = Depends(get_twitter_api),
):
    try:
        tweet = twitter_api.get_status(tweetUrl.split("/")[-1])
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Error{type(e).__name__}: {e}")
    key = Keys(tweet=tweet)
    data = await get_cache(keys=key, feature_store=feature_store)
    if not data:
        request.app.logger.info("Prediction is not exist in feature store")
        twitter = TwitterSentiment(**request.app.model_params)

        prediction = twitter.prediction(tweet.text)
        if prediction:
            result = SentimentResponse(
                sentiment_analyst=prediction, text_input=tweet.text
            )
            background_tasks.add_task(set_cache, result.dict(), key, feature_store)
        else:
            raise HTTPException(status_code=400, detail="Empty prediction")
    else:
        request.app.logger.info("Prediction hits")
        result = SentimentResponse(**data)

    return result

@router.get("/clear")
async def refresh(prefix: str, feature_store: Backend = Depends(get_backend)):
    flag_clear = await feature_store.clear(namespace=prefix)
    if flag_clear and any(flag_clear):
        return "success"
    return "fail"