import logging
from collections import OrderedDict
from typing import List

from .db import DBConfig
from .errors import UnsupportedAppError

class GamestPlugin:
    SETTINGS_TAB_NAME = "Plugin"
    def __init__(self, application):
        self.logger = logging.getLogger(__name__ + '.' + self.__class__.__name__)
        self.application = application

        self.config = DBConfig(owner=self.__class__.__name__)

    def cleanup(self):
        raise NotImplementedError

    @classmethod
    def get_settings_template(cls):
        return OrderedDict()

class GamestPersistentPlugin(GamestPlugin):
    pass

class GamestSessionPlugin(GamestPlugin):
    start_bind = None
    end_bind = None

    def __init__(self, application):
        super().__init__(application)
        self.play_session = application.play_session

class NotificationService(GamestPersistentPlugin):
    user_name = None

    def notify(self, msg):
        raise NotImplementedError

class GameReporterPlugin(GamestSessionPlugin):
    PATH_ENDSWITH : List[str] = []

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

        if self.user_app_ids and (self.play_session.user_app_id not in self.user_app_ids):
            self.logger.debug(
                "User app ID did not match: %s not in %s (configured)",
                self.play_session.user_app_id,
                ', '.join(str(i) for i in self.user_app_ids))
            raise UnsupportedAppError(
                "Current user_app_id: {}. Configured: {}.".format(
                    self.play_session.user_app_id,
                    ', '.join(str(i) for i in self.user_app_ids)))
        elif not self.user_app_ids and not (self.play_session.user_app.path and any(self.play_session.user_app.path.endswith(p) for p in self.PATH_ENDSWITH)):
            self.logger.debug("Path did not match: %s.", self.play_session.user_app.path)
            raise UnsupportedAppError("Current app path does not match a supported path.")

        self.job = None

    @property
    def user_app_ids(self):
        return self.config.getlist('user_app_ids', type=int)

    @property
    def send_begin(self):
        return self.config.getboolean('send_begin', fallback=True)

    @property
    def send_end(self):
        return self.config.getboolean('send_end', fallback=True)

    @property
    def add_status_updates(self):
        return self.config.getboolean('add_status_updates', fallback=True)

    @property
    def interval(self):
        return self.config.get('interval', type=int, fallback=60)

    @classmethod
    def get_settings_template(cls):
        d = super().get_settings_template()
        d[(cls.__name__, 'send_begin')] = {
            'name' : 'Notify on game begin',
            'type' : 'bool',
            'default' : True,
            'hint' : "If checked, send a notification when a game session begins.",
        }
        d[(cls.__name__, 'send_end')] = {
            'name' : 'Notify on game end',
            'type' : 'bool',
            'default' : True,
            'hint' : "If checked, send a notification when a game session ends.",
        }
        d[(cls.__name__, 'add_status_updates')] = {
            'name' : 'Save updates in DB',
            'type' : 'bool',
            'default' : True,
            'hint' : "If checked, status updates from this plugin will be saved to the database.",
        }
        d[(cls.__name__, 'interval')] = {
            'name' : 'Update interval (minutes)',
            'type' : 'text',
            'validate' : int,
            'default' : '60',
            'hint' : "How often to send updates, in minutes.",
        }
        d[(cls.__name__, 'user_app_ids')] = {
            'name' : 'UserApp IDs',
            'type' : 'list',
            'hint' : ("By default, this plugin determines which game is running based on the "
                      "executable path. This can fail if your executable has an unexpected "
                      "filename or if another game has the same filename. In that case, set the "
                      "UserApp ID to the number displayed when gamest is detecting the game, and "
                      "this plugin will activate only for that app.\n"
                      "\n"
                      "To specify multiple app IDs (e.g. for different game versions), put one ID "
                      "per line.")
        }
        return d

    def get_report(self):
        raise NotImplementedError

    def report_update(self, game_end=False):
        self.logger.debug("report_update called")
        if game_end:
            report_text = "{{user_name}} played **{}**:\n\n".format(self.play_session.user_app.app)
        else:
            report_text = "{{user_name}} is playing **{}**:\n\n".format(self.play_session.user_app.app)
        try:
            report_details = self.get_report()
            if report_details:
                if self.add_status_updates:
                    self.play_session.add_status_update(report_details)
                report_text = report_text + report_details
                for s in filter(lambda p: isinstance(p, NotificationService), self.application.persistent_plugins):
                    s.notify(report_text)
            else:
                self.logger.debug("No difference since report.")
        except Exception:
            self.logger.exception("Failed to build update.")
        finally:
            if self.application.RUNNING and self.play_session is self.application.play_session:
                self.job = self.application.after(self.interval*60*1000, self.report_update)

    def onGameStart(self, e):
        self.logger.debug("onGameStart called")
        if self.send_begin:
            self.job = self.application.after(35000, self.report_update)
        else:
            self.job = self.application.after(self.interval*60*1000, self.report_update)

    def onGameEnd(self, e):
        self.logger.debug("onGameEnd called")
        if self.send_end:
            self.report_update(game_end=True)
        self.cleanup()

    def cleanup(self):
        if self.job:
            self.application.after_cancel(self.job)

class IdentifierPlugin(GamestPersistentPlugin):
    def candidates(self):
        return []

    def identify_game(self):
        return None

    def clear_cache(self):
        pass


from pygame.image import load
from pygame import surface
from pygameimagetool import getPath
from scrollimage import Direction, ScrollImage

BACKGROUND = load(getPath("sprites", "background.png"))
FOREGROUND = load(getPath("sprites", "base.png"))
DIRECTION = Direction.HORIZONTAL

class Background(surface.Surface):

    def __init__(self, size, foreground_y, bg_speed, fg_speed, pipes):
        super().__init__(size)
        self.foreground_y = foreground_y
        self.bg = ScrollImage(BACKGROUND, DIRECTION, bg_speed)
        self.fg = ScrollImage(FOREGROUND, DIRECTION, fg_speed)
        self.pipes = pipes

    def draw(self, show_pipes=True):
        bg, fg = self.bg, self.fg
        self.blit(bg, (0, 0))
        show_pipes and self.pipes.group_pipes.draw(self)
        self.blit(fg, (0, self.foreground_y))
        bg.draw()
        fg.draw()

    def move(self):
        self.bg.move()
        self.fg.move()

    def update(self):
        self.draw()
        self.move()


#ITQ模块
import math
from numpy import *
import numpy as np
import torch
import torch.nn as nn
from scipy.io import loadmat
nbits = 32
m = loadmat("E:\北航研究生\实验室\hashing-baselines\datasets\Cifar10-Gist512.mat")

#  pca函数
def pca(dataMat, topNfeat=nbits):    # topNfeat 降维后的维度
    meanVals = mean(dataMat, axis=0)     # 按列求均值，即每一列求一个均值，不同的列代表不同的特征
    # print meanVals
    meanRemoved = dataMat - meanVals   # remove mean     # 去均值，将样本数据的中心点移到坐标原点
    print(meanRemoved)
    covMat = cov(meanRemoved, rowvar=0)         # 计算协方差矩阵
    # print covMat
    eigVals, eigVects = linalg.eig(mat(covMat))   # 计算协方差矩阵的特征值和特征向量
    # print eigVals
    # print eigVects
    eigValInd = argsort(eigVals)            # sort, sort goes smallest to largest  #排序，将特征值按从小到大排列
    # print eigValInd
    eigValInd = eigValInd[:-(topNfeat+1):-1]  # cut off unwanted dimensions      #选择维度为topNfeat的特征值
    # print eigValInd
    redEigVects = eigVects[:,eigValInd]       # reorganize eig vects largest to smallest   #选择与特征值对应的特征向量
    print(redEigVects)
    lowDDataMat = meanRemoved * redEigVects   # transform data into new dimensions    #将数据映射到新的维度上，lowDDataMat为降维后的数据
    print(lowDDataMat)
    reconMat = (lowDDataMat * redEigVects.T) + meanVals         # 对原始数据重构，用于测试
    print(reconMat)
    return lowDDataMat, reconMat

# 迭代优化ITQ
def get_ITQ_binary():
    r = torch.randn(nbits, nbits)
    R = mat(r)
    # print(R.shape)
    U, Sigma, VT = linalg.svd(R)  # sigma只返回奇异值
    R = U[:, 0:nbits]  # R是初始化的旋转矩阵
    lowDDataMat, reconMat = pca(m['X'])
    for ii in range(50):
        Z = lowDDataMat * R
        UX = (-1)*(Z<0)+\
             (+1)*(Z>=0)
    #     UX = mat(temp1)
        C = UX.transpose() * lowDDataMat
        U1, Sigma, VT1 = linalg.svd(C)
        R = VT1 * U1.transpose()
    #     print('iter:%d, loss:%f\n',iter,sum(np.sum(np.array(UX)-Z)))
    #     l = nn.MSELoss()
    #     print(l(torch.Tensor(UX),torch.Tensor(Z)))
    #     print(UX[:1])
    #     print('iter:%d, loss:%f\n',ii,mean(np.sum(np.square(UX-Z))))
    print("ITQ finished")
    R_last = R  # 旋转矩阵
    return UX

# get_ITQ_binary()


import math
import random
import re
import time
import numpy
from numpy.random import choice
import discord

# import enchant
#import ssdeep
# from guess_language import guess_language

import generalconfig as gconf
from .. import events
from .. import util, dbconn
from ..game import players, raids
from ..game import stats, weapons, quests, awards, shields
from ..game.configs import dueserverconfig
from ..game.helpers import imagehelper
from . import gamerules
from ..game.helpers.misc import Ring

# from threading import Lock

SPAM_TOLERANCE = 80
# For awards in the first week. Not permanent.
old_players = open('oldplayers.txt').read()  # For comeback award
testers = open('testers.txt').read()  # For testers award
# spelling_lock = Lock()
awayfromkeyboard = []
awayfromkeyboardmessage = []
criminals = []

def get_spam_level(player, message_content):
    #"""
    #Get's a spam level for a message using a
    #fuzzy hash > 50% means it's probably spam
    #"""
    #message_hash = ssdeep.hash(message_content)
    spam_level = 0
    #spam_levels = [ssdeep.compare(message_hash, prior_hash) for prior_hash in player.last_message_hashes if
                   #prior_hash is not None]
    #if len(spam_levels) > 0:
        #spam_level = max(spam_levels)
    #player.last_message_hashes.append(message_hash)
    #if spam_level > SPAM_TOLERANCE:
        #player.spam_detections += 1
    return spam_level

def original_get_spam_level(player, message_content):
    """
    Get's a spam level for a message using a
    fuzzy hash > 50% means it's probably spam
    """
    if not hasattr(player, "last_message_hashes"):
        player.last_message_hashes = Ring(10)
    message_hash = ssdeep.hash(message_content)
    spam_level = 0
    spam_levels = [ssdeep.compare(message_hash, prior_hash) for prior_hash in player.last_message_hashes if
                   prior_hash is not None]
    if len(spam_levels) > 0:
        spam_level = max(spam_levels)
    player.last_message_hashes.append(message_hash)
    if spam_level > SPAM_TOLERANCE:
        player.spam_detections += 1
    return spam_level

def progress_time(player):
    return time.time() - player.last_progress >= 60

def quest_time(player):
    return time.time() - player.last_quest >= quests.QUEST_COOLDOWN

