from datetime import datetime, timezone
import math
import sys
import traceback

# ________________________________________ Constants ________________________________________

# Parameters:
MAX_DEBT_RATIO                      = 0.8                   # Eg, if 1,000,000 USM are outstanding, users won't be able to redeem FUM unless the ETH pool's (mid) value is >= $1,000,000 / 0.8 = $1,250,000
BUY_SELL_ADJUSTMENTS_HALF_LIFE      = 60                    # Decay rate of our bid/ask related to recent buy/sell activity (eg, rate of buy price, pushed up by buys, dropping back towards oracle buy price): 1.5 -> 1.2247 -> 1.1067
MIN_FUM_BUY_PRICE_HALF_LIFE         = 24 * 60 * 60          # min_fum_buy_price_in_eth() drops by 50% every day

# Price side constants:
MID                                 = 'mid'
BUY                                 = 'buy'
SELL                                = 'sell'

# Gas-saving approximation constants:
APPROXIMATE_TO_SAVE_GAS             = False                 # Switch to True for less accurate, but (hopefully) more gas-efficient calculations
SHIFT                               = 18                    # Number of decimal digits by which fixed-decimal inputs & outputs are shifted.
ONE_SHIFTED                         = 10**SHIFT
ONE_TENTH_SHIFTED                   = ONE_SHIFTED // 10
HALF_TO_THE_ONE_TENTH_SHIFTED       = 933032991536807416    # Basically round(0.5**0.1 * ONE_SHIFTED), except Python doesn't calc that quite precisely so why not hardcode it here (using good old wolframalpha.com)

# ________________________________________ State variables ________________________________________

time                                = datetime(2020, 8, 1, tzinfo=timezone.utc).timestamp()
oracle_eth_buy_price                = 202
oracle_eth_sell_price               = 198
pool_eth                            = 0
usm_holdings                        = {}
fum_holdings                        = {}
mint_burn_adjustment_stored         = 1                     # Price multiplier based on recent mint/burn activity.  Eg, if A just did mint ops driving the ETH sell price down by 0.7x, and B just burned pushing ETH buy price up 1.2x, this factor will be 0.84.  Decays towards 1 over time.
mint_burn_adjustment_timestamp      = 0
fund_defund_adjustment_stored       = 1                     # Same as above, but for funds (increases factor)/defunds (decreases factor).
fund_defund_adjustment_timestamp    = 0
min_fum_buy_price_in_eth_stored     = 0                     # Note that this price is in terms of ETH, not USD/USM.
min_fum_buy_price_timestamp         = 0

# ________________________________________ Main loop ________________________________________

def main():
    input_loop()

def input_loop():
    while True:
        set_min_fum_buy_price_in_eth_if_needed()            # The price calculation here technically may not quite right, because the theoretical FUM price increases (slightly) *during* many ops, as we collect fees...  But #letskeepitsimple
        clear_min_fum_buy_price_if_obsolete()
        print(status_summary())
        print()
        line = input("> ")
        words = line.split()
        try:
            if words[0] == "price":
                # "price 150" -> change the current ETH price (buy and sell) in our simulation to $150, or "price 150/160" -> sell price $150, buy price $160
                prices = map(float, words[1].split('/'))
                set_oracle_eth_price(*prices)
            elif words[0] == "mint":
                # "mint A 10" -> user A adds 10 ETH to the pool, getting back 10 * eth_sell_price newly-minted USM (minus fees)
                user, eth_to_add = words[1], float(words[2])
                usm_minted = mint_usm(user, eth_to_add)
                print("Minted {:,} new USM for {} from {:,} ETH, for {:,} ETH (~${:,}) each.".format(round(usm_minted, 4), user, round(eth_to_add, 6), round(eth_to_add / usm_minted, 8), round(eth_to_add / usm_minted * calc_eth_price(MID), 6)))
            elif words[0] == "burn":
                # "burn A 1000" -> user A burns 1,000 of their USM, getting back (1,000 / eth_buy_price) ETH (minus fees)
                user, usm_to_burn = words[1], float(words[2])
                eth_removed = burn_usm(user, usm_to_burn)
                print("Burned {:,} of {}'s USM for {:,} ETH (~${:,}) each, yielding {:,} ETH.".format(round(usm_to_burn, 4), user, round(eth_removed / usm_to_burn, 8), round(eth_removed / usm_to_burn * calc_eth_price(MID), 6), round(eth_removed, 6)))
            elif words[0] == "fund_eth":
                # "fund_eth B 5" -> user B adds 5 ETH to the pool, getting back a corresponding amount of newly-created FUM (based on the current FUM price, roughly buffer_value() / fum_outstanding())
                user, eth_to_add = words[1], float(words[2])
                fum_created = create_fum_from_eth(user, eth_to_add)
                print("Created {:,} new FUM for {} from {:,} ETH, for {:,} ETH (~${:,}) each.".format(round(fum_created, 4), user, round(eth_to_add, 6), round(eth_to_add / fum_created, 8), round(eth_to_add / fum_created * calc_eth_price(MID), 6)))
            elif words[0] == "fund_usm":
                # "fund_usm B 1000" -> user B returns (burns) 1000 USM to the pool, getting back a corresponding amount of newly-created FUM (based on the current FUM price).  This leaves total pool value unchanged - basically converting USM to FUM, deceasing debt ratio.
                user, usm_to_convert = words[1], float(words[2])
                fum_created = create_fum_from_usm(user, usm_to_convert)
                print("Created {:,} new FUM for {} from {:,} USM, for {:,} USM (~${:,}) each.".format(round(fum_created, 4), user, round(usm_to_convert, 4), round(usm_to_convert / fum_created, 6), round(usm_to_convert / fum_created, 6)))
            elif words[0] == "defund":
                # "defund B 1000" -> user B redeems 1,000 of their FUM, getting back a corresponding amount of ETH (based on the current FUM price)
                user, fum_to_redeem = words[1], float(words[2])
                eth_removed = redeem_fum(user, fum_to_redeem)
                print("Redeemed {:,} of {}'s FUM for {:,} ETH (~${:,}) each, yielding {:,} ETH.".format(round(fum_to_redeem, 4), user, round(eth_removed / fum_to_redeem, 8), round(eth_removed / fum_to_redeem * calc_eth_price(MID), 6), round(eth_removed, 6)))
            elif words[0] == "wait":
                # "wait 300" -> wait 300 seconds (5 minutes)
                wait = float(words[1])
                set_time(time + wait)
            else:
                raise ValueError("Unrecognized command: '{}'".format(words))
        except Exception as err:
            print("Error:", sys.exc_info())
            traceback.print_tb(err.__traceback__)

def status_summary():
    time_string = datetime.utcfromtimestamp(time).strftime('%Y/%m/%d %H:%M:%S')
    min_fum_buy_price_string = "" if min_fum_buy_price_in_eth() == 0 else ", min {:,} ETH (~${:,})".format(round(min_fum_buy_price_in_eth(), 8), round(min_fum_buy_price_in_eth() * calc_eth_price(MID), 6))
    return "{}: {:,} ETH at ${:,} (${:,}/${:,}) = ${:,} pool value, {:,} USM outstanding (${:,}/${:,}, adj {}), buffer = ${:,}, debt ratio = {:.2%}, {:,} FUM outstanding (${:,}/${:,}{}, adj {})\nUSM holdings: {}\nFUM holdings: {}".format(
        time_string, round(pool_eth, 6), round(calc_eth_price(MID), 4), round(oracle_eth_sell_price, 4), round(oracle_eth_buy_price, 4), round(pool_value(), 2), round(usm_outstanding(), 4), round(calc_usm_price(SELL), 6), round(calc_usm_price(BUY), 6), round(mint_burn_adjustment(), 6),
        round(buffer_value(), 2), debt_ratio(), round(fum_outstanding(), 4), round(calc_fum_price(SELL), 6), round(calc_fum_price(BUY), 6), min_fum_buy_price_string, round(fund_defund_adjustment(), 6), usm_holdings, fum_holdings)

# ________________________________________ State-modifying operations ________________________________________

def set_time(new_time):
    global time
    time = new_time

def set_oracle_eth_price(new_price, new_buy_price=None):
    global oracle_eth_sell_price, oracle_eth_buy_price
    oracle_eth_sell_price = new_price
    oracle_eth_buy_price = new_buy_price if new_buy_price is not None else new_price

    if min_fum_buy_price_needs_setting():
        # Set the min FUM buy price (in ETH), to the FUM price in ETH as of the ETH price point where the debt ratio exceeded MAX_DEBT_RATIO.  Eg, suppose pool_eth = 400 and fum_outstanding() = 1,000.  Then, at the moment we exceed MAX_DEBT_RATIO = 0.8, the buffer must contain
        # 400 * (1 - 0.8) = 80 ETH, and therefore the FUM price in ETH is 80 / 1,000 = 0.08.  Eg, suppose USM outstanding = 40,000: then we cross 0.8 at ETH = $125, when total pool value = $50,000, buffer = $10,000, and FUM price = $10,000 / 1,000 = $10 = ($10 / $125) = 0.08 ETH.
        fum_price_in_eth_at_which_we_crossed_max_debt_ratio = (pool_eth * (1 - MAX_DEBT_RATIO)) / fum_outstanding()
        set_min_fum_buy_price_in_eth_if_needed(fum_price_in_eth_at_which_we_crossed_max_debt_ratio)

def mint_usm(user, eth_to_add):
    # Note that minting never pulls us *below* MAX_DEBT_RATIO, since it always moves debt ratio closer to 1.  If debt ratio starts > 1, it stays > 1; if it starts < 1 and > MAX_DEBT_RATIO, it stays > MAX_DEBT_RATIO.
    global pool_eth
    initial_eth_price = calc_eth_price(SELL)
    if pool_eth == 0:
        # This is our very first ETH in the pool, so need to special-case it (otherwise the division below blows up):
        usm_minted = eth_to_add * initial_eth_price
    else:
        # Mint at a sliding-down ETH price (ie, buy USM at a sliding-up USM price).  **BASIC RULE:** anytime pool_eth changes by factor k, eth_price changes by factor 1/k**2.  (Earlier versions of this logic scaled price by 1/k, not 1/k**2; but that results in calls to log()/exp().)
        pool_eth_growth_factor = (pool_eth + eth_to_add) / pool_eth
        usm_minted = pool_eth * initial_eth_price * (1 - 1 / pool_eth_growth_factor)                                # Math: this is an integral - sum of all USM minted at a sliding-down ETH price
        set_mint_burn_adjustment(mint_burn_adjustment() / pool_eth_growth_factor**2)
    pool_eth += eth_to_add
    usm_holdings[user] = usm_holdings.get(user, 0) + usm_minted
    return usm_minted

def burn_usm(user, usm_to_burn, check_debt_ratio=True):
    # Note that burning never pushes us over MAX_DEBT_RATIO (which is < 1), since it always moves debt ratio further from 1.  It can pull us *below* MAX_DEBT_RATIO, but if so the top-level call to clear_min_fum_buy_price_if_obsolete() will take care of it once this op is done.
    global pool_eth
    assert usm_to_burn <= usm_holdings.get(user, 0), "{} doesn't own that many USM".format(user)
    initial_eth_price = calc_eth_price(BUY)
    # Burn at a sliding-up price:
    eth_removed = usm_to_burn / (initial_eth_price + usm_to_burn / pool_eth)                                        # Math: this is an integral - sum of all USM burned at a sliding price.  Follows the same mathematical invariant as above: if pool_eth *= k, eth_price *= 1/k**2.
    assert eth_removed <= pool_eth, "Not enough ETH in the pool"
    if check_debt_ratio:
        assert debt_ratio(eth=pool_eth - eth_removed, usm=usm_outstanding() - usm_to_burn) <= 1, "Burning {:,} USM would leave the debt ratio above 100%".format(usm_to_burn)   # Note: the risk is not this burn op pushing us over 100%, but that a previous price drop might have done so!
    pool_eth_shrink_factor = (pool_eth - eth_removed) / pool_eth
    set_mint_burn_adjustment(mint_burn_adjustment() / pool_eth_shrink_factor**2)
    usm_holdings[user] -= usm_to_burn
    pool_eth -= eth_removed
    return eth_removed