async def player_message(message, player, spam_level):
    """
    W.I.P. Function to allow a small amount of exp
    to be gained from messaging.

    """

    def get_words():
        return re.compile('\w+').findall(message.content)

    if player is not None:

        # Mention the old bot award
        if gconf.DEAD_BOT_ID in message.raw_mentions:
            await awards.give_award(message.channel, player, "SoCold", "They're not coming back.")
        # Art award
        if player.misc_stats["art_created"] >= 100:
            await awards.give_award(message.channel, player, "ItsART")
        if not hasattr(player, "job"):
            player.job = []
        if not hasattr(player, "learning"):
            player.learning = False
        if not hasattr(player, "learning_end"):
            player.learning_end = time.time()
        if not hasattr(player, "tax_status"):
            player.tax_status = False
        if not hasattr(player, "tax_filing"):
            player.tax_filing = time.time()
        if player.learning:
            if int(player.learning_end) <= int(time.time()):
                await util.say(message.author, "Your learning time is finished!")
                player.learning = False
                x = 0
                y = 1
                z = 2
                if "head teacher" in player.job:
                    x += z
                if "teacher" in player.job:
                    x += y
                player.education += (x / z) + 1
        if player.tax_status:
            if (int(time.time() - int(player.tax_filing))) > 172800:
                a = (player.level * ((int(time.time()) - int(player.tax_filing)) / 2) / 10)
                if (player.money - a) < 10000:
                    a = player.money - 10000
                    player.money = 10000
                else:
                    player.money -= a
                player.tax_filing = int(time.time())
                if player.id in criminals:
                    criminals.remove(player.id)
                await util.say(message.author, "You lost %s DUT for failing to file your taxes for your job!" % (a))
            if (int(time.time()) - int(player.tax_filing)) > 86400:
                if not player.id in criminals:
                    criminals.append(player.id)
                    await util.say(message.author, "You have been put on the toil raid list because you haven't filed taxes in a day!")
        if progress_time(player) and spam_level < SPAM_TOLERANCE:

            if len(message.content) > 0:
                player.last_progress = time.time()
            else:
                return
            if not hasattr(player, "language"):
                player.language = "en"

            # Special Awards
            # Comeback award
            if player.id in old_players:
                await awards.give_award(message.channel, player, "CameBack", "Return to DueUtil")
            # Tester award
            if player.id in testers:
                await awards.give_award(message.channel, player, "Tester", ":bangbang: **Something went wrong...**")
            # Donor award
            if player.donor:
                await awards.give_award(message.channel, player, "Donor",
                                        "Donate to DueUtil!!! :money_with_wings: :money_with_wings: :money_with_wings:")
            # DueUtil tech award
            if dbconn.conn()["dueutiltechusers"].find({"_id": player.id}).count() > 0:
                if "DueUtilTech" not in player.awards:
                    player.inventory["themes"].append("dueutil.tech")
                await awards.give_award(message.channel, player, "DueUtilTech", "<https://dueutil.org/>")

            ### DueUtil - the hidden spelling game!
            # The non-thread safe Apsell calls
            # spelling_lock.acquire()
            # DISABLED: Spell checking due to random seg faults (even with locks).
            """lang = guess_language(message.content)
            if lang in enchant.list_languages():
                spelling_dict = enchant.Dict(lang)
            else:
                spelling_dict = enchant.Dict("en_GB")

            spelling_score = 0
            big_word_count = 1
            big_word_spelling_score = 0
            message_words = get_words()
            for word in message_words:
                if len(word) > 4:
                    big_word_count += 1
                if random.getrandbits(1):  # spelling_dict.check(word):
                    spelling_score += 3
                    if len(word) > 4:
                        big_word_spelling_score += 1
                else:
                    spelling_score -= 1
            # spelling_lock.release()
            # We survived?!

            spelling_score = max(1, spelling_score / ((len(message_words) * 3) + 1))
            spelling_avg = player.misc_stats["average_spelling_correctness"]
            1 - abs(spelling_score - spelling_avg)
            spelling_strg = big_word_spelling_score / big_word_count
            # Not really an average (like quest turn avg) (but w/e)
            player.misc_stats["average_spelling_correctness"] = (spelling_avg + spelling_score) / 2

            len_limit = max(1, 120 - len(message.content))
            player.progress(spelling_score / len_limit, spelling_strg / len_limit, spelling_avg / len_limit)
            """

            player.hp = 10 * player.level
#            await check_for_level_up(message, player)
            player.save()

    else:
        players.Player(message.author)
        stats.increment_stat(stats.Stat.NEW_PLAYERS_JOINED)

async def check_for_level_up(ctx, player):
    """
    Handles player level ups.
    """

    exp_for_next_level = gamerules.get_exp_for_next_level(player.level)
    level_up_reward = 0
    while player.exp >= exp_for_next_level:
        player.exp -= exp_for_next_level
        player.level += 1
        level_up_reward += (player.level * 10)
        player.money += level_up_reward

        stats.increment_stat(stats.Stat.PLAYERS_LEVELED)

        exp_for_next_level = gamerules.get_exp_for_next_level(player.level)
    stats.increment_stat(stats.Stat.MONEY_CREATED, level_up_reward)
    if level_up_reward > 0:
        if dueserverconfig.mute_level(ctx.channel) < 0:
            await imagehelper.level_up_screen(ctx.channel, player, level_up_reward)
        else:
            util.logger.info("Won't send level up image - channel blocked.")
        rank = player.rank
        if 1 <= rank <= 10:
            await awards.give_award(ctx.channel, player, "Rank%d" % rank, "Attain rank %d." % rank)

async def team_check_for_level_up(ctx, team):
    """
    Handles team level ups.
    """

    exp_for_next_level = gamerules.team_get_exp_for_next_level(team.level)
    level_up_reward = 0
    while team.exp >= exp_for_next_level:
        team.exp -= exp_for_next_level
        team.level += 1

        exp_for_next_level = gamerules.team_get_exp_for_next_level(team.level)