def create_fum_from_eth(user, eth_to_add):
    # Fund operations never push us over MAX_DEBT_RATIO: they reduce debt ratio.  However, they *can* bring us back *under* MAX_DEBT_RATIO, which means the naive logic here may overcharge an op that pulls debt ratio below MAX_DEBT_RATIO midway through...  But oh well, #letskeepitsimple
    global pool_eth
    if fum_outstanding() == 0:
        # This is our very first FUM created, so need to special-case it (otherwise the division below blows up):
        fum_created = eth_to_add * calc_eth_price(MID)                                                              # No need for any adjustment: the FUM price only matters relative to an existing FUM price, so we can just price the first FUM units at $1
    else:
        # Create at a sliding-up price:
        initial_fum_price = calc_fum_price(BUY)
        initial_eth_price_in_fum = calc_eth_price(MID) / initial_fum_price                                          # Don't apply adjustment to the ETH price - the adjustment should only be applied as the last step in a transaction, not when used indirectly for pricing as here.
        pool_eth_growth_factor = (pool_eth + eth_to_add) / pool_eth
        fum_created = pool_eth * initial_eth_price_in_fum * (1 - 1 / pool_eth_growth_factor)                        # Math: see closely analogous comment in mint_usm() above
        set_fund_defund_adjustment(fund_defund_adjustment() * pool_eth_growth_factor**2)
    pool_eth += eth_to_add
    fum_holdings[user] = fum_holdings.get(user, 0) + fum_created
    return fum_created

def create_fum_from_usm(user, usm_to_convert):
    # To avoid duplication, just implement this as a call to burn_usm() (bypassing the debt ratio check), followed by a call to create_fum_from_eth().  This is a bit hazardous because we could die on an error with half the operation complete, but good enough for govt work:
    eth_converted = burn_usm(user, usm_to_convert, check_debt_ratio=False)
    clear_min_fum_buy_price_if_obsolete()                                                                           # Important so that, if the preceding burn brought us below MAX_DEBT_RATIO, we clear the min FUM buy price before starting the fund operation
    return create_fum_from_eth(user, eth_converted)

def redeem_fum(user, fum_to_redeem):
    global pool_eth
    assert fum_to_redeem <= fum_holdings.get(user, 0), "{} doesn't own that many FUM".format(user)
    initial_fum_price = calc_fum_price(SELL)
    initial_eth_price_in_fum = calc_eth_price(MID) / initial_fum_price                                              # Don't apply adjustment to the ETH price - see similar comment above.
    eth_removed = fum_to_redeem / (initial_eth_price_in_fum + fum_to_redeem / pool_eth)                             # Math: see closely analogous comment in burn_usm() above
    assert debt_ratio(eth=pool_eth - eth_removed) <= MAX_DEBT_RATIO, "Redeeming {:,} FUM would leave the debt ratio above {:.0%}".format(fum_to_redeem, MAX_DEBT_RATIO)
    # Since we've disallowed redeem operations that would push us over MAX_DEBT_RATIO, we don't need to handle that case.  And a redeem can never pull us under MAX_DEBT_RATIO either, because it increases debt ratio.
    pool_eth_shrink_factor = (pool_eth - eth_removed) / pool_eth
    set_fund_defund_adjustment(fund_defund_adjustment() * pool_eth_shrink_factor**2)
    fum_holdings[user] -= fum_to_redeem
    pool_eth -= eth_removed
    return eth_removed

def set_min_fum_buy_price_in_eth_if_needed(price_in_eth=None):
    global min_fum_buy_price_in_eth_stored, min_fum_buy_price_timestamp
    if min_fum_buy_price_needs_setting():
        eth_price = calc_eth_price(MID)
        if price_in_eth is None:
            # If no price is passed in, set it to the current theoretical FUM price (without adjustments), in ETH:
            price_in_usd = calc_fum_price(BUY, adjusted=False, mfbp=False)
            price_in_eth = price_in_usd / eth_price
        else:
            price_in_usd = price_in_eth * eth_price
        print("* Setting min FUM buy price to {:,} ETH (~${:,}), since debt ratio {:.2%} is above {:.0%}.".format(round(price_in_eth, 8), round(price_in_usd, 6), debt_ratio(), MAX_DEBT_RATIO))
        min_fum_buy_price_in_eth_stored = price_in_eth
        min_fum_buy_price_timestamp = time

def clear_min_fum_buy_price_if_obsolete():
    global min_fum_buy_price_in_eth_stored, min_fum_buy_price_timestamp
    if min_fum_buy_price_in_eth() != 0 and debt_ratio() <= MAX_DEBT_RATIO:
        print("* Resetting min FUM buy price to $0, since debt ratio {:.2%} is back below {:.0%}.".format(debt_ratio(), MAX_DEBT_RATIO))
        min_fum_buy_price_in_eth_stored = 0
        min_fum_buy_price_timestamp = None

def set_mint_burn_adjustment(adjustment_factor):
    global mint_burn_adjustment_stored, mint_burn_adjustment_timestamp
    mint_burn_adjustment_stored = adjustment_factor
    mint_burn_adjustment_timestamp = time

def set_fund_defund_adjustment(adjustment_factor):
    global fund_defund_adjustment_stored, fund_defund_adjustment_timestamp
    fund_defund_adjustment_stored = adjustment_factor
    fund_defund_adjustment_timestamp = time

# ________________________________________ Informational USM/FUM utility functions ________________________________________

def usm_outstanding():
    return sum(usm_holdings.values())

def fum_outstanding():
    return sum(fum_holdings.values())

def pool_value(eth=None, eth_price=None):
    if eth is None:
        eth = pool_eth
    if eth_price is None:
        eth_price = calc_eth_price(MID)
    return eth * eth_price

def buffer_value(eth_price=None):
    if eth_price is None:
        eth_price = calc_eth_price(MID)
    return pool_value(eth_price=eth_price) - usm_outstanding()

def debt_ratio(eth=None, usm=None):
    if eth is None:
        eth = pool_eth
    if usm is None:
        usm = usm_outstanding()
    if pool_value(eth=eth) == 0:
        return 0
    else:
        return usm / pool_value(eth=eth)

def calc_eth_price(side, adjusted=True):
    assert side in (MID, BUY, SELL)
    if side == BUY:
        price = oracle_eth_buy_price
        if adjusted:
            price *= max(1, mint_burn_adjustment()) * max(1, fund_defund_adjustment())
    elif side == SELL:
        price = oracle_eth_sell_price
        if adjusted:
            price *= min(1, mint_burn_adjustment()) * min(1, fund_defund_adjustment())
    else:
        price = (oracle_eth_sell_price + oracle_eth_buy_price) / 2
    return price

def calc_fum_price(side, adjusted=True, mfbp=True):
    assert side in (MID, BUY, SELL)
    if fum_outstanding() == 0:
        return 1 if side == BUY else math.nan                                       # If we're pricing our first FUM purchase, just price them at $1, skipping adjustment
    else:
        eth_price = calc_eth_price(side, adjusted=False)                            # adjusted=False because we apply that directly to the FUM price below, not to the ETH price used to calculate the buffer value - that would exaggerate the fee too much
        price = buffer_value(eth_price=eth_price) / fum_outstanding()

    if side == BUY:
        if adjusted:
            price *= max(1, mint_burn_adjustment()) * max(1, fund_defund_adjustment())
        if mfbp:
            price = max(price, min_fum_buy_price_in_eth() * calc_eth_price(MID))
    elif side == SELL:
        if adjusted:
            price *= min(1, mint_burn_adjustment()) * min(1, fund_defund_adjustment())
    return price

def calc_usm_price(side, adjusted=True):
    assert side in (MID, BUY, SELL)
    eth_side = {BUY: SELL, SELL: BUY, MID: MID}[side]
    return calc_eth_price(MID) / calc_eth_price(eth_side, adjusted=adjusted)

def min_fum_buy_price_needs_setting():
    return min_fum_buy_price_in_eth() == 0 and debt_ratio() > MAX_DEBT_RATIO and fum_outstanding() > 0

def min_fum_buy_price_in_eth():
    if min_fum_buy_price_timestamp is None:
        return 0
    elif APPROXIMATE_TO_SAVE_GAS:
        return round(min_fum_buy_price_in_eth_stored * half_exp_approx(((time - min_fum_buy_price_timestamp) / MIN_FUM_BUY_PRICE_HALF_LIFE) * ONE_SHIFTED)) / ONE_SHIFTED
    else:
        return min_fum_buy_price_in_eth_stored * (0.5 ** ((time - min_fum_buy_price_timestamp) / MIN_FUM_BUY_PRICE_HALF_LIFE))

def mint_burn_adjustment():
    if APPROXIMATE_TO_SAVE_GAS:
        power_approx = half_exp_approx(((time - mint_burn_adjustment_timestamp) / BUY_SELL_ADJUSTMENTS_HALF_LIFE) * ONE_SHIFTED, max_power=10) / ONE_SHIFTED
        # Here we use the idea that for  0 < b <= 1 and 0 <= p <= 1, we can crudely approximate b**p by 1 - (1-b)p.  Eg: 0.6**0.5 pulls 0.6 "about halfway" to 1 (0.8); 0.6**0.25 pulls 0.6 "about 3/4 of the way" to 1 (0.9).  So b**p =~ b + (1-p)(1-b) = b + 1 - b - p + bp = 1 - (1-b)p:
        return 1 - (1 - mint_burn_adjustment_stored) * power_approx
    else:
        return mint_burn_adjustment_stored ** (0.5 ** ((time - mint_burn_adjustment_timestamp) / BUY_SELL_ADJUSTMENTS_HALF_LIFE))

def fund_defund_adjustment():
    if APPROXIMATE_TO_SAVE_GAS:
        power_approx = half_exp_approx(((time - fund_defund_adjustment_timestamp) / BUY_SELL_ADJUSTMENTS_HALF_LIFE) * ONE_SHIFTED, max_power=10) / ONE_SHIFTED
        # See parallel comment above:
        return 1 - (1 - fund_defund_adjustment_stored) * power_approx
    else:
        return fund_defund_adjustment_stored ** (0.5 ** ((time - fund_defund_adjustment_timestamp) / BUY_SELL_ADJUSTMENTS_HALF_LIFE))

# ________________________________________ General-purpose utility functions ________________________________________