async def manage_quests(message, player, spam_level):
    """
    Gives out quests!
    """
    if player.learning == True:
        return
    channel = message.channel
    server = message.server
    if time.time() - player.quest_day_start > quests.QUEST_DAY and player.quest_day_start != 0:
        if not hasattr(player, "quests_spawned"):
            player.quests_spawned = 0
        if not hasattr(player, "count"):
            player.count = 0
        if not hasattr(player, "discoin_transactions"):
            player.discoin_transactions = 0
        player.discoin_transactions = 0
        player.quests_completed_today = 0
        player.quest_day_start = 0
        player.quests_spawned = 0
        if player.id != "376166105917554701":
            player.count = 0
        util.logger.info("%s (%s) daily completed quests reset", player.name_assii, player.id)
    if not hasattr(player, "language"):
        player.language = "en"

    # Testing
    if len(quests.get_server_quest_list(channel.server)) == 0:
        quests.add_default_quest_to_server(message.server)
    if quest_time(player):
        if not hasattr(player, "double_quests"):
            player.double_quests = False
        if player.double_quests:
            x = quests.TEAM_MAX_DAILY_QUESTS
        else:
            x = quests.MAX_DAILY_QUESTS
        player.last_quest = time.time()
        if quests.has_quests(channel) and len(
                player.quests) < quests.MAX_ACTIVE_QUESTS and player.quests_completed_today < x and spam_level < SPAM_TOLERANCE:
            server_quests_2 = quests.get_channel_quests(channel)
            r = random.randint(2, 7)
            while r > 0:
                try:
                    if len(server_quests_2) > 0:
                        random.shuffle(server_quests_2)
                except:
                    pass
                r -= 1
            quest_list_2 = []
            quest_chances = []
            for Quest in server_quests_2:
                c = Quest.spawn_chance
                quest_chances.append(int(c * 100))
            quest_chance_sum = sum(quest_chances)
            if not hasattr(player, "quest_spawn_boost"):
                player.quest_spawn_boost = 1
            quest_number_2 = 0
            quest_chances_2 = []
            chosen_quests = []
            check = 0
            while len(server_quests_2) > 0 and check < 20:
                quest = server_quests_2.pop()
                if not hasattr(quest, "rank"):
                    quest.rank = 1000
                #TODO: Take out the loop breaks, as they keep stopping multiple quests from being picked down at TODO 2
                if (quest.rank - 1) == (player.level // 10) or quest.rank == 1000:
                    if quest.channel == message.channel or quest.channel in ("All", "all", message.channel):
                        chosen_quests.append(quest)
                        quest_chances_2.append(((quest.spawn_chance * 100) / quest_chance_sum))
                if len(chosen_quests) == 0:
                    channel_quests = quests.get_channel_quests(channel)
                    while len(channel_quests) > 0:
                        quest = channel_quests.pop()
                        if not hasattr(quest, "rank"):
                            quest.rank = 1000
                        if (quest.rank - 1) == (player.level //10) or quest.rank == 1000:
                            chosen_quests.append(quest)
                            quest_chances_2.append(((quest.spawn_chance * 100) / quest_chance_sum))
                if len(chosen_quests) == 0:
                    channel_quests = quests.get_channel_quests(channel)
                    while len(channel_quests) > 0:
                        quest = channel_quests.pop()
                        if not hasattr(quest, "rank"):
                            quest.rank = 1000
                        chosen_quests.append(quest)
                        quest_chances_2.append(((quest.spawn_chance * 100) / quest_chance_sum))
                check += 1
                quest.save()
            quest_list_3 = []
            q = 0
            check_2 = 0
            for i in chosen_quests:
                if check_2 >= 20:
                    break
                quest_list_3.append(q)
                q += 1
                check_2 += 1
            #TODO 2: Form a loop for quest_2 and insert the quest breaks there
            if len(quest_list_3) and len(quest_chances_2) != 0:
                quest_2 = choice(quest_list_3, 1, quest_chances_2).tolist()
                quest_2 = quest_2[0]
            else:
                try:
                    quest_2 = random.randint(0, int(len(chosen_quests) - 1))
                except:
                    quest_2 = 0
            quest = chosen_quests[quest_2]
            if not hasattr(player, "prestige_coins"):
                player.prestige_coins = 0
            if not hasattr(player, "topdog_attempts"):
                player.topdog_attempts = 0
            if not hasattr(player, "quest_spawn_boost"):
                player.quest_spawn_boost = 1
            a = (player.prestige_coins / 3000 + 1) * (player.topdog_attempts / 250 + 1)
            if not player.learning:
                b = player.quest_spawn_boost
            else:
                b = (player.quest_spawn_boost / 2)
            if random.random() <= .08:
                raid = raids.get_random_raid()
                raids.ActiveRaid(raid.name, (player.level))
            elif random.random() <= (.5 * player.quest_spawn_build_up * b): # add * a if kurono isnt super OP
                new_quest = await quests.ActiveQuest.create(quest.q_id, player)
                stats.increment_stat(stats.Stat.QUESTS_GIVEN)
                player.quest_spawn_build_up = 1
                try:
                    if dueserverconfig.mute_level(message.channel) < 0:
                        await imagehelper.new_quest_screen(channel, new_quest, player)
#                    else:
#                        util.logger.info("Won't send new quest image - channel blocked.")
                except:
                    await util.say(channel, "You recieved a quest but something went wrong with sending the image!")
                util.logger.info("%s has received a quest [%s][s-id: %s]", player.name_assii, new_quest.q_id, message.server.id)
            else:
                player.quest_spawn_build_up += 0.5

async def check_for_recalls(ctx, player):
    """
    Checks for weapons that have been recalled
    """

    current_weapon_id = player.equipped["weapon"]

    weapons_to_recall = [weapon_id for weapon_id in player.inventory["weapons"] + [current_weapon_id]
                         if (weapons.get_weapon_from_id(weapon_id).id == weapons.NO_WEAPON_ID
                             and weapon_id != weapons.NO_WEAPON_ID)]

    if len(weapons_to_recall) == 0:
        return
    if current_weapon_id in weapons_to_recall:
        player.weapon = weapons.NO_WEAPON_ID
    player.inventory["weapons"] = [weapon_id for weapon_id in player.inventory["weapons"] if
                                   weapon_id not in weapons_to_recall]
    recall_amount = sum([weapons.get_weapon_summary_from_id(weapon_id).price for weapon_id in weapons_to_recall])
    player.money += recall_amount
    player.save()
#    await util.say(ctx.channel, (
#        ":bangbang: " + ("One" if len(weapons_to_recall) == 1 else "Some") + " of your weapons has been recalled!\n"
#        + "You get a refund of ``" + util.format_number(recall_amount, money=True, full_precision=True) + "``"))

async def check_afk(ctx, player):
    if not hasattr(player, "afk_message"):
        player.afk_message = 0
    if not hasattr(player, "afk"):
        player.afk = 0
    player.afk = 0
    a = player.id
    b = player.afk_message
    awayfromkeyboard.remove(a)
    awayfromkeyboardmessage.remove(b)
    await util.say(ctx.channel, "Removed your afk!")
    player.afk_message = 0

async def check_afk_ping(ctx, message):
    try:
        player_id = "".join(str(x) for x in message.raw_mentions)
        player = players.find_player(player_id)
        if not hasattr(player, "afk"):
            player.afk = 0
        if not hasattr(player, "afk_message"):
            player.afk_message = 0
        if player.id in awayfromkeyboard and player.afk == 1:
            afk_embed = discord.Embed(title="This player is AFK!", type="rich",
                                    color=gconf.DUE_COLOUR)
            afk_embed.add_field(name="AFK Message:", value=player.afk_message)
            await util.say(ctx.channel, embed=afk_embed)
    except AttributeError:
        pass
    except ValueError:
        pass

async def give_attributes(ctx, player):
    if not hasattr(player, "shield"):
        player.shield = shields.NO_SHIELD_ID
    if not hasattr(player, "shields"):
        player.shields = []
#    if not hasattr(player, "afk"):
#        player.afk = 0
#    if not hasattr(player, "afk_message"):
#        player.afk_message = 0
    if not hasattr(player, "learning"):
        player.learning = False
#    if not hasattr(player, "attribute_holder"):
#        player.attribute_holder = []
#    if not hasattr(player, "broken_status"):
#        player.broken_status = False
#    if hasattr(player, "shield"):
#        if getattr(player, "shield") == shields.NO_SHIELD_ID:
#            player.shield = shields.NO_SHIELD
#    if hasattr(player, "shield"):
#        player.shield = shields.NO_SHIELD
#    if hasattr(player, "shields"):
#        if getattr(player, "shields") == [shields.NO_SHIELD]:
#            player.shields = [shields.NO_SHIELD_2]
#    if hasattr(player, "shields"):
#        if getattr(player, "shields") == []:
#            player.shields = [shields.NO_SHIELD_2]
#    if hasattr(player, "shields"):
#        player.shields = []
#    player.shields = []
#    player.shields.clear()
#    player.inventory["shields"] = []
#    player.inventory["shields"].clear()

async def on_message(message):
    player = players.find_player(message.author.id)
    spam_level = 0
    if player is not None:
        if not player.is_playing(message.server):
            return
        if quest_time(player) or progress_time(player):
            spam_level = get_spam_level(player, message.content)
        await give_attributes(message, player)
    await player_message(message, player, spam_level)
    if player is not None:
        await manage_quests(message, player, spam_level)
        await check_for_recalls(message, player)
#        await check_afk_ping(message, message)
#        if player.afk == 1:
#            await check_afk(message, player)
#            player.afk = 0

events.register_message_listener(on_message)


#!/usr/bin/python3
"""A Python Script that takes in a letter and sends a
POST request to http://0.0.0.0:5000/search_user
with the letter as a parameter."""

import requests
import sys

if __name__ == "__main__":
    # Get the command line argument for the letter
    if len(sys.argv) > 1:
        q = sys.argv[1]
    else:
        q = ""

    # Send the POST request to the API endpoint
    response = requests.post("http://0.0.0.0:5000/search_user", data={"q": q})

    try:
        # Parse the response JSON
        data = response.json()
        if data:
            # Print the ID and name of the user(s) found
            print("[{}] {}".format(data["id"], data["name"]))
        else:
            print("No result")
    except ValueError:
        print("Not a valid JSON")


from tilemap import *
import output
import scripts.bbase2 as script
import omg
import sys
import time
import math

print("-------------")
print("Jade")
print("   by jmickle")
print("-------------")
print("")
print("")
start_time = time.clock()
if (len(sys.argv) < 2):
	#print("Usage: jade.py output_path [config_path]")
	til = script.run()
	print("Generated in "+str(math.floor(100*(time.clock() - start_time))/100)+" seconds.")
	output_path = "c:/users/jmickle/documents/bronze.wad"
	print("Exporting to "+output_path)
	output.to_file(til,output_path,"ZENNODE")

else:
	output_path = sys.argv[1]
	if (len(sys.argv) > 2): config_path = sys.argv[2]
	til = scripts.cavey.run()
	output.to_file(til,output_path)
print("Jade completed in "+str(math.floor(100*(time.clock() - start_time))/100)+" seconds.")


# coding: utf-8

# In[1]:

import random
import numpy as np
import matplotlib.pyplot as plt
import copy
import pandas as pd             #一个强大的分析结构化数据的工具集，用来画Q表

# In[2]:

plt.rcParams['font.sans-serif']=['SimHei']       #使画的图像可以呈现中文
plt.rcParams['axes.unicode_minus']=False

from matplotlib.pyplot import MultipleLocator   #用于设置坐标轴的刻度间隔

# In[3]:

N =80          #个体空间移动的步数            80
NUM = 100      #随机生成的网络总人数          100   10000
n = 5           #初始感染人数                 5
L = 4           #二维网络大小，L*L            4      10
CitySize = 50   #网格边长（画动态图用的）     50
BigCity = 0     #大城市个数
attractUP = 2   #每次增加的地点影响力值       2

d = 1           #个体每次移动的距离，即移动一个城市（或者不移动）
time = 1       #迭代次数（画SIR图用）        60

infection = 0.03        #初始感染率         感染率 = 感染率增长系数 * 城市规模 + 初始感染率  0.03      0.01
recovery = 0.002        #初始恢复率         恢复率 = 恢复率增长系数 * 城市规模 + 初始恢复率  0.002     0.005
attract = 0             #初始地点影响力
infection_add = 0.03         #感染率增长系数      0.03     0.01
recovery_add = 0.0015        #恢复率增长系数      0.0015    0.004

# In[4]:

#定义个体结构，包括每个人的城市位置坐标和一个状态
#S易感者；I感染者；R免疫者
class Agent:
    def __init__(self,a,b,status):
        self.position = [a,b]        #城市位置坐标
        self.status = status         #感染状态

# In[5]:

#定义城市结构，包括他的坐标、感染率和恢复率、影响力
class City:
    def __init__(self,a,b,infection,recovery,attract):
        self.position = [a,b]         #城市坐标
        self.infection = infection    #感染率
        self.recovery = recovery      #恢复率
        self.attract = attract        #影响力

#给每一个城市结构赋值,存入数组CITY
#CITY的标号=i*L+j
#此时城市规模相同
CITY = []     #将所有城市存在数组CITY里
for i in range(0,L,1):
    for j in range(0,L,1):
        CITY.append(City(i,j,infection,recovery,attract))

# In[6]:

#选择几城市，增加他们的城市规模
#提高感染率，增加人口密度
##################①情况不需要：城市规模相同，无医疗包####################
BigCity = int(0.4*L*L)              #大城市个数
# BigCity = L*L                       ####################实验一：不同感染率增长因子用
i = 0
while i<BigCity:         #随机选择BigCity个城市，增加规模：增加地点影响力，增加感染率
    x = random.randint(0,L*L-1)
    if CITY[x].attract==0:
        CITY[x].attract = CITY[x].attract+attractUP*1              #大城市地点影响力增加
        CITY[x].infection = CITY[x].infection+infection_add*1      #大城市感染率增加：感染率=感染因子*城市规模+初始感染率
        X = x//L         #规模增加的城市的横坐标
        Y = x%L          #规模增加的城市的纵坐标
        print('----------------------------')
        print('大城市坐标（%d,%d）  编号:%d'%(X,Y,x))
        i = i+1
    else: i = i
#########################################################################

# In[7]:

#初始化撒人（在L*L的网络中撒NUM个人，随机挑选n个设为感染者）
#NUM、n、L已赋值
#random.randint( , )是闭区间
def initialize(NUM,n,L):
    P_population = []                 #存放每个城市的撒人概率（与城市规模成正比）
    res = []                          #存放所有个体的信息
    sum = 0
    for i in range(L*L):              #归一化，使所有撒人概率之和为1
        P_population.append((CITY[i].attract+1)/(L*L+BigCity*attractUP))
        sum = sum+P_population[i]
    for i in range(L*L):              #计算每个城市撒人概率
        P_population[i] = P_population[i]/sum

    for i in range(NUM):              #撒人
        people_num = np.random.choice(np.arange(0,L*L),p=P_population)     #按概率随机选择个体所在城市的编号
        X_people = people_num//L      #个体所在城市的横坐标
        Y_people = people_num%L       #个体所在城市的纵坐标
        res.append(Agent(X_people,Y_people,'S'))

    for i in range(0,n,1):           #将n个人状态改为感染态I
        res[i].status = 'I'
    return res

agents_before = initialize(NUM,n,L)     #agents_before作为原始数据保存，用作画动态图。保存的是所有个体的城市坐标和状态
agents = copy.deepcopy(agents_before)   #深拷贝  agents保存的是要处理的数据

# In[8]:

##############①情况：城市规模相同，无医疗包################                    （实验二  三选一）（①情况有两个地方要注释掉In[6][11]）

##############②情况：有大城市，无医疗包####################

# In[9]:

##############③情况：有大城市，医疗包均匀分布#############                   （实验二  三选一）
print("在编号为0-%d的城市中："%(L*L-1))
for i in range(0,L*L,1):           #[0,L*L)
    CITY[i].recovery = CITY[i].recovery+recovery_add*1      #增加恢复率：恢复率=恢复率因子*资源个数+初始恢复率
    CITY[i].attract = CITY[i].attract+attractUP*1           #增加地点影响力
print('-----------------------------------')
print("有大城市，医疗包均匀分布")

# In[10]:

# ################④情况：有大城市，医疗包不均匀分布################            （实验二  三选一）（有地方要注释掉In[13]）
# print('医疗包分配给了以下几个城市：')
# for i in range(L*L):          #随机分配医疗资源（可重复分配）
#     x = random.randint(0,L*L-1)
#     CITY[x].recovery = CITY[x].recovery+recovery_add*1    #增加恢复率：恢复率=恢复率因子*资源个数+初始恢复率
#     CITY[x].attract = CITY[x].attract+attractUP*1         #增加地点影响力
#     X = x//L
#     Y = x%L
#     print('---------------------------')
#     print('城市坐标（%d,%d）  编号:%d'%(X,Y,x))

# In[11]:

#打印地点影响力分布图
print('地点影响力分布如下：')
print('-------------------------------------------------------')
for i in range(L):        #纵坐标
    for j in range(L):    #横坐标
        print('%d     '%(CITY[j*L+i].attract),end='')      #其他情况用
#         print('%.1f   '%(CITY[j*L+i].attract),end='')      #①情况用
    print('\n')

# In[12]:

#Q-Learning算法函数 得到Q表（行为-价值表）
WORLD_R = L                                 #二维世界行数
WORLD_C = L                                 #二维世界列数
ACTIONS = ['up','down','left','right','stay']      #上、下、左、右、不动
EPSILON = 0.5                               #贪婪度greedy，EPSILON的概率采取当前最优选择   0.5
ALPHA = 0.2                                 #学习率           0.2
GAMMA = 0.9                                 #可视奖励递减值   0.9
MAX_EPISODES = 5000                         #最大回合数       5000

#函数：建立一个行为价值表，Q表
def build_q_table(world_r,world_c,actions):
    k = 0
    I = np.zeros([world_r * world_c,2],int)
    for i in range(world_r):
        for j in range(world_c):
            I[k,0] = i
            I[k,1] = j
            k+=1
    I = np.transpose(I).tolist()            #DataFrame用矩阵填充index太奇怪了，它是把矩阵转置后填充的,因此要把矩阵转置一次
    table = pd.DataFrame(np.zeros((world_r * world_c, len(actions))),index=I ,columns=actions)          #table是一个二维表
    return table

#计算每个城市被选为终点的概率
Attract_all = 0                 #所有城市地点影响力之和，用于计算每个城市被选为终点的概率
for i in range(0,L*L,1):
    Attract_all = Attract_all + CITY[i].attract
P_end = []                      #每个城市被选为终点的概率
for i in range(0,L*L,1):
    if Attract_all==0:          #等概率选择终点
        P_end.append(1/(L*L))
    else:                       #有概率选择终点
        P_end.append(CITY[i].attract/Attract_all)

#函数：选择终点
def END_num():        #根据地点影响力大小，有概率地随机选择Q-Learning的终点
    END = np.random.choice(np.arange(0,L*L),p=P_end)
    return END

#函数：选择下一步行为
def choose_action(pos_x,pos_y, q_table):                            #根据当前所在位置(pos_x,pos_y)，选择下一步动作
    pos_actions = q_table.loc[(pos_x,pos_y), :]                     #[pos_x,pos_y, :]表示pos_x,pos_y那一行的所有可执行动作都列出来
    if (np.random.rand() > EPSILON) or (max(pos_actions) == 0):     #大于贪婪度或者全0，则随机选择行为
        action_name = np.random.choice(ACTIONS)
    else:                                                 #小于贪婪度，则选择最大Q值，得到它的索引（可能有多个最大值，随机选其中一个）
        action_name = np.random.choice(pos_actions[(pos_actions==max(pos_actions))].index)
    return action_name

#函数：反馈行为（下一步坐标，奖励）
def get_env_feedback(pos_x,pos_y,Action):
    if Action == 'up':                                        #行为：上
        if pos_y == 0:                                        #到顶了，无法继续up
            next_pos_x = pos_x
            next_pos_y = pos_y
            Reward = CITY[next_pos_x*L+next_pos_y].attract    #奖励值=到达的城市的地点影响力
        else:
            next_pos_y = pos_y - 1
            next_pos_x = pos_x
            Reward = CITY[next_pos_x*L+next_pos_y].attract
    elif Action == 'down':                                    #行为：下
        if pos_y == WORLD_R - 1:                              #到底了，无法继续down
            next_pos_x = pos_x
            next_pos_y = pos_y
            Reward = CITY[next_pos_x*L+next_pos_y].attract
        else:
            next_pos_y = pos_y + 1
            next_pos_x = pos_x
            Reward = CITY[next_pos_x*L+next_pos_y].attract
    elif Action == 'left':                                    #行为：左
        if pos_x == 0:                                        #到最左了，无法继续left
            next_pos_x = pos_x
            next_pos_y = pos_y
            Reward = CITY[next_pos_x*L+next_pos_y].attract
        else:
            next_pos_x = pos_x - 1
            next_pos_y = pos_y
            Reward = CITY[next_pos_x*L+next_pos_y].attract
    elif Action == 'right':                                  #行为：右
        if pos_x == WORLD_C - 1:                             #到最右了，无法继续right
            next_pos_x = pos_x
            next_pos_y = pos_y
            Reward = CITY[next_pos_x*L+next_pos_y].attract
        else:
            next_pos_x = pos_x + 1
            next_pos_y = pos_y
            Reward = CITY[next_pos_x*L+next_pos_y].attract
    else:                                                    #行为：不动
        next_pos_x = pos_x
        next_pos_y = pos_y
        Reward = CITY[next_pos_x*L+next_pos_y].attract
    return next_pos_x, next_pos_y, Reward

#主函数：强化学习  Q-learning  获得Q表
def RL():
    q_table = build_q_table(WORLD_R,WORLD_C,ACTIONS)         #实例化Q表
    ############################①情况不需要##############################
    for episode in range(MAX_EPISODES):                      #在总回合数以内，每一回合
        pos_x = random.randint(0,L-1)                        #起始位置随机
        pos_y = random.randint(0,L-1)
        is_terminated = False                                #是否结束
        END = END_num()                                      #随机选择一个终点
        end_pos_x = END//L                                   #终点的横纵坐标
        end_pos_y = END%L
        while not is_terminated:                             #没有结束时，每一步，每次结束返回for循环
            Action = choose_action(pos_x,pos_y, q_table)     #选择动作
            next_pos_x,next_pos_y,Reward = get_env_feedback(pos_x,pos_y,Action)   #获得行动反馈
            q_predict = q_table.loc[(pos_x,pos_y),Action]    #获取当前位置当前动作的Q值,loc通过标签检索，是估计值
                                                             #loc基于标签比如，loc[2,right]：第二行right列
            if next_pos_x != end_pos_x and next_pos_y != end_pos_y:    #如果下一步没有结束
                q_target = Reward + GAMMA * q_table.loc[(next_pos_x,next_pos_y), :].max()   #Q Learning关键公式1，iloc通过序号检索
                                                             #计算真实Q值
            else:                                            #下一步就结束
                q_target = Reward + GAMMA * q_table.loc[(next_pos_x,next_pos_y), :].max()
                is_terminated = True                         #结束，不再进入循环

            q_table.loc[(pos_x,pos_y), Action] += ALPHA * (q_target - q_predict)#Q Learning关键公式2
                                                             #新的Q值=老的Q值+学习率*目标与实际获得Q值的差值
            pos_x = next_pos_x                               #更新位置
            pos_y = next_pos_y
    #######################################################################
    return q_table                                          #更新Q表

#运行
if __name__ == "__main__":                                  #得Q表，并打印
    q_table = RL()
    print('行为-价值表Q-table:')  ###########################
    print(q_table)                ###########################打印Q表

# In[13]:

#函数：城市坐标（a,b）的人的下一步到达的城市坐标（按照强化学习Q表走）
SMALL = q_table.min().min()                 #Q表中的最小Q值
def RANKWALK(a,b):
    pos_actions = q_table.loc[(a,b), :]     #[a,b, :]表示a,b那一行的所有可执行动作都列出来
    POS_ACTIONS = pos_actions.sort_values(ascending=False)      #按降序排列
    POS_ACTIONS_average = 0
#     print (POS_ACTIONS)

    for i in range(0,5,1):                 #算指定坐标对应所有动作的平均值
        POS_ACTIONS_average = POS_ACTIONS_average+POS_ACTIONS[i]
    POS_ACTIONS_average = POS_ACTIONS_average/5
    POS_ACTIONS_small = POS_ACTIONS[4]     #指定坐标对应动作的最小Q值

    for i in range(0,5,1):                 #对数据进行处理
        POS_ACTIONS[i] = POS_ACTIONS[i]-0.985*POS_ACTIONS_small            ##################①②③④用
#         POS_ACTIONS[i] = POS_ACTIONS[i]-SMALL

    ACTIONS_ = POS_ACTIONS.index           #提取所有动作名称

    if max(POS_ACTIONS) != 0:
        ALL = POS_ACTIONS[0]+POS_ACTIONS[1]+POS_ACTIONS[2]+POS_ACTIONS[3]+POS_ACTIONS[4]
        P0 = POS_ACTIONS[0]/ALL            #选择第一个动作的可能性
        P1 = POS_ACTIONS[1]/ALL            #选择第二个动作的可能性
        P2 = POS_ACTIONS[2]/ALL            #选择第三个动作的可能性
        P3 = POS_ACTIONS[3]/ALL            #选择第四个动作的可能性
        P4 = POS_ACTIONS[4]/ALL            #选择第五个动作的可能性

        action_num = np.random.choice(np.arange(0,5),p=[P0,P1,P2,P3,P4])    #按照概率随机选择动作
        Action = ACTIONS_[action_num]

    else:                                  #Q值全为0时，等概率随机选择行为
        Action = np.random.choice(ACTIONS_)

    x,y,Reward = get_env_feedback(a,b,Action)      #获得行动反馈
    return [x,y]        #返回新城市的坐标

# In[14]:

#函数：对每个状态用不同颜色表示:红色感染者I，蓝色易感者S，绿色免疫者R
def colour(a):
    if a == 'S':
        colour = 'blue'
    elif a == 'I':
        colour = 'red'
    elif a == 'R':
        colour = 'lime'
    return colour;

# In[15]:

#函数：以概率p判断是否会被感染或恢复
def bernoulli(p):
    return random.random() <= p

# In[16]:

#SIR模型
def SWEEP(agents,L):
    #先所有个体移动一次
    for i in range(NUM):
        agents[i].position = RANKWALK(agents[i].position[0],agents[i].position[1])
    #再每个城市SIR感染
    for i in range(NUM):
        for j in range(i+1,NUM,1):
            if agents[i].position==agents[j].position and bernoulli(CITY[agents[i].position[0]*L+agents[i].position[1]].infection):
                #在同一个城市，且不是同一个人，且在感染率以内
                if agents[i].status == 'S' and agents[j].status == 'I':
                    agents[i].status = 'I'
                elif agents[i].status == 'I' and agents[j].status == 'S':
                    agents[j].status = 'I'
            elif agents[i].position==agents[j].position and bernoulli(CITY[agents[i].position[0]*L+agents[i].position[1]].recovery):
                #在同一个城市，且在恢复率以内
                if agents[i].status == 'I':
                    agents[i].status = 'R'
                if agents[j].status == 'I':
                    agents[j].status = 'R'
    return agents

# In[17]:

#函数：S，I，R个数统计（转化为占比）
def countSIR(agents):
    result_S = 0
    result_I = 0
    result_R = 0
    for i in range(0,NUM,1):
        if agents[i].status == 'S':
            result_S = result_S+1
        elif agents[i].status == 'I':
            result_I = result_I+1
        elif agents[i].status == 'R':
            result_R = result_R+1
    return(result_S/NUM,result_I/NUM,result_R/NUM)

# print(countSIR(agents)[0])

# In[18]:

#函数：可视化（散点图）
def visualize(agents,L,title):
    for i in range(0,NUM,1):
        try:
            x = agents[i].position[0]*CitySize+random.uniform(3,CitySize-3)
            y = -agents[i].position[1]*CitySize-random.uniform(3,CitySize-3)
            plt.scatter(x,y,s=70,c=colour(agents[i].status),marker='o')
        except:     #排除重合的可能
            continue
    plt.title(title)
    plt.xlim(-1, L*CitySize+1)                    #横纵坐标范围
    plt.ylim(-L*CitySize-1,1)
    plt.xticks(fontsize=15)                       #横纵坐标字体大小
    plt.yticks(fontsize=15)
    x_major_locator=MultipleLocator(CitySize)     #设置x轴坐标的刻度间隔
    y_major_locator=MultipleLocator(CitySize)     #设置y轴坐标的刻度间隔
    ax=plt.gca()
    ax.xaxis.set_major_locator(x_major_locator)
    ax.yaxis.set_major_locator(y_major_locator)
    plt.grid()       #画主刻度网格
    plt.show

# In[26]:

###########################################################################
#个体空间移动轨迹的可视化
#初始城市坐标是（0，0）
get_ipython().run_line_magic('matplotlib', 'inline')
w=100                #随机移动200步
x = np.zeros(w)      #定义x和y的零数组
y = np.zeros(w)
X = np.zeros(w)
Y = np.zeros(w)
x[0] = 0
y[0] = 0
for i in range(1,w,1):       #一个人随机游走w步的城市坐标的轨迹
    x[i] = RANKWALK(int(x[i-1]),int(y[i-1]))[0]
    y[i] = RANKWALK(int(x[i-1]),int(y[i-1]))[1]
for i in range(0,w,1):
    X[i] = x[i]*CitySize+random.uniform(3,CitySize-3)
    Y[i] = -y[i]*CitySize-random.uniform(3,CitySize-3)
plt.plot(X,Y,'r--')
plt.title('个体移动%d次的轨迹'%w)
plt.xlim(-1, L*CitySize+1)
plt.ylim( -L*CitySize-1,1)
x_major_locator=MultipleLocator(CitySize)     #设置x轴坐标的刻度间隔
y_major_locator=MultipleLocator(CitySize)     #设置y轴坐标的刻度间隔
ax=plt.gca()
ax.xaxis.set_major_locator(x_major_locator)
ax.yaxis.set_major_locator(y_major_locator)
plt.grid()       #画主刻度网格
plt.show
############################################################################

# In[20]:

##########################画SIR迭代多次的曲线图#############################
#迭代次数time已赋值
get_ipython().run_line_magic('matplotlib', 'inline')

NUM_S = np.zeros(N)
NUM_I = np.zeros(N)
NUM_R = np.zeros(N)

for i in range(0,time,1):      #迭代time次，画光滑曲线
    result_S = []
    result_I = []
    result_R = []
    agents = initialize(NUM,n,L)    #画完之后再随机生成一组新的数据
    for t in range(0,N,1):
        if t==0:
            result_I.append(n/NUM)
            result_R.append(0/NUM)
            result_S.append((NUM-n)/NUM)
        else:
            agents = SWEEP(agents,L)
            result_SIR = countSIR(agents)
            result_S.append(result_SIR[0])
            result_I.append(result_SIR[1])
            result_R.append(result_SIR[2])

        NUM_S[t] = NUM_S[t]+result_S[t]    #time次 S个数之和
        NUM_I[t] = NUM_I[t]+result_I[t]
        NUM_R[t] = NUM_R[t]+result_R[t]

IMAX_t = 0         #I态波峰到达时间
IMAX_n = 0         #I态波峰人数
RMAX = 0           #感染规模

for i in range(0,N,1):
    NUM_S[i] = NUM_S[i]/time      #取S\I\R个数的平均值
    NUM_I[i] = NUM_I[i]/time
    NUM_R[i] = NUM_R[i]/time
    if IMAX_n<=NUM_I[i]:
        IMAX_n = NUM_I[i]         #求波峰
        IMAX_t = i
RMAX = NUM_R[N-1]                 #求感染规模

#画图
xS = range(N)
xI = range(N)
xR = range(N)
plt.xlim(-1, N+1)
plt.ylim(-0.05,1.05)
plt.xlabel('传播时间t')
plt.ylabel('S\I\R态人数占比')
y_major_locator=MultipleLocator(0.1)          #设置y轴坐标的刻度间隔
ax=plt.gca()
ax.yaxis.set_major_locator(y_major_locator)
plt.plot(xS,NUM_S,label='S',linewidth=1.2,color='blue')
plt.plot(xI,NUM_I,label='I',linewidth=1.2,color='red')
plt.plot(xR,NUM_R,label='R',linewidth=1.2,color='lime')
plt.legend()
plt.show()

#指标输出
print("感染规模占比：%.4f"%(RMAX))
print("I态波峰到达时间：t = %d"%(IMAX_t))
print("I态波峰人数占比：%.4f"%(IMAX_n))

#存入文件，为实验二的补充实验做准备
filename1 = 'Infection_scale.txt'     #将感染规模存入文件，用于实验二补充
with open(filename1,'a') as f1:
    f1.write('%.4f\n'%RMAX)

filename2 = 'Peak_time.txt'           #波峰到达时间存入文件
with open(filename2,'a') as f2:
    f2.write('%d\n'%IMAX_t)

filename3 = 'Peak_num.txt'            #波峰人数存入文件
with open(filename3,'a') as f3:
    f3.write('%.4f\n'%IMAX_n)

# In[21]:

# #####################画散点动图和SIR动图（传播过程可视化）########################
# agents = initialize(NUM,n,L)    #随机生成一组新的数据，用来覆盖原来的数据

# result_S = []
# result_I = []
# result_R = []

# %matplotlib auto

# for t in range(0,N,1):
#     if t==0:
#         result_I.append(n/NUM)
#         result_R.append(0/NUM)
#         result_S.append((NUM-n)/NUM)
#     else:
#         agents = SWEEP(agents,L)
#         result_SIR = countSIR(agents)
#         result_S.append(result_SIR[0])
#         result_I.append(result_SIR[1])
#         result_R.append(result_SIR[2])

#     plt.clf()                   #清空画布
#     plt.subplot(121)            #第一张图
#     visualize(agents,L,'某一次传播过程')
#     plt.subplot(122)            #第二张图
#     plt.xlim(-5, N+5)
#     plt.ylim(-0.05,1.05)
#     y_major_locator=MultipleLocator(0.1)          #设置y轴坐标的刻度间隔
#     ax=plt.gca()
#     ax.yaxis.set_major_locator(y_major_locator)
#     plt.xticks(fontsize=15)
#     plt.yticks(fontsize=15)
#     plt.xlabel('传播时间t')     #x轴标题
#     plt.ylabel('S\I\R态人数占比')
#     xS = range(len(result_S))
#     xI = range(len(result_I))
#     xR = range(len(result_R))
#     plt.plot(xS,result_S,label='易感者S : %.2f'%result_S[t],linewidth=1.7,color='blue')
#     plt.plot(xI,result_I,label='感染者I : %.2f'%result_I[t],linewidth=1.7,color='red')
#     plt.plot(xR,result_R,label='恢复者R : %.2f'%result_R[t],linewidth=1.7,color='lime')
#     plt.rcParams.update({'font.size': 15})     #改变图例大小
#     plt.legend(loc='upper right')
#     plt.pause(0.01)
# plt.show()

# In[22]:

# ######################实验一：改变恢复率增长系数，画图##########################
# ###用③情况
# ###In[6]中①不需要那段注释掉
# ###In[11]中①不需要那段注释掉
# infection = 0.005           #基本感染率
# recovery = 0.003            #基本恢复率
# number = 8                      #因子改变次数
# recovery_ADD = 0                #初始恢复率因子
# colour = ['red','darkorange','gold','lime','cyan','dodgerblue','blueviolet','indigo']
# NUM_I = np.zeros((number,N))         #画图用的二维数组
# NUM_R = np.zeros((number,N))
# mu = np.zeros(number)                #记录每次系数大小
# I_TIME = np.zeros(number)            #记录每次波峰到达时间
# S_SIZE = np.zeros(number)            #记录每次感染规模
# I_NUM = np.zeros(number)             #记录每次波峰人数
# for a in range(number):              #改变系数8次
#     for i in range(L*L):
#         CITY[i].recovery = recovery_ADD*1+recovery    #新的恢复率
#     mu[a] = recovery_ADD                 #记录各个恢复率增长系数
#     recovery_ADD = recovery_ADD+0.001    #每次系数增加
#     for i in range(0,time,1):            #迭代time次，画光滑曲线
#         result_I = []
#         result_R = []
#         agents = initialize(NUM,n,L)     #画完之后再随机生成一组新的数据
#         for t in range(0,N,1):
#             if t==0:
#                 result_I.append(n/NUM)
#                 result_R.append(0/NUM)
#             else:
#                 agents = SWEEP(agents,L)
#                 result_SIR = countSIR(agents)
#                 result_I.append(result_SIR[1])
#                 result_R.append(result_SIR[2])

#             NUM_I[a,t] = NUM_I[a,t]+result_I[t]        #time次之和
#             NUM_R[a,t] = NUM_R[a,t]+result_R[t]

#     IMAX_t = 0         #I态波峰到达时间
#     IMAX_n = 0         #I态波峰人数

#     for i in range(0,N,1):
#         NUM_I[a,i] = NUM_I[a,i]/time         #取S\I\R平均值
#         NUM_R[a,i] = NUM_R[a,i]/time
#         if IMAX_n<=NUM_I[a,i]:
#             IMAX_n = NUM_I[a,i]       #求波峰
#             IMAX_t = i
#     I_TIME[a] = IMAX_t                #波峰时间
#     S_SIZE[a] = NUM_R[a,N-1]          #感染规模
#     I_NUM[a] = IMAX_n                    #波峰人数

# #画不同恢复率增长系数下的感染规模
# plt.figure(dpi=100,figsize=(5,5))      #调整画布大小
# plt.title('不同恢复率增长系数下的感染规模情况')
# for i in range(number):
#     xR = range(N)
#     plt.xlim(-5, N+5)                  #x轴刻度
#     plt.ylim(-0.05,1.05)
#     y_major_locator=MultipleLocator(0.1)          #设置y轴坐标的刻度间隔
#     ax=plt.gca()
#     ax.yaxis.set_major_locator(y_major_locator)
#     plt.xlabel('传播时间t')            #x轴标题
#     plt.ylabel('R态免疫人数占比')
#     plt.plot(xR,NUM_R[i],label='恢复率增长系数=%.3f'%mu[i],linewidth=1.2,color=colour[i])
#     plt.legend()                       #画图例
# plt.show()

# #画不同恢复率增长系数下的I态波峰
# plt.figure(dpi=100,figsize=(10,5))
# plt.title('不同恢复率增长系数下的I态波峰情况')
# for i in range(number):
#     xR = range(N)
#     plt.xlim(-5, N+5)
#     plt.ylim(-0.05,0.45)
#     y_major_locator=MultipleLocator(0.1)   #设置y轴坐标的刻度间隔
#     ax=plt.gca()
#     ax.yaxis.set_major_locator(y_major_locator)
#     plt.xlabel('传播时间t')
#     plt.ylabel('I态感染人数占比')
#     plt.plot(xR,NUM_I[i],label='恢复率增长系数=%.3f'%mu[i],linewidth=1.2,color=colour[i])
#     plt.legend()
# for i in range(number):                #标出波峰点
#     plt.scatter(I_TIME[i],I_NUM[i],s=20,c=colour[i],marker='x')
# plt.show()

# #输出数据
# print("恢复率增长系数：|",end='')
# for i in range(number):
#     print(' %.3f |'%mu[i],end='')
# print('\n-----------------------------------------------------------------------------------------')
# print("  感染规模占比：|",end='')
# for i in range(number):
#     print(' %.3f |'%S_SIZE[i],end='')
# print('\n')
# print("  波峰到达时间：|",end='')
# for i in range(number):
#     print(' %-5d |'%I_TIME[i],end='')
# print('\n')
# print("  波峰人数占比：|",end='')
# for i in range(number):
#     print(' %.3f |'%I_NUM[i],end='')

# In[23]:

# #####实验一：改变感染率增长系数，画图##########################
# ###用②情况，
# ###In[6]中的BigCity要改为L*L
# ###In[11]中①不需要那段注释掉
# infection = 0.003           #基本感染率     0.006（大）    0.003(2)
# recovery = 0.001            #基本恢复率     0.005（大）    0.001（2）
# infection_ADD = 0           #初始感染率增长系数
# number = 8                  #系数改变次数
# colour = ['red','darkorange','gold','lime','cyan','dodgerblue','blueviolet','indigo']
# NUM_I = np.zeros((number,N))         #画图用的二维数组
# NUM_R = np.zeros((number,N))
# beita = np.zeros(number)                #记录每次因子大小
# I_TIME = np.zeros(number)            #记录每次波峰到达时间
# S_SIZE = np.zeros(number)            #记录每次感染规模
# I_NUM = np.zeros(number)             #记录每次波峰人数
# for a in range(number):              #改变因子6次
#     for i in range(L*L):
#         CITY[i].infection = infection_ADD*1+infection    #新的感染率
#     beita[a] = infection_ADD                  #记录各个感染率增长系数
#     infection_ADD = infection_ADD+0.002       #每次系数增加    0.001（大）    0.002（2）
#     for i in range(0,time,1):            #迭代time次，画光滑曲线
#         result_I = []
#         result_R = []
#         agents = initialize(NUM,n,L)     #画完之后再随机生成一组新的数据
#         for t in range(0,N,1):
#             if t==0:
#                 result_I.append(n/NUM)
#                 result_R.append(0/NUM)
#             else:
#                 agents = SWEEP(agents,L)
#                 result_SIR = countSIR(agents)
#                 result_I.append(result_SIR[1])
#                 result_R.append(result_SIR[2])

#             NUM_I[a,t] = NUM_I[a,t]+result_I[t]        #time次之和
#             NUM_R[a,t] = NUM_R[a,t]+result_R[t]

#     IMAX_t = 0         #I态波峰到达时间
#     IMAX_n = 0         #I态波峰人数

#     for i in range(0,N,1):
#         NUM_I[a,i] = NUM_I[a,i]/time         #取S\I\R平均值
#         NUM_R[a,i] = NUM_R[a,i]/time
#         if IMAX_n<=NUM_I[a,i]:
#             IMAX_n = NUM_I[a,i]       #求波峰
#             IMAX_t = i
#     I_TIME[a] = IMAX_t                #波峰时间
#     S_SIZE[a] = NUM_R[a,N-1]          #感染规模
#     I_NUM[a] = IMAX_n                    #波峰人数

# #画不同系数下的感染规模
# plt.figure(dpi=100,figsize=(5,5))      #调整画布大小
# plt.title('不同感染率增长系数下的感染规模情况')
# for i in range(number):
#     xR = range(N)
#     plt.xlim(-5, N+5)                  #x轴刻度
#     plt.ylim(-0.05,1.05)
#     y_major_locator=MultipleLocator(0.1)          #设置y轴坐标的刻度间隔
#     ax=plt.gca()
#     ax.yaxis.set_major_locator(y_major_locator)
#     plt.xlabel('传播时间t')            #x轴标题
#     plt.ylabel('R态恢复人数占比')
#     plt.plot(xR,NUM_R[i],label='感染率增长系数=%.3f'%beita[i],linewidth=1.2,color=colour[i])
#     plt.legend()                       #画图例
# plt.show()

# #画不同系数下的I态波峰
# plt.figure(dpi=100,figsize=(10,5))
# plt.title('不同感染率增长系数下的I态波峰情况')
# for i in range(number):
#     xR = range(N)
#     plt.xlim(-5, N+5)
#     plt.ylim(-0.05,0.4)
#     y_major_locator=MultipleLocator(0.1)          #设置y轴坐标的刻度间隔
#     ax=plt.gca()
#     ax.yaxis.set_major_locator(y_major_locator)
#     plt.xlabel('传播时间t')
#     plt.ylabel('I态感染人数占比')
#     plt.plot(xR,NUM_I[i],label='感染率增长系数=%.3f'%beita[i],linewidth=1.2,color=colour[i])
#     plt.legend()
# for i in range(number):
#     plt.scatter(I_TIME[i],I_NUM[i],s=20,c=colour[i],marker='x')
# plt.show()

# #输出数据
# print("感染率增长系数：|",end='')
# for i in range(number):
#     print(' %.3f |'%beita[i],end='')
# print('\n----------------------------------------------------------------------------------')
# print("  感染规模占比：|",end='')
# for i in range(number):
#     print(' %.3f |'%S_SIZE[i],end='')
# print('\n')
# print("  波峰到达时间：|",end='')
# for i in range(number):
#     print(' %-5d |'%I_TIME[i],end='')
# print('\n')
# print("  波峰人数占比：|",end='')
# for i in range(number):
#     print(' %.3f |'%I_NUM[i],end='')

# In[24]:

# #######################清空文件内容#########################
# filename1 = 'Infection_scale.txt'
# filename2 = 'Peak_time.txt'
# filename3 = 'Peak_num.txt'
# with open(filename1,'r+') as file:    #清空
#     file.truncate(0)
# with open(filename2,'r+') as file:    #清空
#     file.truncate(0)
# with open(filename3,'r+') as file:    #清空
#     file.truncate(0)
# ############################################################

# In[25]:

# ####################实验二补充：四种情况规模、波峰时间人数相对减少量#######################
# #################注意：先清空文件内容，再按顺序分别运行①②③④种情况#########
# filename1 = 'Infection_scale.txt'
# filename2 = 'Peak_time.txt'
# filename3 = 'Peak_num.txt'
# data_n = 4           #数据数量（实验二）
# # data_n = 8           #数据数量(恢复率)（感染率）
# #提取文件数据
# with open(filename1) as f1:                            #提取文件1里的感染规模数据，存入InfectionScale
#     InfectionScale = f1.readlines()
#     InfectionScale = list(map(float,InfectionScale))   #list格式转换为int格式
# with open(filename2) as f2:                            #提取文件2里的波峰时间数据，存入PeakTime
#     PeakTime = f2.readlines()
#     PeakTime= list(map(int,PeakTime))
# with open(filename3) as f3:                            #提取文件3里的波峰人数数据，存入PeakNum
#     PeakNum = f3.readlines()
#     PeakNum= list(map(float,PeakNum))

# #计算相对增加量
# RelativeInfectionScale = np.zeros(data_n-1)            #感染规模相对增减量
# for i in range(data_n-1):
#     RelativeInfectionScale[i] = InfectionScale[i+1]-InfectionScale[i]
# RelativePeakTime = np.zeros(data_n-1)                 #波峰到达时间相对增加量
# for i in range(data_n-1):
#     RelativePeakTime[i] = PeakTime[i+1]-PeakTime[i]
# RelativePeakNum = np.zeros(data_n-1)                  #波峰人数相对增加量
# for i in range(data_n-1):
#     RelativePeakNum[i] = PeakNum[i+1]-PeakNum[i]

# #画图:感染规模相对增加
# plt.figure(dpi=90,figsize=(4,5))                     #调整画布大小（实验二）
# # plt.figure(dpi=90,figsize=(6,5))                     #调整画布大小（恢复率）（感染率）
# plt.ylim(-0.4,0.4)                                   #y轴范围（实验二）
# # plt.ylim(-0.25,0.05)                                 #y轴范围（恢复率）
# # plt.ylim(-0.1,0.3)                                   #y轴范围（感染率）
# plt.xlim(-0.5,data_n-1.5)
# x_major_locator=MultipleLocator(1)
# y_major_locator=MultipleLocator(0.1)                 #设置y轴坐标的刻度间隔（实验二）
# # y_major_locator=MultipleLocator(0.05)                #设置y轴坐标的刻度间隔（恢复率）
# # y_major_locator=MultipleLocator(0.05)                #设置y轴坐标的刻度间隔（感染率）
# ax=plt.gca()
# ax.xaxis.set_major_locator(x_major_locator)
# ax.yaxis.set_major_locator(y_major_locator)
# x_name = ['②-①','③-②','④-③']                #x轴刻度文字 (实验二)
# # x_name = ['△γμ1','△γμ2','△γμ3','△γμ4','△γμ5','△γμ6','△γμ7']      #x轴刻度文字（恢复率）
# # x_name = ['△γβ1','△γβ2','△γβ3','△γβ4','△γβ5','△γβ6','△γβ7']      #x轴刻度文字（感染率）
# _ = plt.xticks(np.arange(0,data_n-1),x_name)
# plt.axhline(0, color='red', linestyle='--')         #y=0基准线
# plt.title('四种情况下的感染规模相对增加量')                 #（实验二）
# # plt.title('相邻恢复率增长系数下的感染规模相对增加量')       #（恢复率）
# # plt.title('相邻感染率增长系数下的感染规模占比相对增加量')       #（感染率）
# plt.ylabel('△Rmax/Nsum')
# for i in range(data_n-1):                           #柱状图上标数字
#     plt.text(i,RelativeInfectionScale[i],'%.4f'%RelativeInfectionScale[i],ha='center',va='center')
# plt.bar(range(data_n-1),RelativeInfectionScale,0.5,color="seagreen")
# plt.show()

# #画图:波峰时间相对增加
# plt.figure(dpi=90,figsize=(4,5))                  #调整画布大小（实验二）
# # plt.figure(dpi=90,figsize=(6,5))                  #调整画布大小（恢复率）（感染率）
# plt.ylim(-20,10)                                  #y轴范围（实验二）
# # plt.ylim(-7,2)                                    #y轴范围（恢复率）
# # plt.ylim(-3,8)                                   #y轴范围（感染率）
# plt.xlim(-0.5,data_n-1.5)
# y_major_locator=MultipleLocator(5)              #设置y轴坐标的刻度间隔（实验二）
# # y_major_locator=MultipleLocator(1)               #设置y轴坐标的刻度间隔（恢复率）（感染率）
# ax=plt.gca()
# ax.xaxis.set_major_locator(x_major_locator)
# ax.yaxis.set_major_locator(y_major_locator)
# x_name = ['②-①','③-②','④-③']             #x轴刻度文字（实验二）
# # x_name = ['△γμ1','△γμ2','△γμ3','△γμ4','△γμ5','△γμ6','△γμ7']     #x轴刻度文字（恢复率）
# # x_name = ['△γβ1','△γβ2','△γβ3','△γβ4','△γβ5','△γβ6','△γβ7']      #x轴刻度文字（感染率）
# _ = plt.xticks(np.arange(0,data_n-1),x_name)
# plt.axhline(0, color='red', linestyle='--')    #y=0基准线
# plt.title('四种情况下的波峰到达时间相对增加量')                  #（实验二）
# # plt.title('不同恢复率增长系数下的波峰到达时间相对增加量')        #（恢复率）
# # plt.title('不同感染率增长系数下的波峰到达时间相对增加量')        #（感染率）
# plt.ylabel('波峰到达时间相对增加量')
# for i in range(data_n-1):                             #柱状图上标数字
#     plt.text(i,RelativePeakTime[i],'%d'%RelativePeakTime[i],ha='center',va='center')
# plt.bar(range(data_n-1),RelativePeakTime,0.5,color="seagreen")
# plt.show()

# #画图:波峰人数相对增加
# plt.figure(dpi=90,figsize=(4,5))                 #调整画布大小（实验二）
# # plt.figure(dpi=90,figsize=(6,5))                 #调整画布大小（恢复率）（感染率）
# plt.ylim(-0.15,0.1)                               #y轴范围 （实验二）
# # plt.ylim(-0.2,0.05)                              #y轴范围 （恢复率）
# # plt.ylim(-0.05,0.1)                              #y轴范围 （感染率）
# plt.xlim(-0.5,data_n-1.5)
# y_major_locator=MultipleLocator(0.05)             #设置y轴坐标的刻度间隔（实验二）
# # y_major_locator=MultipleLocator(0.05)            #设置y轴坐标的刻度间隔（恢复率）
# # y_major_locator=MultipleLocator(0.01)            #设置y轴坐标的刻度间隔（感染率）
# x_major_locator=MultipleLocator(1)
# ax=plt.gca()
# ax.xaxis.set_major_locator(x_major_locator)
# ax.yaxis.set_major_locator(y_major_locator)
# x_name = ['②-①','③-②','④-③']                #x轴刻度文字（实验二）
# # x_name = ['△γμ1','△γμ2','△γμ3','△γμ4','△γμ5','△γμ6','△γμ7']     #x轴刻度文字（恢复率）
# # x_name = ['△γβ1','△γβ2','△γβ3','△γβ4','△γβ5','△γβ6','△γβ7']      #x轴刻度文字（感染率）
# _ = plt.xticks(np.arange(0,data_n-1),x_name)
# plt.axhline(0, color='red', linestyle='--')       #y=0基准线
# plt.title('四种情况下的波峰人数相对增加量')                  #（实验二）
# # plt.title('相邻恢复率增长系数下的波峰人数相对增加量')        #（恢复率）
# # plt.title('相邻感染率增长系数下的波峰人数占比相对增加量')        #（感染率）
# plt.ylabel('△Ipeak/Nsum')
# for i in range(data_n-1):                            #柱状图上标数字
#     plt.text(i,RelativePeakNum[i],'%.4f'%RelativePeakNum[i],ha='center',va='center')
# plt.bar(range(data_n-1),RelativePeakNum,0.5,color="seagreen")
# plt.show()


"""Module defining base FHI-aims input set and generator."""
from __future__ import annotations

import copy
import json
import logging
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, Iterable, List, Sequence, Tuple
from warnings import warn

import numpy as np
from ase.cell import Cell
from ase.calculators.aims import AimsTemplate
from monty.json import MontyDecoder, MontyEncoder
from pymatgen.io.core import InputFile, InputGenerator, InputSet

from atomate2_temp.aims.io.parsers import AimsParseError, read_aims_output
from atomate2_temp.aims.utils.common import (
    CONTROL_FILE_NAME,
    GEOMETRY_FILE_NAME,
    PARAMS_JSON_FILE_NAME,
    TMPDIR_NAME,
    cwd,
)
from atomate2_temp.aims.utils.msonable_atoms import MSONableAtoms

DEFAULT_AIMS_PROPERTIES = [
    "energy",
    "free_energy",
    "forces",
    "stress",
    "stresses",
    "dipole",
    "magmom",
]

__all__ = ["AimsInputSet", "AimsInputGenerator"]

logger = logging.getLogger(__name__)

@dataclass
class AimsInputFile(InputFile):
    """The input file for an FHI-aims calculation

    Parameters
    ----------
    _content_str: str
        The contents of the input file as a string
    """

    _content_str: str = ""

    def get_string(self) -> str:
        """The contents of the input file

        Returns
        -------
        The contents of the input file
        """
        return self._content_str

    def get_str(self) -> str:
        """The contents of the input file

        Returns
        -------
        The contents of the input file
        """
        return self._content_str

    @classmethod
    def from_string(cls, contents: str):
        """Create an input file from the contents string

        Parameters
        ----------
        contents: str
            The contents of the input file
        """
        return cls(contents)

    @classmethod
    def from_str(cls, contents: str):
        """Create an input file from the contents string

        Parameters
        ----------
        contents: str
            The contents of the input file
        """
        return cls(contents)

class AimsInputSet(InputSet):
    """
    A class to represent a set of Aims inputs.
    """

    def __init__(
        self,
        parameters: Dict[str, Any],
        atoms: MSONableAtoms,
        properties: Sequence[str] = ("energy", "free_energy"),
    ):
        """Constructor

        Parameters
        ----------
        parameters: Dict[str, Any]
            The ASE parameters object for the calculation
        atoms: MSONableAtoms
            The atoms objects to create the inputs for
        properties: Sequence[str]
            The properties to calculate for the calculation
        """
        self._parameters = parameters
        self._atoms = MSONableAtoms(atoms)
        self._properties = properties

        aims_control_in, aims_geometry_in = self.get_input_files()
        super().__init__(
            inputs={
                CONTROL_FILE_NAME: aims_control_in,
                GEOMETRY_FILE_NAME: aims_geometry_in,
                PARAMS_JSON_FILE_NAME: json.dumps(
                    self._parameters, indent=2, cls=MontyEncoder
                ),
            }
        )

    def get_input_files(self) -> Tuple(str, str):
        """Get the input file contents for the calculation

        Returns
        -------
        The contents of the control.in and geometry.in file
        """
        with cwd(TMPDIR_NAME, mkdir=True, rmdir=True):
            aims_template = AimsTemplate()
            aims_template.write_input(
                Path("./"), self._atoms, self._parameters, self._properties
            )
            aims_control_in = AimsInputFile.from_file("control.in")
            aims_geometry_in = AimsInputFile.from_file("geometry.in")
        return aims_control_in, aims_geometry_in

    @property
    def control_in(self) -> str:
        """Get the AimsInput object."""
        return self[CONTROL_FILE_NAME]

    @property
    def geometry_in(self) -> str:
        """Get the AimsInput object."""
        return self[GEOMETRY_FILE_NAME]

    @property
    def parameters_json(self) -> str:
        """Get the AimsInput object."""
        return self[PARAMS_JSON_FILE_NAME]

    def set_parameters(self, *args, **kwargs) -> Dict[str, Any]:
        """Set the parameters object for the AimsTemplate

        This sets the parameters object that is passed to an AimsTempalte and
        resets the control.in file

        One can pass a dictionary mapping the aims variables to their values or
        the aims variables as keyword arguments. A combination of the two
        options is also allowed.

        Returns
        -------
        dict
            dictionary with the variables that have been added.
        """
        self._parameters.clear()
        for arg in args:
            self._parameters.update(arg)

        self._parameters.update(kwargs)

        aims_control_in, _ = self.get_input_files()
        self.inputs[CONTROL_FILE_NAME] = aims_control_in
        self.inputs[PARAMS_JSON_FILE_NAME] = json.dumps(
            self._parameters, indent=2, cls=MontyEncoder
        )
        self.__dict__.update(self.inputs)

        return self._parameters

    def remove_parameters(
        self, keys: Iterable[str] | str, strict: bool = True
    ) -> Dict[str, Any]:
        """Remove the aims parameters listed in keys.

        This removes the aims variables from the parameters object.

        Parameters
        ----------
        keys
            string or list of strings with the names of the aims parameters
            to be removed.
        strict
            whether to raise a KeyError if one of the aims parameters to be
            removed is not present.

        Returns
        -------
        Dictionary with the variables that have been removed.
        """
        if isinstance(keys, str):
            keys = [keys]
        for key in keys:
            if strict and key not in self._parameters:
                raise ValueError(f"The key ({key}) is not in self._parameters")
            elif key not in self._parameters:
                continue
            del self._parameters[key]

        return self.set_parameters(**self._parameters)

    def set_atoms(self, atoms: MSONableAtoms):
        """Set the atoms object for this input set."""
        self._atoms = MSONableAtoms(atoms)

        _, aims_geometry_in = self.get_input_files()
        self.inputs[GEOMETRY_FILE_NAME] = aims_geometry_in
        self.__dict__.update(self.inputs)

    def deepcopy(self):
        """Deep copy of the input set."""
        return copy.deepcopy(self)

@dataclass
class AimsInputGenerator(InputGenerator):
    """
    A class to generate Aims input sets.

    Parameters
    ----------
    user_parameters: Dict[str, Any]
        Updates the default parameters for the FHI-aims calculator
    user_kpoints_settings: Dict[str, Any]
        The settings used to create the k-grid parameters for FHI-aims
    """

    user_parameters: Dict[str, Any] = field(default_factory=dict)
    user_kpoints_settings: Dict[str, Any] = field(default_factory=dict)

    def get_input_set(  # type: ignore
        self,
        atoms: MSONableAtoms = None,
        prev_dir: str | Path = None,
        properties: List[str] | Tuple[str] = None,
    ) -> AimsInputSet:
        """Generate an AimsInputSet object.

        Parameters
        ----------
        atoms : MSONableAtoms
            ASE Atoms object to generate the input set for.
        prev_dir: str or Path
            Path to the previous working directory
        properties: list or tuple of str
            System properties that are being calculated

        Returns
        -------
        The input set for the calculation of atoms
        """
        prev_atoms, prev_parameters, prev_results = self._read_previous(prev_dir)
        atoms = atoms if atoms is not None else prev_atoms

        parameters = self._get_input_parameters(atoms, prev_parameters)
        properties = self._get_properties(properties, parameters, prev_results)

        return AimsInputSet(parameters=parameters, atoms=atoms, properties=properties)

    def _read_previous(
        self, prev_dir: str | Path = None
    ) -> tuple[MSONableAtoms, Dict[str, Any], Dict[str, Iterable[float]]]:
        """Read in previous results

        Parameters
        ----------
        prev_dir: str or Path
            The previous directory for the calculation
        """
        prev_atoms = None
        prev_parameters = {}
        prev_results = {}

        if prev_dir:
            # strip hostname from the directory (not good, works only with run_locally.
            # Should be checked with Fireworks, will not for sure work with
            # jobflow_remote)
            prev_dir = prev_dir.split(":")[-1]
            prev_parameters = json.load(
                open(f"{prev_dir}/parameters.json", "rt"), cls=MontyDecoder
            )
            try:
                prev_atoms = read_aims_output(f"{prev_dir}/aims.out")
                prev_results = prev_atoms.calc.results
            except (IndexError, AimsParseError):
                pass

        return prev_atoms, prev_parameters, prev_results

    def _get_properties(
        self,
        properties: Iterable[str] = None,
        parameters: Dict[str, Any] = None,
        prev_results: Dict[str, Iterable[float]] = None,
    ) -> Iterable[str]:
        """Get the properties to calculate

        Parameters
        ----------
        properties
            The currently requested properties
        parameters
            The parameters for this calculation
        prev_results
            The previous calculation results

        Returns
        -------
        The list of properties to calculate
        """
        if properties is None:
            properties = ["energy", "free_energy"]

        for key in prev_results.keys():
            if key not in properties and key in DEFAULT_AIMS_PROPERTIES:
                properties.append(key)

        if "compute_forces" in parameters and "forces" not in properties:
            properties.append("forces")
        if "compute_heat_flux" in parameters and "stresses" not in properties:
            properties.append("stress")
            properties.append("stresses")
        if "stress" not in properties and (
            ("compute_analytical_stress" in parameters)
            or ("compute_numerical_stress" in parameters)
            or ("compute_heat_flux" in parameters)
        ):
            properties.append("stress")

        return properties

    def _get_input_parameters(
        self, atoms: MSONableAtoms, prev_parameters: Dict[str, Any] = None
    ) -> Dict[str, Any]:
        """Create the input parameters

        Parameters
        ----------
        atoms: MSONableAtoms
            The atoms object for the structures
        prev_parameters
            The previous calculation's calculation parameters

        Returns
        -------
        The input object
        """

        # Get the default configuration
        # FHI-aims recommends using their defaults so bare-bones default parameters
        parameters = {
            "xc": "pbe",
            "relativistic": "atomic_zora scalar",
        }

        # Override default parameters with previous parameters
        prev_parameters = (
            {} if prev_parameters is None else copy.deepcopy(prev_parameters)
        )
        prev_parameters.pop("relax_geometry", None)
        prev_parameters.pop("relax_unit_cell", None)

        kpt_settings = copy.deepcopy(self.user_kpoints_settings)
        if "k_grid" in prev_parameters:
            density = self.k2d(atoms, prev_parameters.pop("k_grid"))
            if "density" not in kpt_settings:
                kpt_settings["density"] = density

        parameter_updates = self.get_parameter_updates(atoms, prev_parameters)
        parameters = recursive_update(parameters, parameter_updates)

        # Override default parameters with user_parameters
        parameters = recursive_update(parameters, self.user_parameters)
        if ("k_grid" in parameters) and ("density" in kpt_settings):
            warn(
                "WARNING: the k_grid is set in user_parameters and in the kpt_settings, using the one passed in user_parameters."
            )
        elif np.any(atoms.pbc) and ("k_grid" not in parameters):
            density = kpt_settings.get("density", 5.0)
            even = kpt_settings.get("even", True)
            parameters["k_grid"] = self.d2k(atoms, density, even)
        elif not np.any(atoms.pbc) and "k_grid" in parameters:
            warn("WARNING: removing unnecessary k_grid information")
            del parameters["k_grid"]

        return parameters

    def get_parameter_updates(
        self, atoms: MSONableAtoms, prev_parameters: Dict[str, Any]
    ) -> Dict[str, Any]:
        """
        Updates the parameters for a given calculation type

        Parameters
        ----------
        atoms : MSONableAtoms
            ASE Atoms object.
        prev_parameters: Dict[str, Any]
            Previous calculation parameters.

        Returns
        -------
        A dictionary of updates to apply.
        """
        return prev_parameters

    def d2k(
        self,
        atoms: MSONableAtoms,
        kptdensity: float | Iterable[float] = 5.0,
        even: bool = True,
    ) -> Iterable[float]:
        """Convert k-point density to Monkhorst-Pack grid size.

        inspired by [ase.calculators.calculator.kptdensity2monkhorstpack]

        Parameters
        ----------
        atoms: MSONableAtoms
            Contains unit cell and information about boundary conditions.
        kptdensity: float or list of floats
            Required k-point density.  Default value is 5.0 point per Ang^-1.
        even: bool
            Round up to even numbers.

        Returns
        -------
        Monkhorst-Pack grid size in all directions
        """
        recipcell = atoms.cell.reciprocal()
        return self.d2k_recipcell(recipcell, atoms.pbc, kptdensity, even)

    def k2d(self, atoms: MSONableAtoms, k_grid: Iterable[int]):
        """Generate the kpoint density in each direction from given k_grid.

        Parameters
        ----------
        atoms: MSONableAtoms
            Atoms object of interest.
        k_grid: Iterable[int]
            k_grid that was used.

        Returns
        -------
        Density of kpoints in each direction. result.mean() computes average density
        """
        recipcell = atoms.cell.reciprocal()
        densities = k_grid / (2 * np.pi * np.sqrt((recipcell**2).sum(axis=1)))
        return np.array(densities)

    @staticmethod
    def d2k_recipcell(
        recipcell: Cell,
        pbc: List[bool],
        kptdensity: float | Iterable[float] = 5.0,
        even: bool = True,
    ) -> Iterable[float]:
        """Convert k-point density to Monkhorst-Pack grid size.

        Parameters
        ----------
        recipcell: ASE Cell object
            The reciprocal cell
        pbc: list of Bools
            If element of pbc is True then system is periodic in that direction
        kptdensity: float or list of floats
            Required k-point density.  Default value is 3.5 point per Ang^-1.
        even: bool
            Round up to even numbers.

        Returns
        -------
        Monkhorst-Pack grid size in all directions
        """
        if not isinstance(kptdensity, Iterable):
            kptdensity = 3 * [float(kptdensity)]
        kpts = []
        for i in range(3):
            if pbc[i]:
                k = (
                    2
                    * np.pi
                    * np.sqrt((recipcell[i] ** 2).sum())
                    * float(kptdensity[i])
                )
                if even:
                    kpts.append(2 * int(np.ceil(k / 2)))
                else:
                    kpts.append(int(np.ceil(k)))
            else:
                kpts.append(1)
        return kpts

def recursive_update(d: Dict, u: Dict) -> Dict:
    """
    Update a dictionary recursively and return it.

    Parameters
    ----------
        d: Dict
            Input dictionary to modify
        u: Dict
            Dictionary of updates to apply

    Returns
    -------
    The updated dictionary.

    Example
    -------
        d = {'activate_hybrid': {"hybrid_functional": "HSE06"}}
        u = {'activate_hybrid': {"cutoff_radius": 8}}

        yields {'activate_hybrid': {"hybrid_functional": "HSE06", "cutoff_radius": 8}}}
    """
    for k, v in u.items():
        if isinstance(v, dict):
            d[k] = recursive_update(d.get(k, {}), v)
        elif k == "output" and isinstance(
            v, list
        ):  # for all other keys the list addition is not needed (I guess)
            old_v = d.get(k, [])
            d[k] = old_v + v
        else:
            d[k] = v
    return d


#!/usr/bin/env python

'''
--------------------------------------
Find the longest Open Reading Frame
--------------------------------------
'''
import sys
import logging

from textwrap import wrap
from optparse import OptionParser

from cpmodule import fickett
from cpmodule  import FrameKmer
from cpmodule  import ireader
from cpmodule import find_orfs
from cpmodule.utils import seq_from_bed, index_fasta, bed_or_fasta, coding_prediction,finish_up,signal_handler

__author__ = "Liguo Wang"
__contributor__="Liguo Wang"
__copyright__ = "Copyright 2020, Mayo Clinic"
__credits__ = []
__license__ = "GPLv2"
__version__="3.0.4"
__maintainer__ = "Liguo Wang"
__email__ = "wangliguo78@gmail.com"
__status__ = "Production"

def main():
	usage = "\n%prog  [options]"
	parser = OptionParser(usage,version="%prog " + __version__)
	parser.add_option("-g","--gene",action="store",type="string", dest="gene_file",help="Genomic sequnence(s) of RNA in FASTA (https://en.wikipedia.org/wiki/FASTA_format) or standard 12-column BED (https://genome.ucsc.edu/FAQ/FAQformat.html#format1) format. It is recommended to use *short* and *unique* sequence identifiers (such as Ensembl transcript id) in FASTA and BED file. If this is a BED file, reference genome ('-r/--ref') should be specified. The input FASTA or BED file could be a regular text file or compressed file (*.gz, *.bz2) or accessible URL (http://, https://, ftp://). URL file cannot be a compressed file.")
	parser.add_option("-o","--outfile",action="store",type="string", dest="out_file",help="The prefix of output files.")
	parser.add_option("-d","--logitModel",action="store",dest="logit_model",help="Logistic regression model. The prebuilt models for Human, Mouse, Fly, Zebrafish are availablel. Run 'make_logitModel.py' to build logistic regression model for your own training datset.")
	parser.add_option("-x","--hex",action="store",dest="hexamer_dat",help="The hexamer frequency table. The prebuilt tables for Human, Mouse, Fly, Zebrafish are availablel. Run 'make_hexamer_tab.py' to make this table for your own training dataset.")
	parser.add_option("-r","--ref",action="store",dest="ref_genome",help="Reference genome sequences in FASTA format. Reference genome file will be indexed automatically if the index file ( *.fai) does not exist. Will be ignored if FASTA file was provided to '-g/--gene'.")
	parser.add_option("--antisense",action="store_true",dest="antisense",default=False,help="Logical to determine whether to search for ORFs from the anti-sense strand. *Sense strand* (or coding strand) is DNA strand that carries the translatable code in the 5′ to 3′ direction. default=False (i.e. only search for ORFs from the sense strand)")
	parser.add_option("--start",action="store",type="string", dest="start_codons",default='ATG',help="Start codon (use 'T' instead of 'U') used to define the start of open reading frame (ORF). default=%default")
	parser.add_option("--stop",action="store",type="string", dest="stop_codons",default='TAG,TAA,TGA',help="Stop codon (use 'T' instead of 'U') used to define the end of open reading frame (ORF). Multiple stop codons are separated by ','. default=%default")
	parser.add_option("--min-orf",action="store",type="int", dest="min_orf_len",default=75,help="Minimum ORF length in nucleotides.  default=%default")
	parser.add_option("--top-orf",action="store",type="int", dest="n_top_orf",default=5,help="Number of ORF candidates reported. RNAs may have dozens of putative ORFs, in most cases, the real ORF is ranked (by size) in the top several. It is not necessary to calculate \"Fickett score\", \"Hexamer score\" and \"coding probability\" for every ORF. default=%default")
	parser.add_option("--width",action="store",type="int", dest="line_width",default=100,help="Line width of output ORFs in FASTA format.  default=%default")
	parser.add_option("--log-file",action="store",type="string", dest="log_file",default='CPAT_run_info.log',help="Name of log file. default=\"%default\"")
	parser.add_option("--best-orf",action="store",type="string", dest="mode",default='p',help="Criteria to select the best ORF: \"l\"=length, selection according to the \"ORF length\"; \"p\"=probability, selection according to the \"coding probability\". default=\"%default\"")
	parser.add_option("--verbose",action="store_true",dest="debug",default=False,help="Logical to determine if detailed running information is printed to screen.")
	(options,args)=parser.parse_args()

	for file in ([options.gene_file, options.hexamer_dat, options.logit_model, options.out_file]):
		if not (file):
			parser.print_help()
			sys.exit(0)

	if options.line_width < 1:
		sys.exit(0)

	if options.mode not in ["p", "l"]:
		print ("Please specifiy either \"p\" or \"l\" to --best-orf.", file=sys.stderr)
		sys.exit(0)

	#logging to file
	if options.debug:
		logging.basicConfig(filename='%s' % options.log_file, filemode='w',format = "%(asctime)s [%(levelname)s]  %(message)s",datefmt='%Y-%m-%d %I:%M:%S', level=logging.DEBUG)
	else:
		logging.basicConfig(filename='%s' % options.log_file, filemode='w',format = "%(asctime)s [%(levelname)s]  %(message)s",datefmt='%Y-%m-%d %I:%M:%S', level=logging.INFO)
	#logging to console
	logFormat = logging.Formatter("%(asctime)s [%(levelname)s]  %(message)s",datefmt='%Y-%m-%d %I:%M:%S')
	consoleHandler = logging.StreamHandler()
	consoleHandler.setFormatter(logFormat)
	logging.getLogger().addHandler(consoleHandler)

	logging.info ("Running CPAT version %s..." % ( __version__))
	start_codons = options.start_codons.replace(' ','').split(',')
	stop_codons = options.stop_codons.replace(' ','').split(',')

	SEQOUT = open(options.out_file + '.ORF_seqs.fa', 'w')
	INFOUT = open(options.out_file + '.ORF_info.tsv', 'w')
	NOORF = open(options.out_file + '.no_ORF.txt', 'w')

	logging.info ("Start codons used: [%s]" % ','.join(start_codons))
	logging.info ("Stop codons used: [%s]" % ','.join(stop_codons))

	#build hexamer table from hexamer frequency file
	logging.info ("Reading %s" % options.hexamer_dat)
	coding = {}
	noncoding = {}
	for line in open(options.hexamer_dat):
		line = line.strip()
		fields = line.split()
		if fields[0] == 'hexamer':
			continue
		coding[fields[0]] = float(fields[1])
		noncoding[fields[0]] =  float(fields[2])

	count = 0
	logging.info ("Checking format of \"%s\"" % options.gene_file)
	file_format = bed_or_fasta(options.gene_file)
	if file_format == 'UNKNOWN':
		logging.error("Unknown file format:%s" % options.gene_file)
		sys.exit(0)

	elif file_format == 'FASTA':
		logging.info("Input gene file is in FASTA format")
		if options.ref_genome:
			logging.warning("\"%s\" is a sequence file. The reference genome file \"%s\" will be ignored." % (options.gene_file, options.ref_genome))

		logging.info ("Searching for ORFs ...")
		print ("\t".join(["ID", "mRNA","ORF_strand", "ORF_frame", "ORF_start", "ORF_end","ORF", "Fickett", "Hexamer" ]), file=INFOUT)	## do NOT change these labels, they are R variable names in the model.
		for name,seq in FrameKmer.seq_generator(options.gene_file):
			count += 1
			RNA_len = len(seq)
			#ORF serial number, starting from 1
			orf_sn = 1
			tmp1 = find_orfs.ORFFinder(seq = seq, min_orf = options.min_orf_len)
			ORFs = tmp1.orf_candidates(antisense = options.antisense, n_candidate = options.n_top_orf, start_coden = start_codons, stop_coden = stop_codons)
			if len(ORFs) == 0:
				logging.warning("No ORFs found for %s" % name)
				print (name, file=NOORF)
				continue
			for orf in ORFs:
				# (direction, frame_number+1, orf_start, orf_end, L, sequence)
				orf_seq = orf[-1]
				if orf[0] == '+':
					orf[2] = orf[2] + 1	#change 0-based into 1-based to be consistent with NCBI ORFfinder output (https://www.ncbi.nlm.nih.gov/orffinder/)
				elif orf[0] == '-':
					orf[2] = RNA_len - (orf[2])
					orf[3] = RNA_len - orf[3] + 1

				orf_id = name + '_ORF_' + str(orf_sn) + '\t' + str(RNA_len) + '\t' + '\t'.join([str(i) for i in orf[:-1]])

				fickett_score = fickett.fickett_value(orf_seq)
				hexamer_score = FrameKmer.kmer_ratio(orf_seq,6,3,coding,noncoding)
				print (orf_id + '\t' + str(fickett_score) + '\t' + str(hexamer_score), file=INFOUT)

				print (">" + orf_id, file=SEQOUT)
				print ('\n'.join(wrap(orf_seq, width = options.line_width)), file=SEQOUT)
				orf_sn += 1
			print("%d sequences finished\r" % count, end=' ', file=sys.stderr)
		print ("\n", file=sys.stderr)

	elif file_format == 'BED':
		logging.info("Input gene file is in BED format")
		if not options.ref_genome:
			logging.error("Reference genome file (-r/--ref) must be provided.")
			parser.print_help()
			sys.exit(0)

		logging.info ("Searching for ORFs ...")
		print ("\t".join(["ID", "mRNA","ORF_strand", "ORF_frame", "ORF_start", "ORF_end","ORF", "Fickett", "Hexamer" ]), file=INFOUT)	## do NOT change these labels, they are R variable names in the model.

		index_fasta(options.ref_genome)

		for line in ireader.reader(options.gene_file):
			count +=1
			if line.startswith('track'):continue
			if line.startswith('#'):continue
			if line.startswith('browser'):continue
			name,seq = seq_from_bed(line, options.ref_genome)

			RNA_len = len(seq)
			#ORF serial number, starting from 1
			orf_sn = 1
			tmp1 = find_orfs.ORFFinder(seq = seq, min_orf = options.min_orf_len)
			ORFs = tmp1.orf_candidates(antisense = options.antisense, n_candidate = options.n_top_orf, start_coden = start_codons, stop_coden = stop_codons)
			if len(ORFs) == 0:
				logging.warning("No ORFs found for %s" % name)
				print (line, file=NOORF)
				continue
			for orf in ORFs:
				# (direction, frame_number+1, orf_start, orf_end, L, sequence)
				orf_seq = orf[-1]
				if orf[0] == '+':
					orf[2] = orf[2] + 1	#change 0-based into 1-based to be consistent with NCBI ORFfinder output (https://www.ncbi.nlm.nih.gov/orffinder/)
				elif orf[0] == '-':
					orf[2] = RNA_len - (orf[2])
					orf[3] = RNA_len - orf[3] + 1

				orf_id = name + '_ORF_' + str(orf_sn) + '\t' + str(RNA_len) + '\t' + '\t'.join([str(i) for i in orf[:-1]])

				fickett_score = fickett.fickett_value(orf_seq)
				hexamer_score = FrameKmer.kmer_ratio(orf_seq,6,3,coding,noncoding)
				print (orf_id + '\t' + str(fickett_score) + '\t' + str(hexamer_score), file=INFOUT)

				print (">" + orf_id, file=SEQOUT)
				print ('\n'.join(wrap(orf_seq, width = options.line_width)), file=SEQOUT)
				orf_sn += 1
			print("%d rows finished\r" % count, end=' ', file=sys.stderr)
		print ("\n", file=sys.stderr)

	SEQOUT.close()
	INFOUT.close()

	logging.info ("Calculate coding probability ...")
	coding_prediction(options.logit_model, options.out_file + '.ORF_info.tsv', options.out_file)	#output options.out_file + '.ORF_prob.tsv'

	if options.mode == 'p':
		logging.info ("Select ORF with the highest coding probability ...")
		col_index = 9
	elif options.mode == 'l':
		logging.info ("Select the longest ORF ...")
		col_index = 6

	BEST = open((options.out_file + '.ORF_prob.best.tsv'), 'w')
	best_candidates = {}
	for l in open((options.out_file + '.ORF_prob.tsv'), 'r'):
		l = l.strip()
		if l.startswith('ID'):
			print ("seq_ID\t" + l, file=BEST)
			continue
		f = l.split('\t')
		seq_id = f[0].split('_ORF_')[0]
		prob = float(f[col_index])
		if seq_id not in best_candidates:
			best_candidates[seq_id] = f
		else:
			if prob > float(best_candidates[seq_id][col_index]):
				best_candidates[seq_id] = f

	for k,v in best_candidates.items():
		print (k + '\t' + '\t'.join(v), file=BEST)

	BEST.close()
	logging.info ("Done!")

	finish_up(options.out_file, options.n_top_orf, options.min_orf_len)

if __name__=='__main__':
	main()