def half_exp_approx(power_shifted, max_power=math.inf):
    """Returns a loose but "gas-efficient" approximation of 0.5**power, where:
    1. Both input and output are in fixed-point format, shifted by SHIFT decimal digits.  Eg, input power_shifted = 1400000000000000000 represents power = 1400000000000000000 / 10**18 = 1.4, so returns approx (0.5**1.4) * 10**18 = 378929141627599521
    2. power is rounded to the nearest 10th: power = 0.462 is treated as power = 0.5
    3. Large values of power (> max_power) just return 0, since 0.5**large_power =~ 0"""
    assert power_shifted >= 0
    power_in_tenths = (power_shifted + (ONE_TENTH_SHIFTED // 2)) // ONE_TENTH_SHIFTED   # After this, power_in_tenths must be a non-negative integer
    if power_in_tenths > 10 * max_power:
        return 0
    else:
        return half_to_the_one_tenth_exp_approx(power_in_tenths)

def half_to_the_one_tenth_exp_approx(power):
    """power is an unshifted, non-negative integer; return value is shifted, ie, multiplied by 10**SHIFT."""
    if power == 0:
        return ONE_SHIFTED
    else:
        sqrt = half_to_the_one_tenth_exp_approx(power // 2)
        result = (sqrt * sqrt) // ONE_SHIFTED
        if power % 2 == 1:
            result = (result * HALF_TO_THE_ONE_TENTH_SHIFTED) // ONE_SHIFTED
        return result

main()


import os
import sys
import time

import schedule

def job():
    run_time_cmd = "ps -eo pid,comm,cmd,start,etime | grep 'python main.py' | grep -v grep | sort -k 5 | head -n 1"
    run_time = os.popen(run_time_cmd).read().strip().split(" ")[-1]
    sys.stdout.write("\r" + f"run time : {run_time}")
    sys.stdout.flush()

    run_hour = run_time.split(":")[0]
    if run_time and int(run_hour) >= 60:
        # kill
        print("kill job")
        kill_cmd = "kill -9 `ps -ef|grep 'python main.py'|grep -v 'grep'|awk '{print $2}'`"
        os.system(kill_cmd)
        watchdog = False

watchdog = True

schedule.every(10).minutes.do(job) # 每隔 10 分运行一次 job 函数

while True and watchdog:
    schedule.run_pending()   # 运行所有可以运行的任务
    time.sleep(360)


import itertools
import logging
import math
import random
import cv2
import numpy as np
import torch
from typing import Any, Dict, List, Optional, Tuple
from albumentations import Rotate, DualTransform, ImageOnlyTransform, Resize, RandomScale
from albumentations.augmentations.crops import functional as F_crops
from albumentations.augmentations.geometric import functional as F_geometric
from albumentations.augmentations import functional as F
from torchvision.transforms import functional as F_torchvision
from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD

logger = logging.getLogger(__name__)

class RotateZ(Rotate):
    """Rotate around X axis on randomly selected angle."""
    def __call__(self, *args, force_apply: bool = False, **kwargs) -> Dict[str, Any]:
        # (H, W, D) -> (D, H, W)
        kwargs['image'] = np.transpose(kwargs['image'], (2, 0, 1))

        # Rotate as usual image
        kwargs = super().__call__(*args, force_apply=force_apply, **kwargs)

        # (D, H, W) -> (H, W, D)
        kwargs['image'] = np.transpose(kwargs['image'], (1, 2, 0))

        return kwargs

    def apply_to_mask(self, img, angle=0, x_min=None, x_max=None, y_min=None, y_max=None, **params):
        # Project mask to XZ plane.

        # cv2's y corresponds to H dimention which is shrinked
        # by cos(angle): mask is projected to XZ plane.
        H, W = img.shape[:2]
        img_out = cv2.resize(
            img,
            (0, 0),
            fx=1,
            fy=np.cos(np.deg2rad(angle)),
            interpolation=cv2.INTER_NEAREST,
        )

        # Pad mask to the same size as image by H dimention.
        top_pad = (H - img_out.shape[0]) // 2
        bottom_pad = H - img_out.shape[0] - top_pad
        img_out = cv2.copyMakeBorder(
            img_out,
            top=top_pad,
            bottom=bottom_pad,
            left=0,
            right=0,
            borderType=self.border_mode,
        )

        # Crop mask: if border is cropped, then
        # mask should not only be projected to XZ plane,
        # but also cropped to the same size as image by H dimention.

        # Here only H changes and H corresponds
        # to y here and x in image case because of transpose
        # (see __call__ method).
        if self.crop_border:
            img_out = F_crops.crop(img_out, x_min=0, y_min=x_min, x_max=W, y_max=x_max)

        return img_out

class RandomCropVolumeInside2dMask:
    """Crop a random part of the input.
    """
    def __init__(
        self,
        base_size: Optional[int] = None,
        base_depth: Optional[int] = None,
        scale: Tuple[float, float] = (1.0, 1.0),
        scale_z: Tuple[float, float] = (1.0, 1.0),
        ratio: Tuple[float, float] = (1.0, 1.0),
        value: int = 0,
        mask_value: int = 0,
        crop_mask_index=0,
        always_apply=True,
        p=1.0,
    ):
        if not (base_size is not None or base_depth is not None):
            logger.warning(
                f"Eigher base_size or base_depth should be not None. "
                f"Got base_size={base_size}, base_depth={base_depth}. "
                f"No transform will be performed."
            )

        assert scale[0] > 0 and scale[1] > 0, f"scale should be positive. Got {scale}"
        assert scale_z[0] > 0 and scale_z[1] > 0, f"scale_z should be positive. Got {scale_z}"
        assert ratio[0] > 0 and ratio[1] > 0, f"ratio should be positive. Got {ratio}"
        assert scale[0] <= scale[1], f"scale[0] should be less or equal than scale[1]. Got {scale}"
        assert scale_z[0] <= scale_z[1], f"scale_z[0] should be less or equal than scale_z[1]. Got {scale_z}"
        assert ratio[0] <= ratio[1], f"ratio[0] should be less or equal than ratio[1]. Got {ratio}"

        self.base_size = base_size
        self.base_depth = base_depth
        self.scale = scale
        self.scale_z = scale_z
        self.ratio = ratio
        self.value = value
        self.mask_value = mask_value
        self.crop_mask_index = crop_mask_index

        self.always_apply = always_apply
        self.p = p

    def __call__(self, *args, force_apply: bool = False, **kwargs) -> Dict[str, np.ndarray]:
        # Toss a coin
        if not force_apply and not self.always_apply and random.random() > self.p:
            return kwargs

        # Get random scale and ratio
        scale = random.uniform(self.scale[0], self.scale[1])
        ratio = random.uniform(self.ratio[0], self.ratio[1])

        h_start, h_end = 0, kwargs['image'].shape[0]
        w_start, w_end = 0, kwargs['image'].shape[1]
        if self.base_size is not None:
            # Get height and width
            height = int(self.base_size * scale)
            width = int(self.base_size * scale * ratio)

            # Get crop mask
            crop_mask = kwargs['masks'][self.crop_mask_index]  # (H, W)

            # Crop the mask to ensure that the crop is inside the mask
            h_shift, w_shift = height // 2 + 1, width // 2 + 1
            crop_mask = crop_mask[
                h_shift:-h_shift,
                w_shift:-w_shift,
            ]

            # Get indices of non-zero elements
            nonzero_indices = np.nonzero(crop_mask)

            # Get random index
            random_index = random.randint(0, nonzero_indices[0].shape[0] - 1)
            h_center, w_center = nonzero_indices[0][random_index], nonzero_indices[1][random_index]

            # Shift indices back to compensate crop above
            h_center += h_shift
            w_center += w_shift

            # Get crop indices
            h_start = h_center - height // 2
            h_end = h_start + height
            w_start = w_center - width // 2
            w_end = w_start + width

            # Ensure that crop is inside the image
            assert h_start >= 0 and h_end <= kwargs['image'].shape[0], \
                f"h_start={h_start} and h_end={h_end} should be in [0, {kwargs['image'].shape[0]}]"
            assert w_start >= 0 and w_end <= kwargs['image'].shape[1], \
                f"w_start={w_start} and w_end={w_end} should be in [0, {kwargs['image'].shape[1]}]"

        scale_z = random.uniform(self.scale_z[0], self.scale_z[1])
        z_start, z_end = 0, kwargs['image'].shape[2]
        if self.base_depth is not None:
            # Get depth
            depth = int(self.base_depth * scale_z)

            # Get random z_start
            z_start_max = max(kwargs['image'].shape[2] - depth, 1)
            z_start = np.random.randint(0, z_start_max)

            # Get z_end
            z_end = min(z_start + depth, kwargs['image'].shape[2])

        # Crop data
        kwargs['image'] = kwargs['image'][h_start:h_end, w_start:w_end, z_start:z_end].copy()
        for i in range(len(kwargs['masks'])):
            kwargs['masks'][i] = kwargs['masks'][i][h_start:h_end, w_start:w_end].copy()

        return kwargs

class RandomCropVolume(DualTransform):
    """Crop a random part of the input.
    If image, crop is 3D, if mask, crop is 2D.

    Args:
        height (int): height of the crop.
        width (int): width of the crop.
        depth (int): depth of the crop.
        p (float): probability of applying the transform. Default: 1.

    Targets:
        image, mask

    Image types:
        uint8, float32
    """

    def __init__(self, height, width, depth, always_apply=False, p=1.0):
        super().__init__(always_apply, p)
        self.height = height
        self.width = width
        self.depth = depth

    def apply(self, img, h_start=0, w_start=0, d_start=0, is_mask=False, **params):
        img = F_crops.random_crop(img, self.height, self.width, h_start, w_start)
        if not is_mask:
            z = np.random.randint(d_start, max(img.shape[2] - self.depth, d_start + 1))
            img = img[:, :, z:min(z + self.depth, img.shape[2])]
        return img

    def apply_to_mask(self, img: np.ndarray, **params) -> np.ndarray:
        return self.apply(img, is_mask=True, **{k: cv2.INTER_NEAREST if k == "interpolation" else v for k, v in params.items()})

    def apply_to_bbox(self, bbox, **params):
        return F_crops.bbox_random_crop(bbox, self.height, self.width, **params)

    def apply_to_keypoint(self, keypoint, **params):
        return F_crops.keypoint_random_crop(keypoint, self.height, self.width, **params)

    def get_params(self):
        return {"h_start": random.random(), "w_start": random.random(), "d_start": random.random()}

    def get_transform_init_args_names(self):
        return ("height", "width", "depth")

class CenterCropVolume(DualTransform):
    """Crop the central part of the input.

    Args:
        height (int): height of the crop.
        width (int): width of the crop.
        depth (int): depth of the crop.
        strict (bool): if True, will fail if image size is less than crop size by any dimention,
            otherwise will crop as much as possible. Default: True.
        p (float): probability of applying the transform. Default: 1.

    Targets:
        image, mask

    Image types:
        uint8, float32

    Note:
        It is recommended to use uint8 images as input.
        Otherwise the operation will require internal conversion
        float32 -> uint8 -> float32 that causes worse performance.
    """

    def __init__(self, height, width, depth, strict=True, always_apply=False, p=1.0):
        super().__init__(always_apply, p)
        if not ((height is not None and width is not None) or depth is not None):
            logger.warning(
                f"Eigher height and width or depth should be not None. "
                f"Got height={height}, width={width}, depth={depth}. "
                f"No transform will be performed."
            )
        self.height = height
        self.width = width
        self.depth = depth
        self.strict = strict

    def apply(self, img, is_mask=False, **params):
        if self.height is not None and self.width is not None:
            height, width = img.shape[0], img.shape[1]
            if not self.strict:
                if height < self.height:
                    logger.warning(
                        f"Height of the image is {height} but height of the crop is {self.height}. "
                        f"All the height of the image will be used."
                    )
                if width < self.width:
                    logger.warning(
                        f"Width of the image is {width} but width of the crop is {self.width}. "
                        f"All the width of the image will be used."
                    )
                height, width = \
                    min(img.shape[0], self.height), \
                    min(img.shape[1], self.width)
            img = F_crops.center_crop(img, height, width)
        if not is_mask and self.depth is not None:
            if img.shape[2] < self.depth:
                logger.warning(
                    f"Depth of the image is {img.shape[2]} but depth of the crop is {self.depth}. "
                    f"All the depth of the image will be used."
                )

            if self.strict:
                assert img.shape[2] >= self.depth, \
                    f"Depth of the image is {img.shape[2]} but depth of the crop is {self.depth}. " \

            z_start = max((img.shape[2] - self.depth) // 2, 0)
            z_end = min(z_start + self.depth, img.shape[2])
            img = img[:, :, z_start:z_end]
        return img

    def apply_to_mask(self, img: np.ndarray, **params) -> np.ndarray:
        return self.apply(img, is_mask=True, **{k: cv2.INTER_NEAREST if k == "interpolation" else v for k, v in params.items()})

    def apply_to_bbox(self, bbox, **params):
        return F_crops.bbox_center_crop(bbox, self.height, self.width, **params)

    def apply_to_keypoint(self, keypoint, **params):
        return F_crops.keypoint_center_crop(keypoint, self.height, self.width, **params)

    def get_transform_init_args_names(self):
        return ("height", "width", "depth")

class RandomZSlices(ImageOnlyTransform):
    """Randomly select N (permuted) slices from the input.

    Args:
        n_slices (int): number of slices to select.
        p (float): probability of applying the transform. Default: 1.

    Targets:
        image, mask

    Image types:
        uint8, float32
    """

    def __init__(self, n_slices, permuted=True, always_apply=False, p=1.0):
        super().__init__(always_apply, p)
        self.n_slices = n_slices
        self.permuted = permuted

    def apply(self, img, **params):
        indices = np.random.permutation(img.shape[2])[:self.n_slices]
        if not self.permuted:
            indices = np.sort(indices)
        return img[:, :, indices]

    def get_transform_init_args_names(self):
        return ("n_slices", )

class UniformZSlices(ImageOnlyTransform):
    """Select N (uniformly) slices from the input.

    Args:
        n_slices (int): number of slices to select.
        p (float): probability of applying the transform. Default: 1.

    Targets:
        image, mask

    Image types:
        uint8, float32
    """

    def __init__(self, n_slices, always_apply=False, p=1.0):
        super().__init__(always_apply, p)
        self.n_slices = n_slices

    def apply(self, img, **params):
        indices = np.linspace(0, img.shape[2] - 1, self.n_slices).astype(int)
        return img[:, :, indices]

    def get_transform_init_args_names(self):
        return ("n_slices", )

class ResizeVolume(Resize):
    def __init__(self, height, width, depth=None, interpolation=cv2.INTER_LINEAR, always_apply=False, p=1):
        super().__init__(height, width, interpolation, always_apply, p)
        self.depth = depth

    def apply(self, img, interpolation=cv2.INTER_LINEAR, is_mask=False, **params):
        # Interpolate 2D
        img = super().apply(img, interpolation=interpolation, **params)
        if is_mask:
            return img

        if self.depth is not None:
            # Interpolate depth
            # (H, W, D) -> (D, H, W)
            img = np.transpose(img, (2, 0, 1))

            # Resize as usual image
            img = F_geometric.resize(img, height=self.depth, width=img.shape[1], interpolation=interpolation)

            # (D, H, W) -> (H, W, D)
            img = np.transpose(img, (1, 2, 0))

        return img

    def apply_to_mask(self, img: np.ndarray, **params) -> np.ndarray:
        return self.apply(img, is_mask=True, **{k: cv2.INTER_NEAREST if k == "interpolation" else v for k, v in params.items()})

class ToCHWD(ImageOnlyTransform):
    def apply(self, img: torch.Tensor, **params) -> torch.Tensor:
        # (D, H, W) (grayscale) ->
        # (1, H, W, D)
        return img.unsqueeze(0).permute(0, 2, 3, 1)

class ToWritable(DualTransform):
    def apply(self, img, **params):
        img = img.copy()
        img.setflags(write=True)
        return img

class RandomScaleResize(RandomScale):
    """Same as RandomScale but resize to original size after scaling."""
    def apply(self, img, scale=0, interpolation=cv2.INTER_LINEAR, **params):
        original_h, original_w = img.shape[:2]
        img = super().apply(img, scale, interpolation, **params)
        return F_geometric.resize(img, height=original_h, width=original_w, interpolation=interpolation)

class SubtractDivide:
    """Subtracttion and division is applied by the formula:
    `img = (img - subtract) / divide`

    Targets:
        image

    Image types:
        uint8, float32
    """

    def __init__(
        self,
        subtract: float = 0.0,
        divide: float = 1.0,
        always_apply=True,
        p=1.0,
    ):
        self.subtract = subtract
        self.divide = divide
        self.always_apply = always_apply
        self.p = p

    def __call__(self, *args, force_apply: bool = False, **kwargs) -> Dict[str, np.ndarray]:
        # Toss a coin
        if not force_apply and not self.always_apply and random.random() > self.p:
            return kwargs

        subtract, divide = self.subtract, self.divide

        if (
            "subtract" in kwargs and kwargs["subtract"] is not None and
            "divide" in kwargs and kwargs["divide"] is not None
        ):
            subtract = kwargs["subtract"]
            divide = kwargs["divide"]
        else:
            logger.warning(
                f"Either subtract or divide values are not provided. "
                f"Subtract is set to {subtract} and divide is set to {divide}."
            )

        kwargs['image'] = F.normalize(kwargs['image'], subtract, divide, 1)

        return kwargs

class CopyPastePositive:
    """Copy masked area from one image to another.

    Targets:
        image, mask

    Image types:
        uint8, float32
    """

    def __init__(
        self,
        mask_index: int = 2,
        always_apply=True,
        p=1.0,
    ):
        self.mask_index = mask_index
        self.always_apply = always_apply
        self.p = p

    def __call__(self, *args, force_apply: bool = False, **kwargs) -> Dict[str, np.ndarray]:
        # Toss a coin
        if not force_apply and not self.always_apply and random.random() > self.p:
            return kwargs

        mask = (kwargs['masks'][self.mask_index] > 0) & (kwargs['masks1'][self.mask_index] <= 0)

        kwargs['image'][mask] = kwargs['image1'][mask]
        for i in range(len(kwargs['masks'])):
            kwargs['masks'][i][mask] = kwargs['masks1'][i][mask]

        return kwargs

# https://github.com/albumentations-team/albumentations/pull/1409/files
class MixUp:
    def __init__(
        self,
        alpha = 32.,
        beta = 32.,
        always_apply: bool = False,
        p: float = 0.5,
    ):
        self.alpha = alpha
        self.beta = beta
        self.always_apply = always_apply
        self.p = p

    def __call__(self, *args, force_apply: bool = False, **kwargs) -> Dict[str, np.ndarray]:
        # Toss a coin
        if not force_apply and not self.always_apply and random.random() > self.p:
            return kwargs

        h1, w1, _ = kwargs['image'].shape
        h2, w2, _ = kwargs['image1'].shape
        if h1 != h2 or w1 != w2:
            raise ValueError("MixUp transformation expects both images to have identical shape.")

        r = np.random.beta(self.alpha, self.beta)

        kwargs['image'] = (kwargs['image'] * r + kwargs['image1'] * (1 - r))
        for i in range(len(kwargs['masks'])):
            kwargs['masks'][i] = (kwargs['masks'][i] * r + kwargs['masks1'][i] * (1 - r))

        return kwargs

class CutMix:
    def __init__(
        self,
        width: int = 64,
        height: int = 64,
        always_apply: bool = False,
        p: float = 0.5,
    ):
        self.width = width
        self.height = height
        self.always_apply = always_apply
        self.p = p

    def __call__(self, *args, force_apply: bool = False, **kwargs) -> Dict[str, np.ndarray]:
        # Toss a coin
        if not force_apply and not self.always_apply and random.random() > self.p:
            return kwargs

        h, w, _ = kwargs['image'].shape
        h1, w1, _ = kwargs['image1'].shape
        if (
            h < self.height or
            w < self.width or
            h1 < self.height or
            w1 < self.width
        ):
            raise ValueError("CutMix transformation expects both images to be at least {}x{} pixels.".format(self.max_height, self.max_width))

        # Get random bbox
        h_start = random.randint(0, h - self.height)
        w_start = random.randint(0, w - self.width)
        h_end = h_start + self.height
        w_end = w_start + self.width

        # Copy image and masks region
        kwargs['image'][h_start:h_end, w_start:w_end] = kwargs['image1'][h_start:h_end, w_start:w_end]
        for i in range(len(kwargs['masks'])):
            kwargs['masks'][i][h_start:h_end, w_start:w_end] = kwargs['masks1'][i][h_start:h_end, w_start:w_end]

        return kwargs

class ToFloatMasks:
    def __call__(self, *args, force_apply: bool = False, **kwargs) -> Dict[str, np.ndarray]:
        for i in range(len(kwargs['masks'])):
            kwargs['masks'][i] = kwargs['masks'][i].astype(np.float32)
        return kwargs

class EnlargePositive:
    """Scale up masked area."""
    def __init__(
        self,
        mask_index: int = 2,
        scale_limit: float = 0.1,
        interpolation=cv2.INTER_LINEAR,
        always_apply=False,
        p=1.0,
    ):
        self.mask_index = mask_index
        self.scale_limit = scale_limit
        self.interpolation = interpolation
        self.always_apply = always_apply
        self.p = p

    def __call__(self, *args, force_apply: bool = False, **kwargs) -> Dict[str, np.ndarray]:
        # Toss a coin
        if not force_apply and not self.always_apply and random.random() > self.p:
            return kwargs

        # Get random scale
        scale = 1 + random.uniform(0, self.scale_limit)

        # Get crop bbox around center of mass
        mask = kwargs['masks'][self.mask_index]
        center_h = math.floor(mask.mean(axis=1) * scale)
        center_w = math.floor(mask.mean(axis=0) * scale)
        h_start = max(0, center_h - mask.shape[0] // 2)
        w_start = max(0, center_w - mask.shape[1] // 2)
        h_end = min(kwargs['image'].shape[0], h_start + mask.shape[0])
        w_end = min(kwargs['image'].shape[1], w_start + mask.shape[1])

        # Copy image and masks
        image_larger = kwargs['image'].copy()
        masks_larger = [mask.copy() for mask in kwargs['masks']]

        # Scale up and crop around center of mass
        image_larger = cv2.resize(
            image_larger,
            (0, 0),
            fx=scale,
            fy=scale,
            interpolation=self.interpolation,
        )
        image_larger = F_crops.crop(
            image_larger,
            xmin=w_start,
            ymin=h_start,
            xmax=w_end,
            ymax=h_end,
        )

        for i in range(len(masks_larger)):
            masks_larger[i] = cv2.resize(
                masks_larger[i],
                (0, 0),
                fx=scale,
                fy=scale,
                interpolation=cv2.INTER_NEAREST,
            )
            masks_larger[i] = F_crops.crop(
                masks_larger[i],
                xmin=w_start,
                ymin=h_start,
                xmax=w_end,
                ymax=h_end,
            )

        # Get enlarged mask
        mask = masks_larger[self.mask_index] > 0

        # Copy masked area
        kwargs['image'][h_start:h_end, w_start:w_end][mask] = image_larger[mask]
        for i in range(len(kwargs['masks'])):
            kwargs['masks'][i][h_start:h_end, w_start:w_end][mask] = masks_larger[i][mask]

        return kwargs

class TtaHorizontalFlip:
    def apply(self, batch: torch.Tensor) -> torch.Tensor:
        # (N, C, H, W, D)
        return batch.flip(3)

    def apply_inverse_to_pred(self, batch_pred: torch.Tensor) -> torch.Tensor:
        # (N, C = 1, H, W)
        return batch_pred.flip(3)

class TtaVerticalFlip:
    def apply(self, batch: torch.Tensor) -> torch.Tensor:
        # (N, C, H, W, D)
        return batch.flip(2)

    def apply_inverse_to_pred(self, batch_pred: torch.Tensor) -> torch.Tensor:
        # (N, C = 1, H, W)
        return batch_pred.flip(2)

class TtaRotate90:
    def __init__(self, n_rot) -> None:
        assert n_rot % 4 != 0, f"n_rot should not be divisible by 4. Got {n_rot}"
        self.n_rot = n_rot % 4

    def apply(self, batch: torch.Tensor) -> torch.Tensor:
        # (N, C, H, W, D)
        return batch.rot90(self.n_rot, (2, 3))

    def apply_inverse_to_pred(self, batch_pred: torch.Tensor) -> torch.Tensor:
        # (N, C = 1, H, W)
        return batch_pred.rot90(-self.n_rot, (2, 3))

class TtaRotate:
    def __init__(self, limit_degrees: int = 90, fill_value=0.0) -> None:
        self.limit_degrees = limit_degrees
        self.fill_value = fill_value
        self.angle = None

    @staticmethod
    def _rotate(batch: torch.Tensor, angle: int, fill: float=0) -> torch.Tensor:
        N, C, H, W, D = batch.shape
        # (N, C, H, W, D) -> (N, C, D, H, W)
        batch = batch.permute(0, 1, 4, 2, 3)
        # (N, C, D, H, W) -> (N, C * D, H, W)
        batch = batch.reshape(N, C * D, H, W)
        batch = F_torchvision.rotate(
            batch,
            angle,
            interpolation=F_torchvision.InterpolationMode.BILINEAR,
            expand=False,
            fill=fill,
        )
        # (N, C * D, H, W) -> (N, C, D, H, W)
        batch = batch.reshape(N, C, D, H, W)
        # (N, C, D, H, W) -> (N, C, H, W, D)
        batch = batch.permute(0, 1, 3, 4, 2)
        return batch

    def apply(self, batch: torch.Tensor) -> torch.Tensor:
        assert self.angle is None, "TtaRotate should be applied only once."
        # (N, C, H, W, D)
        self.angle = random.randint(-self.limit_degrees, self.limit_degrees)
        return TtaRotate._rotate(batch, self.angle, fill=self.fill_value)

    def apply_inverse_to_pred(self, batch_pred: torch.Tensor) -> torch.Tensor:
        assert self.angle is not None, "TtaRotate should be applied before TtaRotate.apply_inverse_to_pred."
        # (N, C = 1, H, W) -> (N, C = 1, H, W, D = 1)
        batch_pred = batch_pred.unsqueeze(-1)
        # Fill with NaNs to ignore them in averaging
        batch_pred = TtaRotate._rotate(batch_pred, -self.angle, fill=torch.nan)
        # (N, C = 1, H, W, D = 1) -> (N, C = 1, H, W)
        batch_pred = batch_pred.squeeze(-1)
        self.angle = None
        return batch_pred

class Tta:
    def __init__(self, model, n_random_replays=1, use_hflip=True, use_vflip=True, rotate90_indices=None):
        assert n_random_replays > 0 or use_hflip or use_vflip or rotate90_indices is not None, \
            "At least one of n_random_replays > 0, "\
            "use_hflip or use_vflip or rotate90_indices is not None should be True."
        assert all([i > 0 and i <= 3 for i in rotate90_indices]), \
            f"rotate90_indices should be in [0, 3]. Got {rotate90_indices}"
        assert len(rotate90_indices) == len(set(rotate90_indices)), \
            f"rotate90_indices should not contain duplicates. Got {rotate90_indices}"
        self.model = model

        # Imagenet normalization during training is assumed
        fill_value = (0.0 - sum(IMAGENET_DEFAULT_MEAN) / 3) / (sum(IMAGENET_DEFAULT_STD) / 3)

        # All possible combinations of
        # - flips
        # - rotations on 90 degrees
        # - n_replays rotations on random angle
        rotates90 = [None]
        if rotate90_indices is not None:
            rotates90 = [None] + [
                TtaRotate90(i) for i in rotate90_indices
            ]
        flips = [None]
        if use_hflip:
            flips += [
                TtaHorizontalFlip(),
            ]
        if use_vflip:
            flips += [
                TtaVerticalFlip(),
            ]
        rotates = [None]
        if n_random_replays > 0:
            rotates = [None] + [
                TtaRotate(limit_degrees=45, fill_value=fill_value)
                for _ in range(n_random_replays)
            ]
        self.transforms = [
            rotates90,
            flips,
            rotates,
        ]

    def predict(self, batch: torch.Tensor) -> List[torch.Tensor]:
        preds = []

        # Apply TTA
        for transform_chain in itertools.product(*self.transforms):
            # Direct transform
            batch_aug = batch.clone()
            for transform in transform_chain:
                if transform is not None:
                    batch_aug = transform.apply(batch_aug)

            # Predict
            pred_aug = self.model(batch_aug)
            if pred_aug.ndim == 3:
                pred_aug = pred_aug.unsqueeze(1)

            # Inverse transform
            # Note: order of transforms is reversed
            for transform in reversed(transform_chain):
                if transform is not None:
                    pred_aug = transform.apply_inverse_to_pred(pred_aug)

            preds.append(pred_aug)

        return preds

    def __call__(self, batch: torch.Tensor) -> torch.Tensor:
        preds = self.predict(batch)

        # Average predictions, ignoring NaNs
        preds = torch.stack(preds, dim=0)
        preds = torch.nanmean(preds, dim=0)

        return preds


from cilcdjango.core.media import SharedMediaMixin, make_shared_media_url
from cilcdjango.core.util import get_app_setting

from django import forms
from django.conf import settings
from django.core.urlresolvers import reverse
from django.forms.widgets import flatatt
from django.utils.encoding import force_unicode, smart_unicode
from django.utils.html import escape
from django.utils.safestring import mark_safe
from django.utils.translation import ugettext as _

import datetime
import re

#-------------------------------------------------------------------------------
#  Helpers
#-------------------------------------------------------------------------------

def _update_attr_classes(attrs, classes):
	"""
	Add the classes contained in the `classes` iterable to the "class" key of
	the attributes passed to the widget's constructor in the `attrs` dict.
	"""
	if not attrs:
		attrs = {}
	attrs['class'] = u" ".join(set(re.split(r'\s+', attrs.get('class', '')) + classes)).strip()
	return attrs

#-------------------------------------------------------------------------------
#  Date / Time Widgets
#-------------------------------------------------------------------------------

class DatePickerWidget(forms.DateTimeInput, SharedMediaMixin):
	"""A widget that lets a user use a JavaScript date picker to select a date."""

	def render(self, name, value, attrs=None):

		#  Attempt to turn a previous date value into one capable of being set
		#  as the input field's value attribute.
		try:
			previous_date = value.strftime(self.format)
		except AttributeError:
			if value is None or value == "None":
				previous_date = ""
			else:
				previous_date = value
		current_value = "value='%s'" % escape(previous_date)

		#  Render the widget as a simple text field, which is transformed into a
		#  datepicker widget through JavaScript
		form_html = [u"<input type='text' name='%s' id='id_%s' size='12' class='uses-date-picker' %s />" % (name, name, current_value) ]
		return mark_safe(u'\n'.join(form_html))

	class SharedMedia:

		css = {
			'screen': ('core/css/libs/jquery_ui/jquery.ui.css',)
		}

		js = (
			'core/js/libs/jquery/jquery.ui.js',
			'core/js/widgets/date_picker.js'
		)

class SimpleTimeWidget(forms.TextInput):
	"""A simple time widget that ignores seconds."""

	_CLOCK_ICON = "core/images/clock_icon_tiny.png"
	_ICON_CLASSES = ["clock-icon"]

	def __init__(self, *args, **kwargs):

		default_args = {'size': 4, 'class': 'dt-time-input'}
		try:
			kwargs['attrs'].update(default_args)
		except KeyError:
			kwargs['attrs'] = default_args
		super(SimpleTimeWidget, self).__init__(*args, **kwargs)

		self.time_format = "%H:%M"

	def render(self, name, value, attrs=None):
		"""Show the field's time value in HH:MM format."""

		final_attrs = self.build_attrs(attrs, type=self.input_type, name=name)

		#  Attempt to transform an existing time value into a value capable of
		#  being an attribute of the input tag that this widget is rendered as.
		try:
			value = value.strftime(self.time_format)
		except AttributeError:
			try:
				value = datetime.datetime.strptime(value, "%H:%M").time().strftime(self.time_format)
			except (ValueError, AttributeError, TypeError):
				value = ''
		final_attrs['value'] = force_unicode(value)
		return mark_safe(
			u'<input%(input)s /> <img src="%(icon_src)s" class="%(icon_class)s" alt="%(icon_alt)s" />' % {
				'input': flatatt(final_attrs),
				'icon_src': make_shared_media_url(self._CLOCK_ICON),
				'icon_class': u" ".join(self._ICON_CLASSES),
				'icon_alt': _("clock icon")
			})

class DateTimeWidget(forms.MultiWidget):
	"""A widget that encapsulates a date picker and a time input."""

	def __init__(self, *args, **kwargs):

		kwargs['widgets'] = [ DatePickerWidget(format='%m/%d/%Y'), SimpleTimeWidget() ]
		super(DateTimeWidget, self).__init__(*args, **kwargs)

	def decompress(self, value):
		"""
		Break the datetime.datetime instance in `value` into its date and time
		components.
		"""
		try:
			return [value.date(), value.time()]
		except AttributeError:
			return [None, None]

#-------------------------------------------------------------------------------
#  Text Widgets
#-------------------------------------------------------------------------------

class RichTextEditorWidget(forms.Textarea, SharedMediaMixin):
	"""A textarea that uses CKEditor to allow the user to input rich text."""

	def __init__(self, editor_type="default", auto_focus=False, attrs=None):

		super(RichTextEditorWidget, self).__init__(attrs)
		self.auto_focus  = auto_focus
		self.editor_type = editor_type

	def render(self, name, value, attrs=None):

		final_attrs = self.build_attrs(attrs)
		final_attrs['name'] = name

		#  Set custom classes for the textarea for use by the JavaScript code
		classes = getattr(final_attrs, 'class', [])
		classes.append("uses-rte")
		classes.append("editor-%s" % self.editor_type)
		if self.auto_focus:
			classes.append('auto-focus')
		final_attrs['class'] = " ".join(classes)

		html = [u'<textarea%s>%s</textarea>' % (flatatt(final_attrs), escape(smart_unicode(value or '')))]
		return mark_safe(u'\n'.join(html))

	class SharedMedia:

		js = (
			'core/js/libs/ckeditor/ckeditor.js',
			get_app_setting('RTE_CONFIG_FILE'),
			'core/js/widgets/rich_text.js'
		)

class WordCounterWidget(forms.Textarea, SharedMediaMixin):
	"""
	A textarea that counts the number of words entered and updates this count as
	the user types.
	"""

	_CSS_CLASSES = {
		'widget':    "uses-word-counter",
		'max_words': "max-words_",
		'min_words': "min-words_"
	}

	def __init__(self, min_words=None, max_words=None, attrs=None):
		"""
		Add special classes to the textarea that are used by JavaScript to
		generate the markup for the word counter.
		"""

		#  Add classes for the max and min words, if they are specified
		classes = []
		if min_words or max_words:
			classes.append(self._CSS_CLASSES['widget'])
		if min_words:
			classes.append("%s%s" % (self._CSS_CLASSES['min_words'], min_words))
		if max_words:
			classes.append("%s%s" % (self._CSS_CLASSES['max_words'], max_words))

		super(WordCounterWidget, self).__init__(
			attrs=_update_attr_classes(attrs, classes)
		)

	class SharedMedia:

		js = (
			'core/js/widgets/word_counter.js',
		)


import smtplib
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from config import SMTP_CONFIG
try:
    from mainconfig import ACCOUNT, MAIL_SERVER, CC
except Exception as e:
    from testconfig import ACCOUNT, MAIL_SERVER, CC

def send_email(email=None, subject=None, body=None, include_cc=False):
    mimemsg = MIMEMultipart()
    mimemsg['From'] = 'HR_MANAGER' + '<' + ACCOUNT['email'] + '>'
    mimemsg['To'] = email
    if include_cc and CC:
        mimemsg['Cc'] = CC
    mimemsg['Subject'] = subject
    mimemsg.attach(MIMEText(body, 'plain'))
    try:
        connection = smtplib.SMTP(host=MAIL_SERVER['host'], port=MAIL_SERVER['port'])
        connection.starttls()
        if MAIL_SERVER['host'] == 'smtp.office365.com':
            connection.login(ACCOUNT['email'], ACCOUNT['password'])
        connection.send_message(mimemsg)
        connection.quit()
        return True

    except Exception as e:
        print(e)
        return False


import argparse
import a_star_search as search
from environments import gridEnvP1

parser = argparse.ArgumentParser(prog = 'python3 main.py',
                                 description = 'Implements A* search on a specified dataset.')
parser.add_argument('-d', '--data',
                    default = "./data/100 nodes.csv",
                    dest = "path_to_data",
                    required = True,
                    help = "path to states/environment data.")
parser.add_argument('-s', '--start-state',
                    type = int,
                    required = True,
                    help = "an integer that indicates the starting state.")
parser.add_argument('-g', '--goal-state',
                    type = int,
                    required = True,
                    help = "an integer that indicates the goal state.")
parser.add_argument('-w', '--weight',
                    type = float,
                    required = True,
                    help = "the weight parameter for A* search\
                            to calculate f = w*g + (1-w)*h.\
                            must be between 0 and 1.")
parser.add_argument('-v', '--verbose',
                    action = 'store_true',
                    help = "prints information such as frontier details and best node for every iteration.")

args = vars(parser.parse_args())
filepath = args["path_to_data"]
start = args["start_state"]
goal = args["goal_state"]
w = args["weight"]
v = args["verbose"]

env = gridEnvP1()
env.load(filepath)
env.initialize(start, goal)
result = search.astar_tree_search(env, w, v)

if result:
	result_node = result[0]
	result_frontier = result[1]
	print("Path length: " + str(len(result_node.path)))
	print("Path cost: " + str(result_node.cost))
	print("Path: " + str(result_node.path))
	print("Nodes generated: " + str(result_frontier.generated_nodes))
else:
	print("There is no path")


# -*- coding: utf-8 -*-
from __future__ import unicode_literals

from django.db import migrations, models

class Migration(migrations.Migration):

    dependencies = [
        ('playy', '0020_auto_20151124_2126'),
    ]

    operations = [
        migrations.RenameField(
            model_name='match',
            old_name='match_name2',
            new_name='match_name',
        ),
        migrations.RemoveField(
            model_name='player',
            name='points',
        ),
    ]


import torch

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
TRAIN_DIR = "data/clean/cleaned_train.csv"
PRED_DIR = "data/clean/cleaned_test.csv"
LEARNING_RATE = 1e-3
BATCH_SIZE = 20
NUM_WORKERS = 1
LOAD_MODEL = False
SAVE_MODEL = True
CHECKPOINT = "checkpoints/neural_network/checkpoint.tar"
EPOCHS = 150
FOLDS = 5
EXPERIMENT = False
TRAIN = True
INFERENCE = True


from lexicons import english_fn_database, german_fn_database, french_fn_database

logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
                    datefmt='%m/%d/%Y %H:%M:%S',
                    level=logging.INFO)

logger = logging.getLogger(__name__)

parser = argparse.ArgumentParser()

parser.add_argument("--lexicon_dir", default=None, type=str, required=True,
                    help="the directory with FrameNet lexicon.")

args = parser.parse_args()

def get_lexicon(language, lexicon_dir):
    if language == "english":
        lexicon = english_fn_database.EnglishFrameNet(lexicon_dir)
    if language == "german":
        lexicon = german_fn_database.GermanFrameNet(lexicon_dir)
    if language == "french":
        lexicon = french_fn_database.FrenchFrameNet(lexicon_dir)
    return lexicon

if __name__ == '__main__':
    lexicon = get_lexicon(args.language, args.lexicon_dir)


'''
history =  2013.06.05, 2013.06.14, 2013.06.25, 2013.09.14, 2013.09.15, 2013.12.10, 2015.03.02, 2017.09.28
'''
# there is a "#@" follow the list_name_in_code at the end of the code in processing_code
# because list_name_in_code is the morker of the "file_name" list start_side.
# And "#@" is the marker of the "file_name" list end_side.
# That will be using in modification of the origen processing_code.
#-----------------------------------------------------------------------------------------------
import Raxpy3Libbasic_2015_03_02 as RaxLib
import platform
#-----------------------------------------------------------------------------------------------

def multiprocessing(jobs_list = [], processing_code = "", list_name_in_code = "", CUP_N = 0, deleteTempNorY = 'Y'):

    ori_path = RaxLib.os.getcwd()
    sysstr = platform.system()

    #-----------------------------------------------------------------------------------------------s
    #seperate jobs_list
    RaxLib.os.chdir(ori_path)

    modify_jobs_list = ['@##@n@##@@'.join(x.split('\n')) for x in jobs_list]
    modify_jobs_list = ['#@#@t@#@#@'.join(x.split('\t')) for x in modify_jobs_list]
    list_box = RaxLib.listPartSeparate(modify_jobs_list, CUP_N)

    showTime1 = RaxLib.LoopingTime()
    showTime1.open('build job_part ', list_box)

    i = 1
    for job in list_box:
        t = RaxLib.time.time()

        RaxLib.newOrderfile(['order', 'job'], [[x] for x in job], 'job_part_'+str(i)+'.txt')

        t1 = RaxLib.time.time()-t
        showTime1.showEndTime(t1)
        i += 1

    #-----------------------------------------------------------------------------------------------
    #open job_part file

    file_list = RaxLib.os.listdir(ori_path)
    sub_file_list = [x for x in file_list if 'job_part_' in x]

    for job_file in sub_file_list:

        number = job_file.split('job_part_')[1].split('.txt')[0]

        table = RaxLib.tabfilewithtitle()
        table.open(job_file, 'order')

        A = '['+','.join(["'"+'\\t'.join('\\n'.join(x.split('@##@n@##@@')).split('#@#@t@#@#@'))+"'" for x in table.readrow('Y', 'job') if x != 'job'])+']'

        #-----------------------------------------------------------------------------------------------
        #seperate processing_code

        P1 = processing_code.replace("'\n'","'\\n'").split(list_name_in_code+' =')
        P2 = P1[1].split('#@')
        P2.pop(0)
        B = '#@'.join(P2)
        P1[1] = A +'#@'+ B
        midify_code = (list_name_in_code+' =').join(P1)

        #-----------------------------------------------------------------------------------------------
        #add check_point_file
        end_info = '''
#-----------------------------------------------------------------------------------------------
os.chdir(r"'''+ori_path+'''")
out_put_file_name= 'py_end_'+str('''+str(number)+''')+'.py'
f = open(out_put_file_name, 'w')
f.write(' ')
f.close()
'''
        midify_code = midify_code+'\n'+end_info
        out_put_file_name = 'py_'+str(number)+'.py'
        f = open(out_put_file_name, 'w')
        f.write(midify_code)
        f.close()

    #-----------------------------------------------------------------------------------------------
    #start python file

    file_list = RaxLib.os.listdir(ori_path)
    py_file_list = [x for x in file_list if 'py_' in x]
    print(py_file_list)
    for py_file in py_file_list:
        print('python3 '+py_file+' &')
        RaxLib.os.system('python3 '+py_file+' &')

    #-----------------------------------------------------------------------------------------------
    #check processing end or not
    s = 1
    while s < 43200:
        end_file_list = RaxLib.os.listdir(ori_path)
        checkfilelist = [x for x in end_file_list if 'py_end_' in x]
        if len(checkfilelist) == len(list_box):

            #-----------------------------------------------------------------------------------------------
            #delete job_list_.txt file
            RaxLib.os.chdir(ori_path)
            i = 1
            while i < len(list_box)+1:
                if deleteTempNorY == 'N':
                    pass
                elif deleteTempNorY == 'Y':
                    if sysstr == 'Windows':
                        cmd = 'del'
                    elif sysstr == 'Linux':
                        cmd = 'rm'
                    RaxLib.os.system(cmd+' job_part_'+str(i)+'.txt')
                    RaxLib.os.system(cmd+' py_'+str(i)+'.py')
                    RaxLib.os.system(cmd+' py_end_'+str(i)+'.py')
                else:
                    print("please Enter 'N' or 'Y', Thanks~! ")
                i += 1

            print('processing end')
            return

        else:

            RaxLib.time.sleep(10)
            s += 10

    #-----------------------------------------------------------------------------------------------
    #delete job_list_.txt file
    RaxLib.os.chdir(ori_path)
    i = 1
    while i < len(list_box)+1:
        if deleteTempNorY == 'N':
            pass
        elif deleteTempNorY == 'Y':
            RaxLib.os.system('rm job_part_'+str(i)+'.txt')
            RaxLib.os.system('rm py_'+str(i)+'.py')
            RaxLib.os.system('rm py_end_'+str(i)+'.py')
        else:
            print("please Enter 'N' or 'Y', Thanks~! ")
        i += 1

    print('processing end')
    return

if __name__ == '__main__':
    #-----------------------------------------------------------------------------------------------
    #to list directory in "C:\Program Files" and separate into 3 parts to run
    #then get the response in "./Result"
    import os

    while True:
        ori_path = os.getcwd()
        os.chdir(r"C:\Program Files")
        E_ = [x for x in os.listdir(os.getcwd()) if os.path.isdir(x) is True]

        processing_code = '''

import Raxpy3Libbasic_2015_03_02 as Lib
import os
ori_path = os.getcwd()
os.chdir(r"C:\Program Files")
E_ = [x for x in os.listdir(os.getcwd()) if os.path.isdir(x) is True]#@
for x in E_:
    #t = time.time()
    C = Lib.ReadDirectory(ori_path)
    C.read()
    os.chdir(ori_path)
    if 'Result' not in os.listdir(os.getcwd()):
        os.mkdir('Result')
    os.chdir(ori_path+'\\\\'+'Result')
    f = open(x+'_directory_and_file_list.txt', 'w')
    f.write('\\n'.join(C.file_list))
    f.close()
    #print(time.time()-t)

        '''

        os.chdir(ori_path)
        multiprocessing(E_, processing_code, 'E_', 3, 'N')


from deseq import DESeq
from plots import Plots
import json
import numpy as np
import pandas as pd
from utils import replace_colnames, txt_2_dataframe
from pydeseq2.dds import DeseqDataSet
from pydeseq2.ds import DeseqStats
from sanbomics.tools import id_map
import scanpy as sc
import gseapy as gp
from gseapy.plot import gseaplot
import seaborn as sns
from gseapy import dotplot
import matplotlib.pyplot as plt
from datetime import datetime
from loguru import logger
date = datetime.now().strftime("%Y_%m_%d-%I_%M_%S_%p")

def run_pipeline(CONFIG_PATH: str):
    """_summary_

    Args:
        CONFIG_PATH (str): _description_
    """

    logger.info("Loading configuration parameters.")

    # Open config file
    config = json.load(open(CONFIG_PATH, "r"))

    # DEG Parameters
    control = config["differential_expression"]["Condition"]["control"]
    experiment = config["differential_expression"]["Condition"]["experiment"]
    species = config["differential_expression"]["species"]
    padj = config["differential_expression"]["thresholds_for_significant_genes"]["padj"]
    log2FoldChange = config["differential_expression"]["thresholds_for_significant_genes"]["log2FoldChange"]
    baseMean = config["differential_expression"]["thresholds_for_significant_genes"]["baseMean"]
    count_matrix_path = config["differential_expression"]["count_matrix"]["path"]
    count_matrix_gene_column = config["differential_expression"]["count_matrix"]["gene_id_column_name"]

    # Gene Set enrichment analysis parameters
    perform_GSEA = config["gene_set_enrichment_analysis"]["perform_GSEA"]
    gene_sets = config["gene_set_enrichment_analysis"]["gene_sets"]
    permutations_num = config["gene_set_enrichment_analysis"]["permutations"]
    filter_by_organ = config["gene_set_enrichment_analysis"]["filter_by_organ"]

    # Plotting parameters
    plot_volcano = config["plotting"]["plot_volcano"]["plot"]
    plot_pca = config["plotting"]["plot_pca"]["plot"]
    plot_heatmap = config["plotting"]["plot_heatmap"]["plot"]
    plot_dotplot = config["plotting"]["plot_dotplot"]["plot"]
    plot_gsea = config["plotting"]["plot_gsea"]["plot"]

    ##########################################################
    # Read input data from bash processing
    count_matrix = pd.read_csv(r"{}".format(count_matrix_path), sep="\t")
    # diff_exp_matrix = txt_2_dataframe(diff_exp_matrix_path)

    # Convert ENSG to gene names
    ensg_to_gene_mapper = id_map(species=species)

    # Instatiate DESeq class
    deseq = DESeq(gene_sets, permutations_num,
                  config, ensg_to_gene_mapper)

    # Pre-process count matrix
    count_matrix_processed, metadata = deseq.pre_processing(
        count_matrix=count_matrix, gene_id_col_name=count_matrix_gene_column)

    # Run DEG method
    dds, stat_res, dds_stats_df, significant_genes, ranked_by_stat = deseq.DEG(
        count_matrix_processed=count_matrix_processed, metadata=metadata, experiment=experiment, control=control, padj=padj, log2FoldChange=log2FoldChange, baseMean=baseMean)

    # Run GSEA method if true
    if perform_GSEA:

        go_ontologies, go_ontologies_df = deseq.GSEA(
            ranked_by_stat, filter_by_organ)

        vis = Plots(plot_volcano, plot_pca,
                    plot_heatmap, plot_gsea, plot_dotplot, config)
        vis.run_plots(dds, dds_stats_df,
                      significant_genes, ensg_to_gene_mapper, go_ontologies)
    else:
        vis = Plots(plot_volcano, plot_pca,
                    plot_heatmap, plot_gsea=False, plot_dotplot=False, config=config)
        vis.run_plots(dds, dds_stats_df,
                      significant_genes, ensg_to_gene_mapper)

    logger.info("DESeq pipeline complete.")

if __name__ == "__main__":
    run_pipeline(CONFIG_PATH=r".\src\scripts\config.json")


import tkinter as tk
from tkinter import messagebox
from S_AES import *
import tkinter as tk
from tkinter import messagebox
from time import time
result="error"#默认返回结果
state="encrypt"#默认为加密模式
input_mode="bits"#默认为二进制数输入模式
num_keys=1
Locked_Correct_key:bool=False
work_mode="ECB"
init_vec="1001110011100010"
all_mingwen:list[str]=[]#存储破解模式中接收到的明文组
all_miwen:list[str]=[]#存储破解模式中接收到的密文组
possible_key:list[str]=[]#存储可能的密钥
def multi_keys():
    global num_keys
    num_keys=num_keys+1
    if num_keys==4:
        num_keys=1
    messagebox.showinfo("提示",'现在有{}重密码,请按顺序输入所有的密码'.format(num_keys))
def change_word_mode():
     global work_mode
     if work_mode=="ECB":
        work_mode="CBC"
        messagebox.showinfo("提示","现在是CBC模式")
     else:
        work_mode="ECB"
        messagebox.showinfo("提示","现在是ECB模式")
def clear_records():
    global all_mingwen,all_miwen
    all_miwen.clear()
    all_mingwen.clear()
    possible_key.clear()
    print(all_mingwen,all_mingwen)
    tk.Tk().withdraw()
    messagebox.showinfo("提示","已经清除所有信息")
def Save_info(event,mingwen,miwen):
    global all_mingwen
    global all_miwen
    #print("获取到的明文",mingwen)
    #print("获取到的密文",miwen)
    all_mingwen.append(mingwen)
    all_miwen.append(miwen)
    messagebox.showinfo("提示","保存成功")
def Convert():
    global state
    if state=="encrypt":
        state="decrypt"
        tk.Tk().withdraw()
        messagebox.showinfo("提示","现在是解密模式")
    else:
        state="encrypt"
        tk.Tk().withdraw()
        messagebox.showinfo("提示","现在是加密模式")
def Change_input_mode():
    global input_mode
    if input_mode=="bits":
        input_mode="ASCII"
        tk.Tk().withdraw()
        messagebox.showinfo("提示","现在是字符模式")
    else:
        input_mode="bits"
        tk.Tk().withdraw()
        messagebox.showinfo("提示","现在是二进制模式")
def get_result():#获得转换后的结果
    key=entry1.get()
    info=entry2.get()
    global result
    times=len(key)/16
    if times!=num_keys:
        messagebox.showinfo("提示","密码长度不正确，请检查输入")
    if input_mode=='bits':#二进制模式
        if state=="encrypt":
            if times==1:
                result=Encrypt(info,key[0:16])
            elif times==2:
                result=Encrypt(info,key[0:16])
                result=Encrypt(result,key[16:32])
            elif times==3:
                result=Encrypt(info,key[0:16])
                result=Encrypt(result,key[16:32])
                result=Encrypt(result,key[32:48])
        else:
            if times==1:
                result=Decrypt(info,key)
            elif times==2:
                result=Decrypt(result,key[16:32])
                result=Decrypt(result,key[0:16])
            elif times==3:
                result=Decrypt(info,key[32:48])
                result=Decrypt(result,key[16:32])
                result=Decrypt(result,key[0:16])
    elif input_mode=="ASCII":#字符串模式
        if len(info)%2==1:
            messagebox.showinfo("提示","请输入偶数个字符")
            return None
        info_bit=AsciiToBit(info)#先将字符输入转换为二进制输入
        if work_mode=="CBC":
            info_bit_CBC=[]
            info_bit_CBC_p=""
            for i in range(16):
                info_bit_CBC_p+=str(int(init_vec[i])^int(info_bit[0][i]))
            info_bit_CBC.append(info_bit_CBC_p)
            for i in range(1,len(info_bit)):
                info_bit_CBC_p=""
                for j in range(16):
                    info_bit_CBC_p+=str(int(info_bit_CBC[i-1][j])^int(info_bit[i][j]))
                info_bit_CBC.append(info_bit_CBC_p)
        #print(init_vec,info_bit)
        #print(info_bit_CBC)
            for i in range(len(info_bit)):
                info_bit[i]=info_bit_CBC[i]
        result_bit:list[str]=[]
        for item in info_bit:
            if state=="encrypt":
                if times==1:
                    result_bit_part=Encrypt(item,key[0:16])
                elif times==2:
                    result=Encrypt(item,key[0:16])
                    result_bit_part=Encrypt(result,key[16:32])
                elif times==3:
                    result=Encrypt(item,key[0:16])
                    result=Encrypt(result,key[16:32])
                    result_bit_part=Encrypt(result,key[32:48])
            else:
                if times==1:
                    result_bit_part=Decrypt(item,key)
                elif times==2:
                    result=Decrypt(item,key[16:32])
                    result_bit_part=Decrypt(result,key[0:16])
                elif times==3:
                    result=Decrypt(item,key[32:48])
                    result=Decrypt(result,key[16:32])
                    result_bit_part=Decrypt(result,key[0:16])
            result_bit.append(result_bit_part)
        result=BitToAscii(result_bit)#将二进制的转换结果还原为对应的字符
    result_.set(result)#在文本框中显示结果

def AsciiToBit(string:str):
    string_bit:list[str]=[]
    for i in range(0,len(string),2):
        char_byte1=bin(ord(string[i]))
        char_byte1=char_byte1[2:]#去标志位"0b"
        char_byte1=(8-len(char_byte1))*"0"+char_byte1#补零
        char_byte2=bin(ord(string[i+1]))
        char_byte2=char_byte2[2:]#去标志位"0b"
        char_byte2=(8-len(char_byte2))*"0"+char_byte2#补零
        string_bit.append(char_byte1+char_byte2)
    return string_bit

def BitToAscii(string_bit:list[str]):
    string=""
    string_bit_new=[]
    for item in string_bit:
        string_bit_new.append(item[0:8])
        string_bit_new.append(item[8:16])
    for byte in string_bit_new:
        char=chr(int(byte,2))
        string+=char
    return string
def Violent_Crack():
    global Locked_Correct_key
    Locked_Correct_key=False
    print("获取到的明文列表",all_mingwen)
    print("获取到的密文列表",all_miwen)
    start_time=time()
    for k in range(len(all_mingwen)):
        if Locked_Correct_key==False:
            test_mingwen=all_mingwen[k]
            test_miwen=all_miwen[k]
            print(f"正在对第{k+1}对明密文进行破解")
            find_possbile_key(test_mingwen,test_miwen)
            print("现在可能的密钥有",possible_key)
    end_time=time()
    run_time=end_time-start_time
    #print(len(possible_key))
    if len(possible_key)==1:
        messagebox.showinfo("提示","破解用时{}\n密钥为{}".format(run_time,possible_key))
    elif len(possible_key)==0:
        messagebox.showinfo("提示","没有找到合理密钥,请检查输入")
    else:#输入较少时可能用多个密钥符合要求
        str1="密钥为"+possible_key[0]
        for i in range(len(possible_key)-1):
            str1=str1+"或"+possible_key[i+1]
        tk.Tk().withdraw()
        messagebox.showinfo("提示",str1+f"\n破解用时{run_time}")
    possible_key.clear()
def find_possbile_key(mingwen,miwen):
    global Locked_Correct_key
    i=0
    global possible_key
    while i<65536:#最大的16位二进制数转换为十进制即为65536
        test_key=bin(i)
        test_key=test_key[2:]#去标志位"0b"
        length=len(test_key)
        gap=16-length
        test_key=gap*"0"+test_key#补零
        if Encrypt(mingwen,test_key)==miwen:
            #print(f"发现密钥{test_key}满足条件")
            if test_key not in possible_key:#样本较少时可能有多个密钥满足条件
                possible_key.append(test_key)
            else:#一旦出现重复即可确定唯一正确的密钥
                possible_key.clear()
                possible_key.append(test_key)
                print("已经锁定唯一正确的密钥",test_key)
                Locked_Correct_key=True
                break
        i+=1
def create_new_windows():
    new_window=tk.Toplevel(root)
    new_window.title("破解模式")
    new_window.geometry('400x200+500+500')
    mingwen_va=tk.StringVar()#明文文本框中的字符串变量
    miwen_va=tk.StringVar()#密文文本框中的字符串变量
    entry_mingwen=tk.Entry(new_window,textvariable=mingwen_va)
    entry_miwen=tk.Entry(new_window,textvariable=miwen_va)
    entry_mingwen.place(relx=0.05,rely=0.05,width=200,height=20)
    entry_miwen.place(relx=0.05,rely=0.15,width=200,height=20)
    tk.Label(new_window,text="此处输入明文",font=('Simhei.ttf',10)).place(relx=0.6,rely=0.05,width=100,height=20)
    tk.Label(new_window,text="此处输入密文",font=('Simhei.ttf',10)).place(relx=0.6,rely=0.15,width=100,height=20)
    button5=tk.Button(new_window,text="开始破解",fg="black",font=("SimHei.ttf",10),command=Violent_Crack)
    button5.place(relx=0.05,rely=0.35,width=100,height=20)
    button6=tk.Button(new_window,text="保存并继续添加",fg="black",font=("SimHei.ttf",10))
    button6.place(relx=0.05,rely=0.22,width=100,height=20)
    button6.bind("<Button-1>",lambda event:Save_info(event,entry_mingwen.get(),entry_miwen.get()))#使用lambda方法来传入Save_info函数需要的参数
    button7=tk.Button(new_window,text="清除保存的信息",fg="black",font=("SimHei.ttf",10),command=clear_records)
    button7.place(relx=0.05,rely=0.45,width=100,height=20)
root=tk.Tk()
root.title("S-AES算法演示")
root.geometry("640x360+200+200")
info_va=tk.StringVar()#待转换信息文本框中的字符串变量
key_va=tk.StringVar()#密钥文本框中的字符串变量
result_=tk.StringVar()#结果文本框中的字符串变量
entry1=tk.Entry(root,textvariable=key_va,width=350)
entry2=tk.Entry(root,textvariable=info_va,width=200)
entry3=tk.Entry(root,state='readonly',text=result_,width=200)
entry1.place(relx=0.1,rely=0.1,width=350,height=20)
entry2.place(relx=0.1,rely=0.2,width=200,height=20)
entry3.place(relx=0.1,rely=0.3,width=200,height=20)
tk.Label(root,text="此处输入密钥",font=('Simhei.ttf',10)).place(relx=0.7,rely=0.1,width=100,height=20)
tk.Label(root,text="此处输入明文/密文",font=('Simhei.ttf',10)).place(relx=0.4,rely=0.2,width=120,height=20)
tk.Label(root,text="转换后的结果",font=('Simhei.ttf',10)).place(relx=0.4,rely=0.3,width=100,height=20)
Button_fun1=tk.Button(root,text="转换加密/解密模式",fg="black",font=("SimHei.ttf",10),command=Convert)
Button_fun1.place(relx=0.25,rely=0.4,width=125,height=20)
Button_fun2=tk.Button(root,text="开始计算",fg="black",font=("SimHei.ttf",10),command=get_result)
Button_fun2.place(relx=0.1,rely=0.4,width=75,height=20)
Button_fun3=tk.Button(root,text="转换输入模式",fg='black',font=("SimHei.ttf",10),command=Change_input_mode)
Button_fun3.place(relx=0.1,rely=0.5,width=100,height=20)
Button_fun4=tk.Button(root,text="进入破解模式",fg="black",font=("SimHei.ttf",10),command=create_new_windows)
Button_fun4.place(relx=0.28,rely=0.5,width=100,height=20)
Button_fun5=tk.Button(root,text="转换多重密码",font=("SimHei.ttf",10),command=multi_keys)
Button_fun5.place(relx=0.1,rely=0.6,width=100,height=20)
Button_fun6=tk.Button(root,text="转换工作模式",font=("SimHei.ttf",10),command=change_word_mode)
Button_fun6.place(relx=0.28,rely=0.6,width=100,height=20)
root.mainloop()


# -*- coding: utf-8 -*-
from __future__ import unicode_literals

from unittest import TestCase

try:
    from unittest.mock import MagicMock, patch
except ImportError:
    from mock import MagicMock, patch

from bundlewrap.exceptions import ItemDependencyError, NodeAlreadyLockedException, RepositoryError
from bundlewrap.group import Group
from bundlewrap.items import Item
from bundlewrap.node import ApplyResult, apply_items, _flatten_group_hierarchy, Node, NodeLock
from bundlewrap.operations import RunResult
from bundlewrap.repo import Repository
from bundlewrap.utils import names

class MockNode(object):
    name = "mocknode"

class MockBundle(object):
    name = "mock"
    bundle_dir = ""
    bundle_data_dir = ""
    items = []

class MockItem(Item):
    BUNDLE_ATTRIBUTE_NAME = "mock"
    ITEM_TYPE_NAME = "type1"
    NEEDS_STATIC = []
    _APPLY_RESULT = Item.STATUS_OK

    def apply(self, *args, **kwargs):
        return self._APPLY_RESULT
del Item.__reduce__  # we don't need the custom pickle-magic for our
                     # MockItems

def get_mock_item(itype, name, deps_static, deps):
    bundle = MockBundle()
    bundle.node = MockNode()
    item = MockItem(bundle, name, {'needs': deps}, skip_validation=True)
    item.ITEM_TYPE_NAME = itype
    item.NEEDS_STATIC = deps_static
    return item

class ApplyItemsTest(TestCase):
    """
    Tests bundlewrap.node.apply_items.
    """
    def test_self_loop(self):
        i1 = get_mock_item("type1", "name1", [], ["type1:name1"])
        i2 = get_mock_item("type1", "name2", [], [])
        node = MagicMock()
        node.items = [i1, i2]
        with self.assertRaises(ItemDependencyError):
            list(apply_items(node))

    def test_direct_loop(self):
        i1 = get_mock_item("type1", "name1", [], ["type1:name2"])
        i2 = get_mock_item("type1", "name2", [], ["type1:name1"])
        node = MagicMock()
        node.items = [i1, i2]
        with self.assertRaises(ItemDependencyError):
            list(apply_items(node))

    def test_nested_loop(self):
        i1 = get_mock_item("type1", "name1", [], ["type1:name2"])
        i2 = get_mock_item("type1", "name2", [], ["type1:name3"])
        i3 = get_mock_item("type1", "name3", [], ["type1:name4"])
        i4 = get_mock_item("type1", "name4", [], ["type1:name1"])
        node = MagicMock()
        node.items = [i1, i2, i3, i4]
        with self.assertRaises(ItemDependencyError):
            list(apply_items(node))

    def test_implicit_loop(self):
        i1 = get_mock_item("type1", "name1", [], ["type1:name2"])
        i2 = get_mock_item("type1", "name2", [], ["type1:"])
        node = MagicMock()
        node.items = [i1, i2]
        with self.assertRaises(ItemDependencyError):
            list(apply_items(node))

    def test_simple_order(self):
        i1 = get_mock_item("type1", "name1", [], ["type1:name2"])
        i2 = get_mock_item("type1", "name2", [], ["type1:name3"])
        i3 = get_mock_item("type1", "name3", [], [])

        node = MagicMock()
        node.items = [i1, i2, i3]

        results = list(apply_items(node))

        self.assertEqual(results[0][0], "type1:name3")
        self.assertEqual(results[1][0], "type1:name2")
        self.assertEqual(results[2][0], "type1:name1")

    def test_implicit_order(self):
        i1 = get_mock_item("type1", "name1", [], ["type1:name2"])
        i2 = get_mock_item("type1", "name2", [], [])
        i3 = get_mock_item("type2", "name3", ["type1:"], [])

        node = MagicMock()
        node.items = [i1, i2, i3]

        results = list(apply_items(node))

        self.assertEqual(results[0][0], "type1:name2")
        self.assertEqual(results[1][0], "type1:name1")
        self.assertEqual(results[2][0], "type2:name3")

    def test_apply_parallel(self):
        i1 = get_mock_item("type1", "name1", [], ["type1:name2"])
        i2 = get_mock_item("type1", "name2", [], ["type1:name3"])
        i3 = get_mock_item("type1", "name3", [], [])

        node = MagicMock()
        node.items = [i1, i2, i3]

        results = list(apply_items(node, workers=2))

        self.assertEqual(results[0][0], "type1:name3")
        self.assertEqual(results[1][0], "type1:name2")
        self.assertEqual(results[2][0], "type1:name1")

    def test_apply_interactive(self):
        i1 = get_mock_item("type1", "name1", [], ["type1:name2"])
        i2 = get_mock_item("type1", "name2", [], ["type1:name3"])
        i3 = get_mock_item("type1", "name3", [], [])

        node = MagicMock()
        node.items = [i1, i2, i3]

        results = list(apply_items(node, interactive=True, profiling=True))

        self.assertEqual(results[0][0], "type1:name3")
        self.assertEqual(results[1][0], "type1:name2")
        self.assertEqual(results[2][0], "type1:name1")

class ApplyResultTest(TestCase):
    """
    Tests bundlewrap.node.ApplyResult.
    """
    def test_correct(self):
        item_results = (
            ("item_id", Item.STATUS_OK, None),
        )
        output_result = ApplyResult(MagicMock(), item_results)
        self.assertEqual(output_result.correct, 1)
        self.assertEqual(output_result.fixed, 0)
        self.assertEqual(output_result.skipped, 0)
        self.assertEqual(output_result.failed, 0)

    def test_fixed(self):
        item_results = (
            ("item_id", Item.STATUS_FIXED, None),
        )
        output_result = ApplyResult(MagicMock(), item_results)
        self.assertEqual(output_result.correct, 0)
        self.assertEqual(output_result.fixed, 1)
        self.assertEqual(output_result.skipped, 0)
        self.assertEqual(output_result.failed, 0)

    def test_skipped(self):
        item_results = (
            ("item_id", Item.STATUS_SKIPPED, None),
        )
        output_result = ApplyResult(MagicMock(), item_results)
        self.assertEqual(output_result.correct, 0)
        self.assertEqual(output_result.fixed, 0)
        self.assertEqual(output_result.skipped, 1)
        self.assertEqual(output_result.failed, 0)

    def test_failed(self):
        item_results = (
            ("item_id", Item.STATUS_FAILED, None),
        )
        output_result = ApplyResult(MagicMock(), item_results)
        self.assertEqual(output_result.correct, 0)
        self.assertEqual(output_result.fixed, 0)
        self.assertEqual(output_result.skipped, 0)
        self.assertEqual(output_result.failed, 1)

    def test_bs(self):
        item_results = (
            ("item_id", 777, None),
        )
        with self.assertRaises(RuntimeError):
            ApplyResult(MagicMock(), item_results)

class FlattenGroupHierarchyTest(TestCase):
    """
    Tests bundlewrap.node._flatten_group_hierarchy.
    """
    def test_reorder_chain(self):
        group1 = MagicMock()
        group1.name = "group1"
        group2 = MagicMock()
        group2.name = "group2"
        group3 = MagicMock()
        group3.name = "group3"

        group3.subgroups = [group2]
        group2.subgroups = [group1]
        group1.subgroups = []

        self.assertEqual(
            _flatten_group_hierarchy([group1, group2, group3]),
            ["group3", "group2", "group1"],
        )

    def test_reorder_short_chain(self):
        group1 = MagicMock()
        group1.name = "group1"
        group2 = MagicMock()
        group2.name = "group2"
        group3 = MagicMock()
        group3.name = "group3"

        group3.subgroups = [group2]
        group2.subgroups = []
        group1.subgroups = []

        order = _flatten_group_hierarchy([group1, group2, group3])
        self.assertTrue(order.index("group3") < order.index("group2"))

    def test_loop(self):
        group1 = MagicMock()
        group1.name = "group1"
        group2 = MagicMock()
        group2.name = "group2"
        group3 = MagicMock()
        group3.name = "group3"

        group3.subgroups = [group2]
        group2.subgroups = [group1]
        group1.subgroups = [group3]

        with self.assertRaises(RuntimeError):
            _flatten_group_hierarchy([group1, group2, group3])

class InitTest(TestCase):
    """
    Tests initialization of bundlewrap.node.Node.
    """
    @patch('bundlewrap.node.validate_name', return_value=False)
    def test_bad_bundle_name(self, *args):
        with self.assertRaises(RepositoryError):
            Node("name")

class NodeTest(TestCase):
    """
    Tests bundlewrap.node.Node.
    """
    @patch('bundlewrap.node.ApplyResult')
    @patch('bundlewrap.node.NodeLock')
    @patch('bundlewrap.node.apply_items')
    def test_apply(self, apply_items, NodeLock, ApplyResult):
        repo = Repository()
        n = Node("node1", {})
        repo.add_node(n)
        result = MagicMock()
        ApplyResult.return_value = result
        NodeLock.__enter__ = lambda x: x
        NodeLock.__exit__ = lambda x: x
        self.assertEqual(n.apply(), result)
        self.assertEqual(apply_items.call_count, 1)
        assert ApplyResult.call_count == 1

    def test_bundles(self):
        repo = Repository()
        repo.bundle_names = ("bundle1", "bundle2", "bundle3")
        n = Node("node1", {})
        repo.add_node(n)
        g1 = Group("group1", {'bundles': ("bundle1", "bundle2")})
        repo.add_group(g1)
        g2 = Group("group2", {'bundles': ("bundle3",)})
        repo.add_group(g2)
        with patch('tests.unit.node_tests.Node.groups', new=(g1, g2)):
            self.assertEqual(
                tuple(names(n.bundles)),
                ("bundle1", "bundle2", "bundle3"),
            )

    def test_hostname_defaults(self):
        n = Node("node1", {})
        self.assertEqual(n.hostname, "node1")
        n = Node("node2", {'hostname': "node2.example.com"})
        self.assertEqual(n.hostname, "node2.example.com")

class NodeLockTest(TestCase):
    """
    Tests bundlewrap.node.NodeLock.
    """
    @patch('bundlewrap.node.ask_interactively')
    def test_locked(self, ask_interactively):
        node = MagicMock()
        runres = RunResult()
        runres.return_code = 1
        node.run.return_value = runres
        with self.assertRaises(NodeAlreadyLockedException):
            with NodeLock(node, False, ignore=False):
                pass