"""
Binance Pybot v1.1 (21-9-5)
https://github.com/rulibar/binance-pybot
"""

import os
import time
from calendar import timegm as timegm
import numpy
import random
import logging
import talib
from binance.client import Client

# instance vars
api_key = ""
api_secret = ""
client = Client(api_key, api_secret, tld='us')

asset = "ETH"; base = "BTC"
interval_mins = 30

# strategy vars
storage = dict()

# set up logger
def set_log_file():
    # Set up the log folders
    gmt = time.gmtime()
    yy = str(gmt.tm_year)[2:]; mm = str(gmt.tm_mon); dd = str(gmt.tm_mday)
    if len(mm) == 1: mm = "0" + mm
    if len(dd) == 1: dd = "0" + dd
    path = "./logs/"
    if not os.path.isdir(path): os.mkdir(path)
    path += "{}/".format(yy + mm)
    if not os.path.isdir(path): os.mkdir(path)
    # Set the log destination and format
    fileh = logging.FileHandler("./logs/{}/{}.log".format(yy + mm, yy + mm + dd), "a")
    formatter = logging.Formatter("%(levelname)s %(asctime)s - %(message)s")
    fileh.setFormatter(formatter)
    logger.handlers = [fileh]

logging.basicConfig(level=logging.INFO)
logging.Formatter.converter = time.gmtime
logger = logging.getLogger()
set_log_file()

# set up trading bot
def fix_dec(float_in):
    float_out = "{:.8f}".format(float_in)
    while float_out[-1] == "0": float_out = float_out[:-1]
    if float_out[-1] == ".": float_out = float_out[:-1]
    return float_out

def shrink_list(list_in, size):
    if len(list_in) > size: return list_in[-size:]
    return list_in

class Portfolio:
    def __init__(self, candle, positions, funds):
        self.ts = candle['ts_end']
        self.asset = positions['asset'][1]
        self.base = positions['base'][1]
        self.price = candle['close']
        self.positionValue = self.price * self.asset
        self.size = self.base + self.positionValue
        self.funds = funds
        if funds > self.size or funds == 0: self.funds = float(self.size)
        self.sizeT = float(self.funds)
        self.rin = self.price * self.asset / self.size
        self.rinT = self.price * self.asset / self.sizeT

class Instance:
    def __init__(self, asset, base, interval_mins):
        self.next_log = 0
        self.ticks = 0; self.days = 0; self.trades = 0
        self.exchange = "binance"
        self.base = str(base)
        self.asset = str(asset)
        self.pair = self.asset + self.base
        self.interval = int(interval_mins)
        logger.info("New trader instance started on {} {} {}m.".format(self.exchange.title(), self.pair, self.interval))
        self.get_params()

        self.candles_raw = self._get_candles_raw()
        self.candles = self._get_candles()
        self.candles_raw = shrink_list(self.candles_raw, 2 * self.interval)
        self.candles_raw_unused = self._get_raw_unused()

        self.deposits_pending = set()
        self.withdrawals_pending = set()
        self.earliest_pending = 0

        self.candle_start = None
        self.positions_start = None
        self.positions_init_ts = 0
        self.positions = self.get_positions()
        self.positions_f = {'asset': list(self.positions['asset'])}
        self.positions_f['base'] = list(self.positions['base'])
        self.positions_t = {'asset': list(self.positions['asset'])}
        self.positions_t['base'] = list(self.positions['base'])
        p = Portfolio(self.candles[-1], self.positions, float(self.params['funds']))
        self.last_order = {"type": "none", "amt": 0, "pt": self.candles[-1]['close']}
        self.signal = {"rinTarget": p.rinT, "rinTargetLast": p.rinT, "position": "none", "status": 0, "apc": p.price, "target": p.price, "stop": p.price}
        self.performance = {"bh": 0, "change": 0, "W": 0, "L": 0, "wSum": 0, "lSum": 0, "w": 0, "l": 0, "be": 0, "aProfits": 0, "bProfits": 0, "cProfits": 0}
        self.init(p)

    def _get_candles_raw(self):
        # get enough 1m candles to create 600 historical candles
        data = self.get_historical_candles(self.pair, "1m", 600 * self.interval)
        data.pop()
        for i in range(len(data)): data[i] = self.get_candle(data[i])
        return data

    def _get_candles(self):
        # convert historical 1m candles into historical candles
        candles = list(); candle_new = dict()
        candles_raw_clone = list(self.candles_raw)
        for i in range(self.interval - 2): candles_raw_clone.pop()
        for i in range(len(candles_raw_clone)):
            order = i % self.interval
            candle_raw = candles_raw_clone[- 1 - i]

            if order == 0:
                candle_new = candle_raw
                continue

            if candle_raw["high"] > candle_new["high"]:
                candle_new["high"] = candle_raw["high"]
            if candle_raw["low"] < candle_new["low"]:
                candle_new["low"] = candle_raw["low"]
            candle_new["volume"] += candle_raw["volume"]

            if order == self.interval - 1:
                candle_new["open"] = candle_raw["open"]
                candle_new["ts_start"] = candle_raw["ts_start"]
                candles.append(candle_new)

        return candles[::-1]

    def _get_raw_unused(self):
        # get unused historical 1m candles
        raw_unused = -1
        str_out = str()
        data = self.candles_raw[-2 * self.interval:]
        for i in range(len(data)):
            candle_raw = data[i]
            if raw_unused > -1:
                raw_unused += 1
            if candle_raw["ts_end"] == self.candles[-1]["ts_end"]:
                raw_unused += 1
                continue

            if raw_unused > 0: str_out += "    {}\n".format(candle_raw)

        return raw_unused

    def get_historical_candles_method(self, symbol, interval, start_str):
        data, err = list(), str()
        try: data = client.get_historical_klines(symbol, interval, start_str)
        except Exception as e: err = e
        return data, err

    def get_historical_candles(self, symbol, interval, n_candles):
        tries = 0
        while True:
            data, err = self.get_historical_candles_method(symbol, interval, "{} minutes ago UTC".format(n_candles))
            tries += 1
            if len(data) == 0:
                if tries <= 3:
                    err_msg = "Error getting historical candle data. Retrying in 5 seconds..."
                    if err != "": err_msg += "\n'{}'".format(err)
                    logger.error(err_msg)
                if tries == 3: logger.error("(Future repeats of this error hidden to avoid spam.)")
                time.sleep(5)
            else: break
        if tries > 3: logger.error("Failed to get historical candle data {} times.".format(tries - 1))
        return data

    def get_candle(self, data):
        # data is a kline list from Binance
        candle = {
            "ts_start": int(data[0]),
            "open": round(float(data[1]), 8),
            "high": round(float(data[2]), 8),
            "low": round(float(data[3]), 8),
            "close": round(float(data[4]), 8),
            "volume": round(float(data[5]), 8),
            "ts_end": int(data[6])}
        return candle

    def limit_buy(self, amt, pt):
        try:
            logger.warning("Trying to buy {} {} for {} {}. (price: {})".format(fix_dec(amt), self.asset, fix_dec(round(amt * pt, self.pt_dec)), self.base, fix_dec(pt)))
            self.last_order = {"type": "buy", "amt": amt, "pt": pt}
            client.order_limit_buy(symbol = self.pair, quantity = "{:.8f}".format(amt), price = "{:.8f}".format(pt))
        except Exception as e:
            logger.error("Error buying.\n'{}'".format(e))

    def limit_sell(self, amt, pt):
        try:
            logger.warning("Trying to sell {} {} for {} {}. (price: {})".format(fix_dec(amt), self.asset, fix_dec(round(amt * pt, self.pt_dec)), self.base, fix_dec(pt)))
            self.last_order = {"type": "sell", "amt": amt, "pt": pt}
            client.order_limit_sell(symbol = self.pair, quantity = "{:.8f}".format(amt), price = "{:.8f}".format(pt))
        except Exception as e:
            logger.error("Error selling.\n'{}'".format(e))

    def bso(self, p):
        # buy/sell/other
        s = self.signal

        rbuy = s['rinTarget'] - s['rinTargetLast']
        order_size = 0
        if rbuy * p.asset >= 0:
            order_size = abs(rbuy * p.funds)
            if order_size > p.base: order_size = p.base
        if rbuy * p.asset < 0:
            rbuy_asset = rbuy / s['rinTargetLast']
            order_size = abs(rbuy_asset * p.asset * p.price)
        if order_size < self.min_order: order_size = 0

        if order_size > 0:
            if rbuy > 0: pt = (1 + 0.0015) * p.price
            else: pt = (1 - 0.0015) * p.price
            pt = round(pt, self.pt_dec)
            if rbuy > 0: amt = order_size / pt
            else: amt = order_size / p.price
            amt = round(0.995 * amt * 10**self.amt_dec - 2) / 10**self.amt_dec
            if rbuy > 0: self.limit_buy(amt, pt)
            if rbuy < 0: self.limit_sell(amt, pt)
        if rbuy == 0: order_size = 0
        if order_size == 0:
            if self.ticks == 1: logger.info("Waiting for a signal to trade...")
            self.last_order = {"type": "none", "amt": 0, "pt": p.price}

    def close_orders(self):
        # close open orders
        try:
            orders = client.get_open_orders(symbol = self.pair)
            for order in orders:
                client.cancel_order(symbol = self.pair, orderId = order['orderId'])
        except Exception as e:
            logger.error("Error closing open orders.\n'{}'".format(e))

    def update_vars(self):
        # Get preliminary vars
        self.ticks += 1
        self.days = (self.ticks - 1) * self.interval / (60 * 24)

        try: data = client.get_symbol_info(self.pair)['filters']
        except Exception as e:
            logger.error("Error getting symbol info.\n'{}'".format(e))
            return
        min_order = float(data[2]['minQty']) * self.candles[-1]['close']
        self.min_order = 3 * max(min_order, float(data[3]['minNotional']))
        amt_dec = 8
        for char in reversed(data[2]['stepSize']):
            if char == "0": amt_dec -= 1
            else: break
        self.amt_dec = amt_dec
        pt_dec = 8
        for char in reversed(data[0]['tickSize']):
            if char == "0": pt_dec -= 1
            else: break
        self.pt_dec = pt_dec

    def get_params(self):
        # import and process params from config.txt
        params = dict()
        with open("config.txt") as cfg:
            par = [l.split()[0] for l in cfg.read().split("\n")[2:-1]]
            for p in par:
                p = p.split("=")
                if len(p) != 2: continue
                params[str(p[0])] = str(p[1])

        # check values
        funds = float(params['funds'])
        if funds < 0:
            logger.warning("Warning! Maximum amount to invest should be zero or greater.")
            params['funds'] = "0"

        logs_per_day = float(params['logs_per_day'])
        if logs_per_day < 0:
            logger.warning("Warning! Logs per day should be zero or greater.")
            params['logs_per_day'] = "1"

        log_dws = str(params['log_dws'])
        if log_dws not in {"yes", "no"}:
            logger.warning("Warning! Log deposits and withdrawals set to 'yes'.")
            params['log_dws'] = "yes"

        # check for additions and removals
        if self.ticks == 0: self.params = dict()

        keys_old = {key for key in self.params}
        keys_new = {key for key in params}

        keys_added = {key for key in keys_new if key not in keys_old}
        keys_removed = {key for key in keys_old if key not in keys_new}

        if len(keys_added) > 0:
            logger.info("{} parameter(s) added.".format(len(keys_added)))
            for key in keys_added: logger.info("    \"{}\": {}".format(key, params[key]))
        if len(keys_removed) > 0:
            logger.info("{} parameter(s) removed.".format(len(keys_removed)))
            for key in keys_removed: logger.info("    \"{}\"".format(key))

        # check for changes
        keys_remaining = {key for key in keys_old if key in keys_new}
        keys_changed = set()

        for key in keys_remaining:
            if params[key] != self.params[key]: keys_changed.add(key)

        if self.ticks == 0:
            keys_changed.add('funds'); keys_changed.add('logs_per_day'); keys_changed.add('log_dws')

        if "funds" in keys_changed:
            if params['funds'] == "0": logger.info("No maximum investment amount specified.")
            else: logger.info("Maximum investment amount set to {} {}.".format(params['funds'], self.base))
            self.params['funds'] = params['funds']
            keys_changed.remove('funds')
        if "logs_per_day" in keys_changed:
            if params['logs_per_day'] == "0": logger.info("Log updates turned off.")
            elif params['logs_per_day'] == "1": logger.info("Logs updating once per day.")
            else: logger.info("Logs updating {} times per day".format(params['logs_per_day']))
            self.params['logs_per_day'] = params['logs_per_day']
            keys_changed.remove('logs_per_day')
        if "log_dws" in keys_changed:
            if params['log_dws'] == "yes": logger.info("Deposit and withdrawal logs enabled.")
            else: logger.info("Deposit and withdrawal logs disabled.")
            self.params['log_dws'] = params['log_dws']
            keys_changed.remove('log_dws')

        if len(keys_changed) > 0:
            logger.info("{} parameter(s) changed.".format(len(keys_changed)))
            for key in keys_changed:
                logger.info("    \"{}\": {} -> {}".format(key, self.params[key], params[key]))
                self.params[key] = params[key]

    def get_new_candle(self):
        # Create a new candle from 1m candles
        candle_new = dict()
        for i in range(self.interval):
            candle_raw = self.candles_raw[- 1 - i]

            if i == 0:
                candle_new = candle_raw
                continue

            if candle_raw["high"] > candle_new["high"]:
                candle_new["high"] = candle_raw["high"]
            if candle_raw["low"] < candle_new["low"]:
                candle_new["low"] = candle_raw["low"]
            candle_new["volume"] += candle_raw["volume"]

            if i == self.interval - 1:
                candle_new["open"] = candle_raw["open"]
                candle_new["ts_start"] = candle_raw["ts_start"]
                self.candles.append(candle_new)
                self.candles_raw_unused = 0
                self.candles = shrink_list(self.candles, 5000)

    def get_positions(self):
        positions = {"asset": [self.asset, 0], "base": [self.base, 0]}
        try: data = client.get_account()["balances"]
        except Exception as e:
            logger.error("Error getting account balances.\n'{}'".format(e))
            return self.positions
        for i in range(len(data)):
            asset = data[i]["asset"]
            if asset not in {self.asset, self.base}: continue
            free = float(data[i]["free"])
            locked = float(data[i]["locked"])
            total = free + locked
            if asset == self.asset: positions['asset'][1] = total
            if asset == self.base: positions['base'][1] = total

        if self.ticks == 0:
            self.positions_init_ts = int(1000 * time.time())
            self.earliest_pending = int(self.positions_init_ts)

        return positions

    def update_f(self, p, apc):
        if apc == 0:
            if self.ticks != 1: return
            apc = p.price
        r = self.performance
        s = self.signal
        pos_f = self.positions_f
        pos_t = self.positions_t

        size = p.base + apc * p.asset
        rin = apc * p.asset / size
        sizeT = p.funds * (1 - s['rinTargetLast']) + apc * p.asset
        rinT = apc * p.asset / sizeT

        if self.ticks == 1: size_f = 1; size_t = 1
        else:
            size_f = pos_f['base'][1] + apc * pos_f['asset'][1]
            size_t = pos_t['base'][1] + apc * pos_t['asset'][1]
            if s['rinTarget'] == 0 and p.positionValue < self.min_order:
                profit = size_t - 1
                if profit >= 0: r['wSum'] += profit; r['W'] += 1; self.trades += 1
                if profit < 0: r['lSum'] += profit; r['L'] += 1; self.trades += 1
                if r['W'] != 0: r['w'] = r['wSum'] / r['W']
                if r['L'] != 0: r['l'] = r['lSum'] / r['L']
                size_t = 1

        base_f = (1 - rin) * size_f; base_t = (1 - rinT) * size_t
        asset_f = (rin / apc) * size_f; asset_t = (rinT / apc) * size_t

        pos_f['base'][1] = base_f; pos_t['base'][1] = base_t
        pos_f['asset'][1] = asset_f; pos_t['asset'][1] = asset_t

    def get_dws(self):
        diffasset_dw = 0; diffbase_dw = 0
        ts_last = self.candles[-2]['ts_end']
        if len(self.deposits_pending) == 0 and len(self.withdrawals_pending) == 0: self.earliest_pending = ts_last
        start_time = self.earliest_pending - 1000
        if self.ticks == 1: start_time = self.earliest_pending - 1000 * 60 * 60 * 24 * 7

        def process_d(deposit):
            amt = deposit['amount']
            diffasset = 0; diffbase = 0
            if self.params['log_dws'] == "yes":
                logger.warning("Deposit of {} {} detected.".format(fix_dec(amt), deposit['asset']))
            if deposit['asset'] == self.base: diffbase += amt
            else: diffasset += amt
            return diffasset, diffbase

        def process_w(withdrawal):
            amt = withdrawal['amount'] + withdrawal['transactionFee']
            diffasset = 0; diffbase = 0
            if self.params['log_dws'] == "yes":
                logger.warning("Withdrawal of {} {} detected.".format(fix_dec(amt), withdrawal['asset']))
            if withdrawal['asset'] == self.base: diffbase -= amt
            else: diffasset -= amt
            return diffasset, diffbase

        # get dws
        try:
            deposits = client.get_deposit_history(startTime = start_time)
            withdrawals = client.get_withdraw_history(startTime = start_time)

            # deal with differing formats
            if type(deposits) is dict:
                deposits = deposits['depositList']

            if type(withdrawals) is dict:
                withdrawals = withdrawals['withdrawList']

            for deposit in deposits:
                if 'asset' not in deposit:
                    deposit['asset'] = str(deposit['coin'])
                    deposit['creator'] = ''
                    deposit['amount'] = float(deposit['amount'])

            for withdrawal in withdrawals:
                if 'asset' not in withdrawal:
                    withdrawal['asset'] = str(withdrawal['coin'])
                    withdrawal['withdrawOrderId'] = ''
                    withdrawal['amount'] = float(withdrawal['amount'])
                    withdrawal['transactionFee'] = float(withdrawal['transactionFee'])

                if type(withdrawal['applyTime']) is str:
                    applyTime = time.strptime(withdrawal['applyTime'], '%Y-%m-%d %H:%M:%S')
                    applyTime = 1000 * timegm(applyTime)
                    withdrawal['applyTime'] = int(applyTime)
        except Exception as e:
            logger.error("Error getting deposits and withdrawals.\n'{}'".format(e))
            return 0, 0
        deposits = [d for d in deposits if d['asset'] in {self.asset, self.base}]
        withdrawals = [w for w in withdrawals if w['asset'] in {self.asset, self.base}]

        # add new pending trades to pending set
        # check for and process old pending trades that were filled
        for deposit in deposits:
            id = deposit['txId']
            if deposit['status'] < 1:
                self.deposits_pending.add(id)
                self.earliest_pending = start_time
                continue
            if id not in self.deposits_pending: continue
            if deposit['status'] > 0:
                diffasset, diffbase = process_d(deposit)
                diffasset_dw += diffasset
                diffbase_dw += diffbase
                self.deposits_pending.remove(id)
        for withdrawal in withdrawals:
            id = withdrawal['id']
            if withdrawal['status'] < 0:
                self.withdrawals_pending.add(id)
                self.earliest_pending = start_time
                continue
            if id not in self.withdrawals_pending: continue
            if withdrawal['status'] > -1:
                diffasset, diffbase = process_w(withdrawal)
                diffasset_dw += diffasset
                diffbase_dw += diffbase
                self.withdrawals_pending.remove(id)

        # skip the last section if it's the first tick
        if self.ticks == 1: return diffasset_dw, diffbase_dw

        # process dws received and completed in the last candle
        deposits = [d for d in deposits if d['insertTime'] > ts_last]
        deposits = [d for d in deposits if d['status'] > 0]
        withdrawals = [w for w in withdrawals if w['applyTime'] > ts_last]
        withdrawals = [w for w in withdrawals if w['status'] > -1]

        for deposit in deposits:
            id = deposit['txId']
            diffasset, diffbase = process_d(deposit)
            diffasset_dw += diffasset
            diffbase_dw += diffbase
        for withdrawal in withdrawals:
            id = withdrawal['id']
            diffasset, diffbase = process_w(withdrawal)
            diffasset_dw += diffasset
            diffbase_dw += diffbase

        return diffasset_dw, diffbase_dw

    def get_trades(self, p):
        diffasset_trad = 0; diffbase_trad = 0
        ts_last = self.candles[-2]['ts_end']
        l = self.last_order
        s = self.signal
        limit = int(ts_last)
        if self.ticks == 1: limit = self.positions_init_ts

        # Get trades
        try: trades = reversed(client.get_my_trades(symbol = self.pair, limit = 20))
        except Exception as e:
            logger.error("Error getting trade info.\n'{}'".format(e))
            return 0, 0, p.price
        trades = [t for t in trades if t['time'] > limit]

        # process trades
        if len(trades) > 0:
            str_out = "{} new trade(s) found.".format(len(trades))
            for trade in trades:
                str_out += "\n    {}".format(trade)
                qty = float(trade['qty'])
                price = float(trade['price'])
                if not trade['isBuyer']: qty *= -1
                diffasset_trad += qty
                diffbase_trad -= qty * price

        rbuy = s['rinTarget'] - s['rinTargetLast']
        rTrade = 0
        apc = 0
        if diffasset_trad != 0: apc = -diffbase_trad / diffasset_trad
        if l['amt'] != 0: rTrade = abs(diffasset_trad / l['amt'])
        if diffasset_trad > 0:
            log_amt = "{} {}".format(fix_dec(diffasset_trad), self.asset)
            log_size = "{} {}".format(fix_dec(diffasset_trad * apc), self.base)
            if l['type'] != "buy":
                logger.info("Manual buy detected.")
                rTrade = 0
            elif abs(rTrade - 1) > 0.1:
                logger.info("Buy order partially filled.")
            logger.warning("{} bought for {}.".format(log_amt, log_size))
        elif diffasset_trad < 0:
            log_amt = "{} {}".format(fix_dec(-diffasset_trad), self.asset)
            log_size = "{} {}".format(fix_dec(-diffasset_trad * apc), self.base)
            if l['type'] != "sell":
                logger.info("Manual sell detected")
                rTrade = 0
            elif abs(rTrade - 1) > 0.1:
                logger.info("Sell order partially filled.")
            logger.warning("{} sold for {}.".format(log_amt, log_size))

        if self.ticks == 1 or diffasset_trad != 0:
            s['rinTargetLast'] += rTrade * rbuy
            self.update_f(p, apc)

        return diffasset_trad, diffbase_trad, apc

    def get_dwts(self, p):
        # get deposits, withdrawals, and trades
        s = self.signal
        diffasset = round(self.positions['asset'][1] - self.positions_last['asset'][1], 8)
        diffbase = round(self.positions['base'][1] - self.positions_last['base'][1], 8)

        # get dws and trades
        diffasset_dw, diffbase_dw = self.get_dws()
        diffasset_trad, diffbase_trad, apc = self.get_trades(p)

        # get unknown changes
        diffasset_expt = round(diffasset_dw + diffasset_trad, 8)
        diffbase_expt = round(diffbase_dw + diffbase_trad, 8)
        diffasset_unkn = diffasset - diffasset_expt
        diffbase_unkn = diffbase - diffbase_expt

        # process unknown changes
        if self.params['log_dws'] == "yes":
            if diffasset_unkn > 0: logger.info("{} {} has become available.".format(fix_dec(diffasset_unkn), self.asset))
            elif diffasset_unkn < 0: logger.info("{} {} has become unavailable.".format(fix_dec(-diffasset_unkn), self.asset))
            if diffbase_unkn > 0: logger.info("{} {} has become available.".format(fix_dec(diffbase_unkn), self.base))
            elif diffbase_unkn < 0: logger.info("{} {} has become unavailable.".format(fix_dec(-diffbase_unkn), self.base))

        # set position and apc
        if apc == 0: apc = p.price
        if p.positionValue > self.min_order:
            if s['position'] != "long":
                s['position'] = "long"; s['apc'] = apc
        elif s['position'] != "none":
            s['position'] = "none"; s['apc'] = apc

        return

    def get_performance(self, p):
        if self.ticks == 1:
            self.candle_start = dict(self.candles[-1])
            self.positions_start = {key: value[:] for key, value in self.positions.items()}
        r = self.performance
        s = self.signal
        pos_f = self.positions_f
        pos_t = self.positions_t
        c_start = self.candle_start
        p_start = self.positions_start
        p_start_size = p_start['base'][1] + c_start['close'] * p_start['asset'][1]

        pfake_size = pos_f['base'][1] + p.price * pos_f['asset'][1]
        ptrade_size = pos_t['base'][1] + p.price * pos_t['asset'][1]

        r['bh'] = (p.price - c_start['close']) / c_start['close']
        r['change'] = (p.size - p_start_size) / p_start_size
        r['bProfits'] = pfake_size - 1
        r['aProfits'] = (1 + r['bProfits']) / (1 + r['bh']) - 1
        r['cProfits'] = ptrade_size - 1

        W = int(r['W']); w = float(r['w'])
        L = int(r['L']); l = float(r['l'])
        if r['cProfits'] >= 0: W += 1; wSum = r['wSum'] + r['cProfits']; w = wSum / W
        if r['cProfits'] < 0: L += 1; lSum = r['lSum'] + r['cProfits']; l = lSum / L
        r['be'] = W * w + L * l

    def log_update(self, p):
        r = self.performance

        hr = "~~~~~"
        tpd = float()
        if self.days != 0: tpd = self.trades / self.days
        winrate = float()
        if r['W'] + r['L'] != 0: winrate = r['W'] / (r['W'] + r['L'])
        header_1 = "{} {} {} {}".format(3 * hr, self.bot_name, self.version, 9 * hr)[:50]
        header_2 = "{} {} {} {}m {}".format(3 * hr, self.exchange.title(), self.pair, self.interval, 9 * hr)[:50]
        trades = "{} trades ({} per day)".format(int(self.trades), round(tpd, 2))
        currency = "{} {}".format(fix_dec(p.base), self.base)
        price = "{} {}".format(fix_dec(p.price), self.base)
        assets = "{} {}".format(fix_dec(p.asset), self.asset)
        assetvalue = "{} {}".format(fix_dec(p.positionValue), self.base)
        accountvalue = "{} {}".format(fix_dec(p.size), self.base)
        boteff = "{}% {},".format(round(100 * r['be'], 2), self.base)
        boteff += " {}% {}".format(round(100 * ((1 + r['be']) / (1 + r['bh'])) - 100, 2), self.asset)
        botprof = "{}% {},".format(round(100 * r['bProfits'], 2), self.base)
        botprof += " {}% {}".format(round(100 * ((1 + r['bProfits']) / (1 + r['bh'])) - 100, 2), self.asset)

        logger.info(header_1)
        logger.info(header_2)
        logger.info("Days running: {} | Trades completed: {}".format(round(self.days, 2), trades))
        logger.info("Currency: {} | Current price: {}".format(currency, price))
        logger.info("Assets: {} | Value of assets: {}".format(assets, assetvalue))
        logger.info("Value of account: {}".format(accountvalue))
        logger.info("    Win rate: {}%".format(round(100 * winrate, 2)))
        logger.info("    Wins: {} | Average win: {}%".format(r['W'], round(100 * r['w'], 2)))
        logger.info("    Losses: {} | Average loss: {}%".format(r['L'], round(100 * r['l'], 2)))
        logger.info("    Current profits: {}%".format(round(100 * r['cProfits'], 2)))
        logger.info("    Bot efficiency: {}".format(boteff))
        logger.info("Bot profits: {}".format(botprof))
        logger.info("Buy and hold: {}%".format(round(100 * r['bh'], 2)))

    def init(self, p):
        # Binance Pybot 20/100 SXS
        self.bot_name = "Binance Pybot"
        self.version = "1.1"
        logger.info("Analyzing the market...")

        # vars
        ma_lens = [[20, 100]]
        storage["ma_lens"] = ma_lens
        logger.info("Ready to start trading...")

    def strat(self, p):
        """ strategy / trading algorithm
        - Use talib for indicators
        - Talib objects require numpy.array objects as input
        - s stands for signal, rinTarget stands for 'ratio invested target'
        - Set s['rinTarget'] between 0 and 1. 0 is 0%, 1 is 100% invested
        """
        # vars
        s = self.signal
        ma_lens = storage["ma_lens"]
        num_sigs = len(ma_lens)
        close_data = numpy.array([c['close'] for c in self.candles])

        # get sigs
        SMAs = dict()
        class SMA:
            def __init__(self, ma_len):
                self.ma_len = ma_len
                self.arr = talib.SMA(close_data, timeperiod = ma_len)
                self.price = self.arr[-1]
                SMAs[ma_len] = self.price

        class SXS:
            def __init__(self, ma_len_pair):
                self.pair = ma_len_pair
                self.trend = "bear"
                if ma_len_pair[0] not in SMAs:
                    mas = SMA(ma_len_pair[0])
                self.mas = SMAs[ma_len_pair[0]]
                if ma_len_pair[1] not in SMAs:
                    mal = SMA(ma_len_pair[1])
                self.mal = SMAs[ma_len_pair[1]]
                if self.mas > self.mal: self.trend = "bull"

        sigs = list()
        for i in range(num_sigs):
            sig = SXS(ma_lens[i])
            sigs.append(sig)
        num_sigs = len(sigs)

        # set rinTarget
        nBulls = 0; nBears = 0
        s['rinTarget'] = 0
        for i in range(num_sigs):
            if sigs[i].trend == "bear": nBears += 1
            if sigs[i].trend == "bull":
                nBulls += 1
                s['rinTarget'] += 1 / num_sigs

    def ping(self):
        # check if its time for a new candle
        if (1000 * time.time() - self.candles_raw[-1]["ts_end"]) < 60000: return
        data = self.get_historical_candles(self.pair, "1m", 2)
        data_top = self.get_candle(data[0])

        # New raw candle?
        if data_top["ts_end"] != self.candles_raw[-1]["ts_end"]:
            self.candles_raw_unused += 1
            self.candles_raw.append(data_top)
            self.candles_raw = self.candles_raw[-2 * self.interval:]

        # New candle?
        if self.candles_raw_unused == self.interval:
            # Preliminary setup
            set_log_file()
            self.close_orders()
            self.update_vars()
            self.get_params()
            self.get_new_candle()

            self.positions_last = {key: value[:] for key, value in self.positions.items()}
            self.positions = self.get_positions()
            p = Portfolio(self.candles[-1], self.positions, float(self.params['funds']))
            self.get_dwts(p)
            self.get_performance(p)

            # Log output
            if self.params['logs_per_day'] == "0": self.next_log = self.days + 1
            if self.days >= self.next_log:
                self.log_update(p)
                self.next_log += 1 / float(self.params['logs_per_day'])

            # Trading strategy, buy/sell/other
            self.strat(p)
            self.bso(p)

ins = Instance(asset, base, interval_mins)
while True:
    ins.ping()
    time.sleep(0.5)


import random
import numpy as np

import torch
import torch.optim as optim
import torch.nn.functional as F
from torch.autograd import Variable

from model import DQN, DuelingDQN
from replay_buffer import ReplayBuffer, PrioritizedReplayBuffer
from rl_config import RLConfig

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

class DQNAgent:
    def __init__(self,
            state_size,
            action_size,
            config=RLConfig()):
        self.seed = random.seed(config.seed)
        self.state_size = state_size
        self.action_size = action_size
        self.batch_size = config.batch_size
        self.batch_indices = torch.arange(config.batch_size).long().to(device)
        self.samples_before_learning = config.samples_before_learning
        self.learn_interval = config.learning_interval
        self.parameter_update_interval = config.parameter_update_interval
        self.per_epsilon = config.per_epsilon
        self.tau = config.tau
        self.gamma = config.gamma

        if config.useDuelingDQN:
            self.qnetwork_local = DuelingDQN(state_size, action_size, config.seed).to(device)
            self.qnetwork_target = DuelingDQN(state_size, action_size, config.seed).to(device)
        else:
            self.qnetwork_local = DQN(state_size, action_size, config.seed).to(device)
            self.qnetwork_target = DQN(state_size, action_size, config.seed).to(device)
        self.optimizer = optim.Adam(self.qnetwork_local.parameters(), lr=config.learning_rate)

        self.doubleDQN = config.useDoubleDQN
        self.usePER = config.usePER
        if self.usePER:
            self.memory = PrioritizedReplayBuffer(config.buffer_size,config.per_alpha)
        else:
            self.memory = ReplayBuffer(config.buffer_size)

        self.t_step = 0

    def act(self, state, eps=0.):
        state = torch.from_numpy(state).float().unsqueeze(0).to(device)
        self.qnetwork_local.eval()
        with torch.no_grad():
            action_values = self.qnetwork_local(state)
        self.qnetwork_local.train()

        if random.random() < eps:
            return random.choice(np.arange(self.action_size))
        else:
            return np.argmax(action_values.cpu().data.numpy())

    def step(self, state, action, reward, next_state, done, beta):
        self.memory.add(state, action, reward, next_state, done)
        self.t_step += 1
        if self.t_step % self.learn_interval == 0:
            if len(self.memory) > self.samples_before_learning:
                state = torch.from_numpy(state).float().unsqueeze(0).to(device)
                next_state = torch.from_numpy(next_state).float().unsqueeze(0).to(device)
                target = self.qnetwork_local(state).data
                old_val = target[0][action]
                target_val = self.qnetwork_target(next_state).data
                if done:
                    target[0][action] = reward
                else:
                    target[0][action] = reward + self.gamma * torch.max(target_val)
                if self.usePER:
                    states, actions, rewards, next_states, dones, weights, indices = self.memory.sample(self.batch_size, beta)
                else:
                    indices=None
                    weights=None
                    states, actions, rewards, next_states, dones = self.memory.sample(self.batch_size)

                self.learn(states, actions, rewards, next_states, dones, indices, weights, self.gamma)

    def learn(self, states, actions, rewards, next_states, dones, indices, weights, gamma):
        states = torch.from_numpy(np.vstack(states)).float().to(device)
        actions = torch.from_numpy(np.vstack(actions)).long().to(device)
        rewards = torch.from_numpy(np.vstack(rewards)).float().to(device)
        next_states = torch.from_numpy(np.vstack(next_states)).float().to(device)
        dones = torch.from_numpy(np.vstack(dones.astype(np.uint8))).float().to(device)
        Q_targets_next = self.qnetwork_target(next_states).detach()

        if self.doubleDQN:
            # choose the best action from the local network
            next_actions = self.qnetwork_local(next_states).argmax(dim=-1)
            Q_targets_next = Q_targets_next[self.batch_indices, next_actions]
        else:
            Q_targets_next = Q_targets_next.max(1)[0]

        Q_targets = rewards + gamma * Q_targets_next.reshape((self.batch_size, 1)) * (1 - dones)

        pred = self.qnetwork_local(states)
        Q_expected = pred.gather(1, actions)

        if self.usePER:
            errors = torch.abs(Q_expected - Q_targets).data.numpy() + self.per_epsilon
            self.memory.update_priorities(indices, errors)

        self.optimizer.zero_grad()
        loss = F.mse_loss(Q_expected, Q_targets)
        loss.backward()
        self.optimizer.step()

        if self.t_step % self.parameter_update_interval == 0:
            self.soft_update(self.qnetwork_local, self.qnetwork_target, self.tau)

    def soft_update(self, qnetwork_local, qnetwork_target, tau):
        for local_param, target_param in zip(qnetwork_local.parameters(), qnetwork_target.parameters()):
            target_param.data.copy_(tau * local_param.data + (1.0 - tau) * target_param.data)


from __future__ import print_function
import numpy as np
import fitsio
from legacypipe.image import LegacySurveyImage
from legacypipe.bits import DQ_BITS
from legacypipe.survey import create_temp

class HscImage(LegacySurveyImage):
    def __init__(self, survey, ccd):
        ccd.plver = 'xxx'
        ccd.procdate = 'xxx'
        ccd.plprocid = 'xxx'
        super().__init__(survey, ccd)
        self.dq_hdu = 2
        self.wt_hdu = 3
        # Adjust zeropoint for exposure time
        self.ccdzpt += 2.5 * np.log10(self.exptime)

        # FIXME -- these are just from DECam
        self.zp0 = dict(
            g = 26.610,
            r = 26.818,
            i = 26.758,
            z = 26.484,
        )
        self.k_ext = dict(g = 0.17,
                          r = 0.10,
                          i = 0.08,
                          z = 0.06,
                          )

    def get_extension_list(self, fn, debug=False):
        return [self.image_hdu,]

    def calibration_good(self, primhdr):
        return True

    '''
    def get_psfex_unmerged_filename(self):
        basefn = os.path.basename(self.fn_base)
        basedir = os.path.dirname(self.fn_base)
        base = basefn.split('.')[0]
        fn = base + '-psfex.fits'
        fn = os.path.join(self.calibdir, 'psfex-single', basedir, base, fn)
        return fn
    def get_splinesky_unmerged_filename(self):
        basefn = os.path.basename(self.fn_base)
        basedir = os.path.dirname(self.fn_base)
        base = basefn.split('.')[0]
        fn = base + '-splinesky.fits'
        fn = os.path.join(self.calibdir, 'sky-single', basedir, base, fn)
        return fn
    '''
    def compute_filenames(self):
        self.dqfn = self.imgfn
        self.wtfn = self.imgfn

    def get_expnum(self, primhdr):
        return primhdr['EXPID']

    def get_fwhm(self, hdr, hdu):
        return self.primhdr['SEEING']

    def get_wcs(self, hdr=None):
        from astrometry.util.util import Sip
        if hdr is None:
            hdr = self.read_image_header()
        wcs = Sip(hdr)
        # Correction: ccd ra,dec offsets from zeropoints/CCDs file
        dra,ddec = self.dradec
        # debug('Applying astrometric zeropoint:', (dra,ddec))
        r,d = wcs.get_crval()
        wcs.set_crval((r + dra / np.cos(np.deg2rad(d)), d + ddec))
        wcs.version = ''
        wcs.plver = ''
        #phdr = self.read_image_primary_header()
        #wcs.plver = phdr.get('PLVER', '').strip()
        return wcs

    def colorterm_ps1_to_observed(self, ps1stars, band):
        """ps1stars: ps1.median 2D array of median mag for each band"""
        from legacypipe.ps1cat import ps1_to_hsc
        return ps1_to_hsc(ps1stars, band)

    def read_image(self, header=False, **kwargs):
        img = super().read_image(header=header, **kwargs)
        if header:
            img,hdr = img
        img[np.logical_not(np.isfinite(img))] = 0.
        if header:
            img = img,hdr
        return img

    def remap_dq(self, dq, header):
        return remap_hsc_bitmask(dq, header)

    def read_invvar(self, dq=None, **kwargs):
        # HSC has a VARIANCE map (not a weight map)
        v = self._read_fits(self.wtfn, self.wt_hdu, **kwargs)
        iv = 1./v
        iv[v==0] = 0.
        iv[np.logical_not(np.isfinite(iv))] = 0.
        #! this can happen
        iv[np.logical_not(np.isfinite(np.sqrt(iv)))] = 0.
        return iv

    def funpack_files(self, imgfn, maskfn, imghdu, maskhdu, todelete):
        # Before passing files to SourceExtractor / PsfEx, filter our mask image
        # because it marks DETECTED pixels with a mask bit.
        tmpimgfn,tmpmaskfn = super().funpack_files(imgfn, maskfn, imghdu, maskhdu, todelete)
        #print('Dropping mask bit 5 before running SE')
        m = fitsio.read(tmpmaskfn)
        m &= ~(1 << 5)
        tmpmaskfn = create_temp(suffix='.fits')
        todelete.append(tmpmaskfn)
        fitsio.write(tmpmaskfn, m, clobber=True)
        return tmpimgfn, tmpmaskfn

    def validate_version(self, *args, **kwargs):
        return True
    def check_image_header(self, imghdr):
        pass

def remap_hsc_bitmask(dq, header):
    new_dq = np.zeros(dq.shape, np.int16)
    # MP_BAD  =                    0
    # MP_SUSPECT =        7
    # MP_NO_DATA =        8
    # MP_CROSSTALK =      9
    # MP_UNMASKEDNAN =   11
    new_dq |= (DQ_BITS['badpix'] * ((dq & ((1<<0) | (1<<7) | (1<<8) | (1<<9) | (1<<11))) != 0))
    # MP_SAT  =                    1
    new_dq |= (DQ_BITS['satur' ] * ((dq & (1<<1)) != 0))
    # MP_INTRP=                    2
    new_dq |= (DQ_BITS['interp'] * ((dq & (1<<2)) != 0))
    # MP_CR   =                    3
    new_dq |= (DQ_BITS['cr'] * ((dq & (1<<3)) != 0))
    #MP_EDGE =                    4
    new_dq |= (DQ_BITS['edge'] * ((dq & (1<<4)) != 0))
    '''
    MP_DETECTED =       5
    MP_DETECTED_NEGATIVE = 6
    MP_NOT_DEBLENDED = 10
    '''
    return new_dq


#coding=utf8
import numpy as np
np.random.seed(1337)  # for reproducibility
import re
import h5py
import os
from nltk import tokenize
from keras.preprocessing.text import Tokenizer, text_to_word_sequence
from attention import Attention_input1, Attention_input2
from keras.preprocessing.text import Tokenizer
from keras.preprocessing.sequence import pad_sequences
from keras.utils.np_utils import to_categorical
from keras.layers import Reshape, Dense, Input, Flatten, Dropout, merge, BatchNormalization
from keras.layers import TimeDistributed, LSTM, GRU, Bidirectional
from keras.models import Model
from keras.optimizers import SGD, Adadelta, Adam, RMSprop
from keras.callbacks import EarlyStopping, ModelCheckpoint
from keras.layers.core import Reshape, RepeatVector

from keras.callbacks import EarlyStopping
from keras.models import Sequential, Model
from keras.layers import Activation, Dense, Dropout, Embedding, Flatten, Input, Merge, Convolution1D, MaxPooling1D

GLOVE_DIR = '../data/'

MAX_SEQUENCE_LENGTH = 140
MAX_NB_WORDS = 10000
EMBEDDING_DIM = 200
VALIDATION_SPLIT = 0.1
NB_EPOCH = 100
NB_CLASS = 3
DIM_HIDDEN = 128
DIM_LSTM = 128

# datamode = 'mul'
datamode = 'single'

if datamode == 'mul':
    DATA_PATH = '../data/MSVA_multiple_17024.h5'
    BATCH_SIZE = 128
else:
    DATA_PATH = '../data/MSVA_single_4511.h5'
    BATCH_SIZE = 32

def load_data():
    read_file = h5py.File(DATA_PATH, 'r')
    texts = read_file['txt_data'][:]
    labels = read_file['label'][:]
    scenes = read_file['scene_data'][:]
    objects = read_file['object_data'][:]
    return texts,labels,scenes,objects
def split_data(data,VALIDATION_SPLIT):
    nb_validation_samples = int(VALIDATION_SPLIT * data.shape[0])
    data_train = data[:-(nb_validation_samples * 2)]
    data_val = data[-(nb_validation_samples * 2):-(nb_validation_samples)]
    data_test = data[-nb_validation_samples:]
    return data_train,data_val,data_test

def dp_txt(txt):
    # nonEnglish_regex = re.compile('[^a-zA-Z0-9\\?\\!\\,\\.@#\\+\\-=\\*\'\"><&\\$%\\(\\)\\[\\]:;]+')
    hashtag_pattern = re.compile('#[a-zA-Z0-9]+')
    at_pattern = re.compile('@[a-zA-Z0-9]+')
    http_pattern = re.compile("((http|ftp|https)://)(([a-zA-Z0-9\._-]+\.[a-zA-Z]{2,6})|([0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}))(:[0-9]{1,4})*(/[a-zA-Z0-9\&%_\./-~-]*)?")
    txt = txt.strip()
    txt_hashtag = re.sub(hashtag_pattern, '', txt)
    txt_nonat = re.sub(at_pattern, '', txt_hashtag)
    txt_nonhttp = re.sub(http_pattern, '', txt_nonat)
    txt = txt_nonhttp
    return txt

def fun():
    texts,labels,scenes,objects = load_data()
    new_texts = []
    for idx in range(len(texts)):
        text = texts[idx]
        text = dp_txt(str(text))
        new_texts.append(text)
    texts = new_texts

    tokenizer = Tokenizer(nb_words=MAX_NB_WORDS)
    tokenizer.fit_on_texts(texts)
    sequences = tokenizer.texts_to_sequences(texts)
    word_index = tokenizer.word_index

    text_data = pad_sequences(sequences, maxlen=MAX_SEQUENCE_LENGTH)
    labels = to_categorical(np.asarray(labels))

    # print('Text tensor shape:', text_data.shape)
    # print('Label tensor shape:', labels.shape)
    # print('Scene tensor shape:', scenes.shape)
    # print('Object tensor shape:', objects.shape)
    # # split the text_data into a training set and a validation set
    rand = np.arange(labels.shape[0])
    np.random.shuffle(rand)
    indices = rand

    text_data = text_data[indices]
    labels = labels[indices]
    scenes = scenes[indices]
    objects = objects[indices]
    text_train,text_val,text_test = split_data(text_data,VALIDATION_SPLIT)
    label_train,label_val,label_test = split_data(labels,VALIDATION_SPLIT)
    scene_train,scene_val,scene_test = split_data(scenes,VALIDATION_SPLIT)
    object_train,object_val,object_test = split_data(objects,VALIDATION_SPLIT)

    text_shape = text_train.shape[1:]
    scene_shape = scene_train.shape[1:]
    object_shape = object_train.shape[1:]

    embeddings_index = {}
    f = open(os.path.join(GLOVE_DIR, 'glove.6B.200d.txt'))
    for line in f:
        values = line.split()
        word = values[2]
        coefs = np.asarray(values[1], dtype='float32')
        embeddings_index[word] = coefs
    f.close()

    nb_words = min(MAX_NB_WORDS, len(word_index))
    embedding_matrix = np.zeros((nb_words + 1, EMBEDDING_DIM))
    for word, i in word_index.items():
        if i > MAX_NB_WORDS:
            continue
        embedding_vector = embeddings_index.get(word)
        if embedding_vector is not None:
            embedding_matrix[i] = embedding_vector

    embedding_layer = Embedding(nb_words + 1,
                                EMBEDDING_DIM,
                                weights=[embedding_matrix],
                                input_length=MAX_SEQUENCE_LENGTH,
                                trainable=True)

    save_best = ModelCheckpoint('../../model/{}.hdf5'.format('my_weight'), save_best_only=True)
    elstop = EarlyStopping(monitor='val_loss', min_delta=1e-4, patience=5)

    # Image Sence
    scene_input = Input(shape=scene_shape, dtype='float32')
    img_scene = Dense(DIM_HIDDEN, activation='relu')(scene_input)
    img_scene_encoder = RepeatVector(text_shape[0], name='scene-repeat')(img_scene)

    # Image Object
    object_input = Input(shape=object_shape, dtype='float32')
    img_object = Dense(DIM_HIDDEN, activation='relu')(object_input)
    img_object_encoder = RepeatVector(text_shape[0], name='object-repeat')(img_object)

    # Text
    txt_input = Input(shape=text_shape, dtype='float32')
    txt = embedding_layer(txt_input)
    txt_hidden = (LSTM(DIM_HIDDEN, return_sequences=True, name='tweet-lstm'))(txt)
    txt_att = Attention_input2(name='att_so')([txt_hidden, img_object_encoder, img_scene_encoder])

    # Merge
    img_txt = merge([img_scene, img_object, txt_att], mode='concat')
    img_txt = Dense(DIM_HIDDEN, activation='relu')(img_txt)
    img_txt_loss = Dense(NB_CLASS, activation='softmax', name='main_output')(img_txt)

    model = Model(input=[txt_input, scene_input, object_input], output=[img_txt_loss])
    model.compile(loss='categorical_crossentropy', optimizer='RMSprop',
                  metrics=['acc', 'fmeasure'])
    model.fit([text_train, scene_train, object_train], [label_train],
              validation_data=([text_val, scene_val, object_val], [label_val]),
              nb_epoch=NB_EPOCH, batch_size=BATCH_SIZE, callbacks=[elstop,save_best], verbose=1)

    model.load_weights('../../model/{}.hdf5'.format('my_weight'))
    score = model.evaluate([text_test, scene_test, object_test], label_test, verbose=0)

    print('results：', score[1], score[2])
    return score[1:]

if __name__ == '__main__':
    fun()


import numpy as np
import queue

from scipy.misc import imsave

from screen import grab_frame

class GameReader(object):
    """
    The GameReader object parses 15 frames at a time and returns the game state
    that shall be interpreted by the game loop.

    Constructor initializes the object, frame rate and the queue that stores
    our frames.

    Inputs:
     - fps: An integer represeting how many frames we would like to read per
       second.
     - max_queue_size: An integer represeting how many frames we will be
       holding in memory.
    """
    def __init__(self, fps, max_queue_size):
        self.fps = fps
        self.max_queue_size = max_queue_size
        self.frame_queue = queue.Queue(maxsize=max_queue_size)
        self.frame_shape = None
        self.is_dead = False

        #To debug frames
        self.imsave_img_count = 0
        self.imsave_queue_count = 0
        self.save_dir = "../debug_frames/"

    """
    Grabs the next frame in the game and updates the frame queue. The frame
    queue must not have more than 20 frames. If it is full, we discard the
    oldest frame and insert the newest one.

    Returns the oldest frame in the queue.
    """
    def update_frames(self):
        popped_frame = None
        frame = grab_frame()
        if self.frame_shape == None:
            self.frame_shape = frame.shape
        if self.frame_queue.full():
            popped_frame = self.frame_queue.get()
            self.frame_queue.put(frame)
        else:
            self.frame_queue.put(frame)
        return popped_frame

    """
    Returns the average grayscale value from every pixel curerntly stored in the
    frame queue.
    """
    def avg_luminosity(self):
        frame_length = float(len(self.frame_queue))
        return np.sum(np.array(list(self.frame_queue.queue))) / frame_length

    """
    Returns whether or not the player is dead, indicating the restarting of a
    new trial.
    """
    def isdead(self):
        if self.frame_shape == None:
            return True
        prev_frame_val = np.inf
        flag_count = 0
        for img in list(self.frame_queue.queue):
            #Sum up every pixel value
            W, H = self.frame_shape
            frame_val = float(np.sum(img)) / (W * H)
            print(frame_val)

            """
            For some reason, the screen that appears when the player loses,
            appears to be perfect black to the naked eye, but is in actuality,
            no where near. The screen has a grayscale value of about 13 out of
            a maximum 255.
            """
            if frame_val < 0.06:
                return True
        return False
    """
    Analyzes the queue of frames and returns the state of the game.

    TODO: Need to possibly update
    Here are the prospective states
     - playing: The agent is current in play and is running normally. In other
       words, the player is currently acting normally and is neither dead,
       recovering dead nor transitioning from level to level.
     - dead: The agent has made a mistake and is now resetting itself to be
       playing again.
     - transition: The game is now in a transitional phase and we must wait
       until it is not to begin acting again.
       # NOTE: Might not need this
    """
    def parse(self):
        pass

    def save_queue(self):
        for img in list(self.frame_queue.queue):
            img_name = ("queue" + str(self.imsave_queue_count) + "img" +
                str(self.imsave_img_count) + ".png")

            mod_dir = self.save_dir + img_name
            print("Saving" + img_name)
            imsave(mod_dir, img)

            self.imsave_img_count += 1
        self.imsave_queue_count += 1


import argparse
import os
import time
from functools import partial

import torch
import pickle
import numpy as np

import torch.nn as nn
from torch.utils.tensorboard import SummaryWriter
from torch.optim import Adam
from tqdm import tqdm

import model.utils as model_utils
import utils
from data.dataset import SeparationDataset
from data.musdb import get_musdb_folds
from data.utils import crop_targets, random_amplify
from test import evaluate, validate
from model.waveunet import Waveunet

def main(args):
    #torch.backends.cudnn.benchmark=True # This makes dilated conv much faster for CuDNN 7.5

    # MODEL
    num_features = [args.features*i for i in range(1, args.levels+1)] if args.feature_growth == "add" else \
                   [args.features*2**i for i in range(0, args.levels)]
    target_outputs = int(args.output_size * args.sr)
    model = Waveunet(args.channels, num_features, args.channels, args.instruments, kernel_size=args.kernel_size,
                     target_output_size=target_outputs, depth=args.depth, strides=args.strides,
                     conv_type=args.conv_type, res=args.res, separate=args.separate)

    if args.cuda:
        model = model_utils.DataParallel(model)
        print("move model to gpu")
        model.cuda()

    print('model: ', model)
    print('parameter count: ', str(sum(p.numel() for p in model.parameters())))

    writer = SummaryWriter(args.log_dir)

    ### DATASET
    musdb = get_musdb_folds(args.dataset_dir,args.musdb_version)
    # If not data augmentation, at least crop targets to fit model output shape
    crop_func = partial(crop_targets, shapes=model.shapes)
    # Data augmentation function for training
    augment_func = partial(random_amplify, shapes=model.shapes, min=0.7, max=1.0)
    train_data = SeparationDataset(musdb, "train", args.instruments, args.sr, args.channels, model.shapes, True, args.hdf_dir, audio_transform=augment_func)
    val_data = SeparationDataset(musdb, "val", args.instruments, args.sr, args.channels, model.shapes, False, args.hdf_dir, audio_transform=crop_func)
    test_data = SeparationDataset(musdb, "test", args.instruments, args.sr, args.channels, model.shapes, False, args.hdf_dir, audio_transform=crop_func)

    dataloader = torch.utils.data.DataLoader(train_data, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, worker_init_fn=utils.worker_init_fn)

    ##### TRAINING ####

    # Set up the loss function
    if args.loss == "L1":
        criterion = nn.L1Loss()
    elif args.loss == "L2":
        criterion = nn.MSELoss()
    else:
        raise NotImplementedError("Couldn't find this loss!")

    # Set up optimiser
    optimizer = Adam(params=model.parameters(), lr=args.lr)

    # Set up training state dict that will also be saved into checkpoints
    state = {"step" : 0,
             "worse_epochs" : 0,
             "epochs" : 0,
             "best_loss" : np.Inf}

    # LOAD MODEL CHECKPOINT IF DESIRED
    if args.load_model is not None:
        print("Continuing training full model from checkpoint " + str(args.load_model))
        state = model_utils.load_model(model, optimizer, args.load_model, args.cuda)

    print('TRAINING START')
    while state["worse_epochs"] < args.patience:
        print("Training one epoch from iteration " + str(state["step"]))
        avg_time = 0.
        model.train()
        with tqdm(total=len(train_data) // args.batch_size) as pbar:
            np.random.seed()
            for example_num, (x, targets) in enumerate(dataloader):
                if args.cuda:
                    x = x.cuda()
                    for k in list(targets.keys()):
                        targets[k] = targets[k].cuda()

                t = time.time()

                # Set LR for this iteration
                utils.set_cyclic_lr(optimizer, example_num, len(train_data) // args.batch_size, args.cycles, args.min_lr, args.lr)
                writer.add_scalar("lr", utils.get_lr(optimizer), state["step"])

                # Compute loss for each instrument/model
                optimizer.zero_grad()
                outputs, avg_loss = model_utils.compute_loss(model, x, targets, criterion, compute_grad=True)

                optimizer.step()

                state["step"] += 1

                t = time.time() - t
                avg_time += (1. / float(example_num + 1)) * (t - avg_time)

                writer.add_scalar("train_loss", avg_loss, state["step"])

                if example_num % args.example_freq == 0:
                    input_centre = torch.mean(x[0, :, model.shapes["output_start_frame"]:model.shapes["output_end_frame"]], 0) # Stereo not supported for logs yet
                    writer.add_audio("input", input_centre, state["step"], sample_rate=args.sr)

                    for inst in outputs.keys():
                        writer.add_audio(inst + "_pred", torch.mean(outputs[inst][0], 0), state["step"], sample_rate=args.sr)
                        writer.add_audio(inst + "_target", torch.mean(targets[inst][0], 0), state["step"], sample_rate=args.sr)

                pbar.update(1)

        # VALIDATE
        val_loss = validate(args, model, criterion, val_data)
        print("VALIDATION FINISHED: LOSS: " + str(val_loss))
        writer.add_scalar("val_loss", val_loss, state["step"])

        # EARLY STOPPING CHECK
        checkpoint_path = os.path.join(args.checkpoint_dir, "checkpoint_" + str(state["step"]))
        if val_loss >= state["best_loss"]:
            state["worse_epochs"] += 1
        else:
            print("MODEL IMPROVED ON VALIDATION SET!")
            state["worse_epochs"] = 0
            state["best_loss"] = val_loss
            state["best_checkpoint"] = checkpoint_path

        # CHECKPOINT
        print("Saving model...")
        model_utils.save_model(model, optimizer, state, checkpoint_path)

        state["epochs"] += 1

    #### TESTING ####
    # Test loss
    print("TESTING")

    # Load best model based on validation loss
    state = model_utils.load_model(model, None, state["best_checkpoint"], args.cuda)
    test_loss = validate(args, model, criterion, test_data)
    print("TEST FINISHED: LOSS: " + str(test_loss))
    writer.add_scalar("test_loss", test_loss, state["step"])

    # Mir_eval metrics
    test_metrics = evaluate(args, musdb["test"], model, args.instruments)

    # Dump all metrics results into pickle file for later analysis if needed
    with open(os.path.join(args.checkpoint_dir, "results.pkl"), "wb") as f:
        pickle.dump(test_metrics, f)

    # Write most important metrics into Tensorboard log
    avg_SDRs = {inst : np.mean([np.nanmean(song[inst]["SDR"]) for song in test_metrics]) for inst in args.instruments}
    avg_SIRs = {inst : np.mean([np.nanmean(song[inst]["SIR"]) for song in test_metrics]) for inst in args.instruments}
    for inst in args.instruments:
        writer.add_scalar("test_SDR_" + inst, avg_SDRs[inst], state["step"])
        writer.add_scalar("test_SIR_" + inst, avg_SIRs[inst], state["step"])
    overall_SDR = np.mean([v for v in avg_SDRs.values()])
    writer.add_scalar("test_SDR", overall_SDR)
    print("SDR: " + str(overall_SDR))

    writer.close()

if __name__ == '__main__':
    ## TRAIN PARAMETERS
    parser = argparse.ArgumentParser()
    parser.add_argument('--instruments', type=str, nargs='+', default=["bass", "drums", "other", "vocals"],
                        help="List of instruments to separate (default: \"bass drums other vocals\")")
    parser.add_argument('--cuda', action='store_true',
                        help='Use CUDA (default: False)')
    parser.add_argument('--num_workers', type=int, default=1,
                        help='Number of data loader worker threads (default: 1)')
    parser.add_argument('--features', type=int, default=32,
                        help='Number of feature channels per layer')
    parser.add_argument('--log_dir', type=str, default='logs/waveunet',
                        help='Folder to write logs into')
    parser.add_argument('--dataset_dir', type=str, default="/mnt/windaten/Datasets/MUSDB18HQ",
                        help='Dataset path')
    parser.add_argument('--hdf_dir', type=str, default="hdf",
                        help='Dataset path')
    parser.add_argument('--checkpoint_dir', type=str, default='checkpoints/waveunet',
                        help='Folder to write checkpoints into')
    parser.add_argument('--load_model', type=str, default=None,
                        help='Reload a previously trained model (whole task model)')
    parser.add_argument('--lr', type=float, default=1e-3,
                        help='Initial learning rate in LR cycle (default: 1e-3)')
    parser.add_argument('--min_lr', type=float, default=5e-5,
                        help='Minimum learning rate in LR cycle (default: 5e-5)')
    parser.add_argument('--cycles', type=int, default=2,
                        help='Number of LR cycles per epoch')
    parser.add_argument('--batch_size', type=int, default=4,
                        help="Batch size")
    parser.add_argument('--levels', type=int, default=6,
                        help="Number of DS/US blocks")
    parser.add_argument('--depth', type=int, default=1,
                        help="Number of convs per block")
    parser.add_argument('--sr', type=int, default=44100,
                        help="Sampling rate")
    parser.add_argument('--channels', type=int, default=2,
                        help="Number of input audio channels")
    parser.add_argument('--kernel_size', type=int, default=5,
                        help="Filter width of kernels. Has to be an odd number")
    parser.add_argument('--output_size', type=float, default=2.0,
                        help="Output duration")
    parser.add_argument('--strides', type=int, default=4,
                        help="Strides in Waveunet")
    parser.add_argument('--patience', type=int, default=20,
                        help="Patience for early stopping on validation set")
    parser.add_argument('--example_freq', type=int, default=200,
                        help="Write an audio summary into Tensorboard logs every X training iterations")
    parser.add_argument('--loss', type=str, default="L1",
                        help="L1 or L2")
    parser.add_argument('--conv_type', type=str, default="gn",
                        help="Type of convolution (normal, BN-normalised, GN-normalised): normal/bn/gn")
    parser.add_argument('--res', type=str, default="fixed",
                        help="Resampling strategy: fixed sinc-based lowpass filtering or learned conv layer: fixed/learned")
    parser.add_argument('--separate', type=int, default=1,
                        help="Train separate model for each source (1) or only one (0)")
    parser.add_argument('--feature_growth', type=str, default="double",
                        help="How the features in each layer should grow, either (add) the initial number of features each time, or multiply by 2 (double)")
    parser.add_argument('--musdb_version', type=str, default="HQ",
                        help="What version of the MUSDB do we use? HQ or compressed?")

    args = parser.parse_args()

    main(args)


"""
about:
generate a simulation.input file given the (modified with Z0/Z1) CIF file

usage:
python simgen.py MOF.cif ads.pdb

Haoyuan Chen, 03.17.2020
"""

import sys
from ase.io import *
import numpy as np

cifname = sys.argv[1]
adsname = sys.argv[2]

# can change these values
cutoff = 12.0
ncycles = 1000
ninit = 0
nprint = 100
nmovie = 10
temp = 100.0
nads = 1

def unitcells(cifname,cutoff):
    c = read(cifname)
    cells = [1,1,1]
    if c.cell.orthorhombic:
        for i in range(3):
            cells[i] = np.ceil(2*cutoff/c.cell[i][i])
    else:
    # below is placeholder
    # will implement the MIC for triclinic later
    # https://pdfs.semanticscholar.org/a5bd/f949911c0d7333ffec488b0589e88cd722b6.pdf
        for i in range(3):
            cells[i] = np.ceil(2*cutoff/c.cell[i][i])
    return cells

cells = unitcells(cifname,cutoff)

f = open('simulation.input','w')
f.write('SimulationType  MonteCarlo\nNumberOfCycles  %d\nNumberOfInitializationCycles  %d\nPrintEvery  %d\n'%(ncycles,ninit,nprint))
f.write('Forcefield  GenericMOFs\nCutOffVDW  %.1f\n'%(cutoff))
f.write('Framework  0\nFrameworkName  %s\nUnitCells  %d %d %d\n'%(cifname.replace('.cif',''),cells[0],cells[1],cells[2]))
f.write('Movies  yes\nWriteMoviesEvery  %d\nExternalTemperature  %.1f\n'%(nmovie,temp))
f.write('Component 0 MoleculeName  %s\n            MoleculeDefinition  local\n            StartingBead  0\n            TranslationProbability  1.0\n            RotationProbability  1.0\n            ReinsertionProbability  1.0\n            CreateNumberOfMolecules  %d\n'%(adsname.replace('.pdb',''),nads))
f.close()


import numpy
import astropy.io.fits as fits
from numina.core import Result, Requirement, Parameter
from numina.types.obsresult import ObservationResultType
from numina.core.recipes import BaseRecipe
from numina.core.validator import range_validator
from clodiadrp.products import MasterBias, MasterFlat

class Flat(BaseRecipe):

    obresult = Requirement(ObservationResultType, "Observation Result")
    master_bias = Requirement(MasterBias, "Master Bias")
    polynomial_degree = Parameter(5, 'Polynomial degree of arc calibration',
                                  as_list=True, nelem='+',
                                  validator=range_validator(minval=1)
                                  )

    @master_bias.validator
    def custom_validator(self, value):
        print('func', self, value)
        return True

    master_flat = Result(MasterFlat)

    def run(self, rinput):

        # Here the raw images are processed
        # and a final image myframe is created
        obresult = rinput.obresult
        fr0 = obresult.frames[0].open()
        data = numpy.ones_like(fr0[0].data)
        hdu = fits.PrimaryHDU(data, header=fr0[0].header)
        myframe = fits.HDUList([hdu])
        #
        result = self.create_result(master_flat=myframe)
        return result


# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.3'
#       jupytext_version: 0.8.6
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

import numpy as np
from numpy import dot
from numpy.linalg import norm
from boto3 import client as boto3_client
import json
import logging
from botocore.client import Config
from dataclasses import dataclass
from scorer.pre_process import preprocess_text
from s3client.s3 import S3Manager
import os
import pickle

logger = logging.getLogger()

config = Config(connect_timeout=60, read_timeout=240, retries={"max_attempts": 0},)
lambda_client = boto3_client("lambda", config=config)

@dataclass
class TextSegment:
    id: str
    text: str
    speaker: str

@dataclass
class Score(TextSegment):
    score: float

def upload_fv(fv_list, Request, context_id, instance_id):
    try:
        bucket = "io.etherlabs." + os.getenv("ACTIVE_ENV", "staging2") + ".contexts"
        s3_path = (
            context_id + "/feature_vectors/" + instance_id + "/" + Request.id + ".pkl"
        )
        logger.info("The path used for s3.", extra={"S3": s3_path, "bucket": bucket})
        s3_obj = S3Manager(bucket_name=bucket)
        s3_obj.upload_object(pickle.dumps(fv_list), s3_path)
    except Exception as e:
        logger.info("Uploading failed ", extra={"exception:": e})
        return False
    return True

def get_score(
    mind_id: str,
    mind_dict,
    Request: TextSegment,
    context_id,
    instance_id,
    for_pims=False,
) -> Score:
    score = []
    pre_processed_input = preprocess_text(Request.text)
    lambda_function = "mind-" + mind_id
    transcript_text = Request.text
    if len(pre_processed_input) != 0:
        mind_input = json.dumps({"text": pre_processed_input})
        mind_input = json.dumps({"body": mind_input})
        logger.info("sending request to mind service")
        if for_pims is False:
            transcript_score = get_feature_vector(
                mind_input,
                lambda_function,
                mind_dict,
                Request,
                context_id,
                instance_id,
            )
        else:
            response = get_feature_vector(
                mind_input,
                lambda_function,
                mind_dict,
                Request,
                context_id,
                instance_id,
                store_features=True,
            )
            if response is True:
                return True
            else:
                return False
    else:
        return True
        transcript_score = 0.00001
        logger.warn("processing transcript: {}".format(transcript_text))
        logger.warn("transcript too small to process. Returning default score")
    # hack to penalize out of domain small transcripts coming as PIMs - word level
    if len(transcript_text.split(" ")) < 40:
        transcript_score = 0.1 * transcript_score
    score = 1 / transcript_score
    return score

def getClusterScore(mind_vec, sent_vec):
    n1 = norm(mind_vec, axis=1).reshape(1, -1)
    n2 = norm(sent_vec, axis=1).reshape(-1, 1)
    dotp = dot(sent_vec, mind_vec).squeeze(2)
    segment_scores = dotp / (n2 * n1)
    return segment_scores

def get_feature_vector(
    mind_input,
    lambda_function,
    mind_dict,
    Request,
    context_id,
    instance_id,
    store_features=False,
):
    invoke_response = lambda_client.invoke(
        FunctionName=lambda_function,
        InvocationType="RequestResponse",
        Payload=mind_input,
    )
    out_json = invoke_response["Payload"].read().decode("utf8").replace("'", '"')
    data = json.loads(json.loads(out_json)["body"])
    response = json.loads(out_json)["statusCode"]
    if store_features is True:
        vector_list = data["sent_feats"][0]
        upload_fv(vector_list, Request, context_id, instance_id)
        return True

    feats = list(mind_dict["feature_vector"].values())
    mind_vector = np.array(feats).reshape(len(feats), -1)
    transcript_score = 0.00001
    transcript_mind_list = []
    transcript_score_list = []
    if response == 200:
        logger.info("got {} from mind server".format(response))
        feature_vector = np.array(data["sent_feats"][0])
        if len(feature_vector) > 0:
            # For paragraphs, uncomment below LOC
            # feature_vector = np.mean(np.array(feature_vector),0).reshape(1,-1)
            batch_size = min(10, feature_vector.shape[0])
            for i in range(0, feature_vector.shape[0], batch_size):
                mind_vec = np.expand_dims(np.array(mind_vector), 2)
                sent_vec = feature_vector[i : i + batch_size]

                cluster_scores = getClusterScore(mind_vec, sent_vec)

                batch_scores = cluster_scores.max(1)
                transcript_score_list.extend(batch_scores)

                minds_selected = cluster_scores.argmax(1)
                transcript_mind_list.extend(minds_selected)
            transcript_score = np.mean(transcript_score_list)
            logger.info(
                "Mind Selected is {}".format(
                    {
                        ele: transcript_mind_list.count(ele)
                        for ele in set(transcript_mind_list)
                    }
                )
            )
    else:
        logger.debug(
            "Invalid response from mind service for input: {}".format(mind_input)
        )
        logger.debug("Returning default score")

    return transcript_score


from __future__ import print_function
import argparse
import skimage
import skimage.io
import skimage.transform
from PIL import Image
from math import log10
import sys
import shutil
import os
import re
from struct import unpack
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torch.nn.functional as F
from torch.autograd import Variable
from torch.utils.data import DataLoader
from models.GANet_deep import GANet
from torch.utils.data import DataLoader
from dataloader.dataset import DatasetFromList
# from dataloader.data import get_test_set
import numpy as np
from natsort import natsorted

parser = argparse.ArgumentParser(description='PyTorch GANet Example')
parser.add_argument('--max_disp', type=int, default=192, help="max disp")
parser.add_argument('--resume', type=str, default='', help="resume from saved model")
parser.add_argument('--cuda', type=bool, default=True, help='use cuda?')
parser.add_argument('--dataset', type=str, default='sceneflow', help='sceneflow, kitti2015')
parser.add_argument('--data_path', type=str, required=True, help="data root")
parser.add_argument('--save_path', type=str, default='./result/', help="location to save result")
parser.add_argument('--threshold', type=float, default=3.0, help="threshold of error rates")
parser.add_argument('--multi_gpu', type=int, default=0, help="multi_gpu choice")
parser.add_argument('--model', type=str, default='GANet_deep', help="model to train")
parser.add_argument('--batch_size', type=int, default=3)
parser.add_argument('--within_max_disp', action='store_true')

opt = parser.parse_args()

print(opt)

cuda = opt.cuda
# cuda = True
if cuda and not torch.cuda.is_available():
    raise Exception("No GPU found, please run without --cuda")

# print('===> Loading datasets')
# test_set = get_test_set(opt.data_path, opt.test_list, [opt.crop_height, opt.crop_width], false, opt.kitti, opt.kitti2015)
# testing_data_loader = DataLoader(dataset=test_set, num_workers=opt.threads, batch_size=opt.testBatchSize, shuffle=False)

print('===> Building model')
model = GANet(opt.max_disp)
print('Number of model parameters: {}'.format(sum([p.data.nelement() for p in model.parameters()])))

if cuda:
    model = torch.nn.DataParallel(model).cuda()

if opt.resume:
    if os.path.isfile(opt.resume):
        print("=> loading checkpoint '{}'".format(opt.resume))
        checkpoint = torch.load(opt.resume)
        model.load_state_dict(checkpoint['state_dict'], strict=False)

    else:
        print("=> no checkpoint found at '{}'".format(opt.resume))

def test(input1, input2):
    _, _, height, width = input1.size()
    input1 = Variable(input1, requires_grad=False)
    input2 = Variable(input2, requires_grad=False)

    model.eval()
    if cuda:
        input1 = input1.cuda()
        input2 = input2.cuda()
    with torch.no_grad():
        prediction = model(input1, input2)

    temp = prediction
    # if height <= opt.crop_height and width <= opt.crop_width:
    #     temp = temp[:, opt.crop_height - height: opt.crop_height, opt.crop_width - width: opt.crop_width]
    # else:
    #     temp = temp[:, :, :]
    # skimage.io.imsave(savename, (temp * 256).astype('uint16'))
    return temp

if __name__ == "__main__":
    if opt.dataset == 'sceneflow':
        file_path = opt.data_path
        directory = os.path.join(file_path, 'frame_finalpass', 'TEST')
        sub_folders = [os.path.join(directory, subset) for subset in os.listdir(directory) if
                       os.path.isdir(os.path.join(directory, subset))]

        seq_folders = []
        for sub_folder in sub_folders:
            seq_folders += [os.path.join(sub_folder, seq) for seq in os.listdir(sub_folder) if
                            os.path.isdir(os.path.join(sub_folder, seq))]

        left_data = []
        for seq_folder in seq_folders:
            left_data += [os.path.join(seq_folder, 'left', img) for img in
                          os.listdir(os.path.join(seq_folder, 'left'))]

        left_data = natsorted(left_data)
        right_data = [left.replace('left', 'right') for left in left_data]
        disp_data = [left.replace('frame_finalpass', 'disparity').replace('.png', '.pfm') for left in left_data]

        directory = os.path.join(file_path, 'occlusion', 'TEST', 'left')
        occ_data = [os.path.join(directory, occ) for occ in os.listdir(directory)]
        occ_data = natsorted(occ_data)
    elif opt.dataset == 'kitti2015':
        file_path = opt.data_path
        left_data = natsorted(
            [os.path.join(file_path, 'image_2/', img) for img in os.listdir(os.path.join(file_path, 'image_2/')) if
             img.find('_10') > -1])
        right_data = [img.replace('image_2/', 'image_3/') for img in left_data]
        disp_data = [img.replace('image_2/', 'disp_noc_0/') for img in left_data]
        occ_data = None
    elif opt.dataset == 'middlebury':
        file_path = opt.data_path
        left_data = [os.path.join(file_path, obj, 'im0.png') for obj in os.listdir(file_path)]
        right_data = [os.path.join(file_path, obj, 'im1.png') for obj in os.listdir(file_path)]
        disp_data = [os.path.join(file_path, obj, 'disp0GT.pfm') for obj in os.listdir(file_path)]
        occ_data = [os.path.join(file_path, obj, 'mask0nocc.png') for obj in os.listdir(file_path)]

        left_data = natsorted(left_data)
        right_data = natsorted(right_data)
        disp_data = natsorted(disp_data)
        occ_data = natsorted(occ_data)
    elif opt.dataset == 'scared':
        file_path = opt.data_path
        left_data = natsorted(
            [os.path.join(file_path, 'left/', img) for img in os.listdir(os.path.join(file_path, 'left'))])
        right_data = [img.replace('left/', 'right/') for img in left_data]
        disp_data = [img.replace('left/', 'disparity/').replace('.png', '.tiff') for img in left_data]
        occ_data = [img.replace('left/', 'occlusion/') for img in left_data]

    test_data = DatasetFromList(left_data, right_data, disp_data, occ_data)
    test_dataloader = DataLoader(test_data, batch_size=opt.batch_size, shuffle=False, num_workers=2)

    avg_error = 0.0
    avg_wrong = 0.0
    avg_total = 0.0
    for index, (left, right, disp) in enumerate(test_dataloader):
        prediction = test(left, right)
        disp = disp.squeeze(1).cuda()  # NxHxW
        if opt.within_max_disp:
            mask = torch.logical_and(disp > 0.0, disp < opt.max_disp)
        else:
            mask = disp > 0.0

        error = F.l1_loss(prediction[mask], disp[mask], reduction='none')
        wrong = torch.sum(error > opt.threshold).item()
        total = error.numel()
        avg_error = avg_error + torch.mean(error)
        avg_wrong += wrong
        avg_total += total
        print("===> Iteration {}: ".format(index) + " ==> EPE Error: {:.4f}, Error Rate: {:.4f}".format(
            torch.mean(error), wrong / total * 100))

        # np.save('left.npy', left.data.cpu())
        # np.save('right.npy', right.data.cpu())
        # np.save('disp_pred.npy', prediction)
        # np.save('disp.npy', disp)

        # import matplotlib.pyplot as plt
        # plt.figure()
        # plt.imshow(disp[0])
        # plt.figure()
        # plt.imshow(prediction[0])
        # plt.show()
        # asdf

    avg_error = avg_error / len(test_dataloader)
    avg_rate = avg_wrong / avg_total * 100
    print("===> Total {} Frames ==> AVG EPE Error: {:.4f}, AVG Error Rate: {:.4f}".format(len(test_dataloader),
                                                                                avg_error, avg_rate))


from __future__ import print_function
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.autograd import Variable
import os
import time
import shutil
import itertools
import numpy as np
from tqdm import tqdm
from utils import Bar, Logger, AverageMeter, regression_accuracy, mkdir_p, savefig

# Training settings
parser = argparse.ArgumentParser(description='Observer Network')
parser.add_argument('--train-batch', default=128, type=int, metavar='N',
                    help='train batchsize')
parser.add_argument('--test-batch', default=100, type=int, metavar='N',
                    help='test batchsize')
parser.add_argument('--lr', type=float, default=0.01, metavar='LR',
                    help='learning rate (default: 0.0001)')
parser.add_argument('--schedule', type=int, nargs='+', default=[10, 30, 50, 80],
                        help='Decrease learning rate at these epochs.')
parser.add_argument('--gamma', type=float, default=0.5, help='LR is multiplied by gamma on schedule.')
parser.add_argument('--epochs', type=int, default=100, metavar='N',
                    help='number of epochs to train (default: 10)')
parser.add_argument('--seed', type=int, default=1, metavar='S',
                    help='random seed (default: 1)')
parser.add_argument('-j', '--workers', default=4, type=int, metavar='N',
                    help='number of data loading workers (default: 4)')
parser.add_argument('-c', '--checkpoint', default='checkpoint', type=str, metavar='PATH',
                    help='path to save checkpoint (default: checkpoint)')
parser.add_argument('--n-classes', default=3, type=int, metavar='N',
                    help='number of classes')

args = parser.parse_args()
state = {k: v for k, v in args._get_kwargs()}
args.cuda = torch.cuda.is_available()

torch.manual_seed(args.seed)
if args.cuda:
    torch.cuda.manual_seed(args.seed)

best_acc = float("-inf")

kwargs = {'num_workers': 1, 'pin_memory': True} if args.cuda else {}

# from models.cnned import SegNet
from new_models.model import TrajPredictor

# model = SegNet(n_classes=args.n_classes, is_unpooling=True)
model = TrajPredictor()

if args.cuda:
    model = torch.nn.DataParallel(model).cuda()
if args.cuda:
    ckpt = torch.load('./checkpoint_half/checkpoint.pth.tar')
    model.load_state_dict(ckpt['state_dict'])
else:
    ckpt = torch.load('./checkpoint_half/checkpoint.pth.tar', map_location=lambda storage, loc: storage)
    from collections import OrderedDict
    cpu_state_dict = OrderedDict()
    for k, v in ckpt['state_dict'].items():
        name = k[7:] # remove `module.`
        cpu_state_dict[name] = v
    model.load_state_dict(cpu_state_dict)

lossfun = nn.MSELoss()

def mkdir(folder):
    if os.path.exists(folder):
        shutil.rmtree(folder)
    os.makedirs(folder)

policy = 'single_obs_1_food'
test_folder = 'data_arr_progressive_' + policy
test_files = os.listdir(test_folder)
for p_file in tqdm(test_files):
    filepath = os.path.join(test_folder, p_file)
    num_frames = len(os.listdir(filepath))
    for frame_idx in range(num_frames):
        arr_filename = 'frame_' + str(frame_idx) + '.npy'
        arr_filepath = os.path.join(filepath, arr_filename)
        arr = np.load(arr_filepath)

        arr = arr / 255.0
        arr = torch.tensor(arr)
        arr = arr.unsqueeze(0)
        arr = arr.permute(0, 3, 1, 2)
        if args.cuda:
            arr = arr.cuda()
        arr = Variable(arr)
        output = model(arr.float())

        output = output[0].permute(1, 2, 0).data.cpu().numpy()
        output = (output * 255.0).astype('uint8')

        output_filename = 'output_' + str(frame_idx) + '.npy'
        output_filepath = os.path.join(filepath, output_filename)
        np.save(output_filepath, output)


# -*- coding: utf-8 -*-

from __future__ import print_function

import os
import sys
import argparse
import time
import math
import pandas as pd
import shutil
import multiprocessing

import tensorboard_logger as tb_logger
import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn
from torchvision import transforms, datasets
from tqdm import tqdm
from torchsummary import summary
import numpy as np
import faiss

from data_utils import DynamicLabelDatasetIndex
from data_utils import MyTransform
from util import AverageMeter, ProgressMeter
from util import adjust_learning_rate
from util import set_optimizer, save_model
from util import txt_logger, set_parser, parser_processing, suppress_std, set_save_folder
from networks.resnet_big import SupConResNet, SupCEResNet
from losses import SupConLoss

try:
    import apex
    from apex import amp, optimizers
except ImportError:
    pass

def additional_argument(parser):
    parser.add_argument('--used_attributes', type=int, nargs='+', default=[-1], help='Most informative attributes')
    parser.add_argument('--print_freq', type=int, default=100, help='print freq')

    # temperature
    parser.add_argument('--temp', type=float, default=0.1,
                        help='temperature for loss function')

    parser.add_argument('--threshold', type=float, default=0.1)
    parser.add_argument('--warmup_epoch', type=int, default=100)
    parser.add_argument('--perform_cluster_epoch', type=int, default=1)
    parser.add_argument('--low_dim', type=int, default=128, help="latent dim")
    parser.add_argument('--num_cluster', type=str, default="2500", help="num of cluster")

    return

def parse_option():
    parser_str = 'argument for training'
    opt = set_parser(parser_str, additional_argu_func=additional_argument, if_linear=False)

    parser_processing(opt, default_save_folder=True)

    if "," in opt.num_cluster:
        opt.num_cluster = [int(i) for i in opt.num_cluster.split(',')]
    else:
        opt.num_cluster = [int(opt.num_cluster)]
    return opt

def set_loader(opt):
    # costomize transform
    mytransform = MyTransform(opt)
    if opt.method in ['SupCon', 'SimCLR']:
        train_transform = mytransform.train_transform(ssl=True)
    elif opt.method == 'CE':
        train_transform = mytransform.train_transform(ssl=False)
    else:
        raise NotImplementedError("Method {} is not supported for transform operation.".format(opt.method))

    val_transform = mytransform.val_transform()

    # contruct dataset
    if opt.dataset in ['ut-zap50k', 'ut-zap50k-sub', 'CUB', 'aPascal', 'Wider', 'Google_Image', 'imagenet100', 'full_imagenet']:
        train_meta_df = pd.read_csv(os.path.join(opt.data_folder, opt.instruction, opt.meta_file_train), index_col=0)

        train_dataset = DynamicLabelDatasetIndex(df=train_meta_df, data_path=opt.data_root_folder,
                                gran_lvl=opt.gran_lvl, transform = train_transform)
        train_dataset2 = DynamicLabelDatasetIndex(df=train_meta_df, data_path=opt.data_root_folder,
                                gran_lvl=opt.gran_lvl, transform = val_transform)
    else:
        raise ValueError(opt.dataset)

    train_sampler = None
    opt.n_cls = train_dataset.num_class
    train_loader = torch.utils.data.DataLoader(
        train_dataset, batch_size=opt.batch_size, shuffle=True,
        num_workers=opt.num_workers, pin_memory=False, sampler=train_sampler)

    train_loader2 = torch.utils.data.DataLoader(
        train_dataset2, batch_size=opt.batch_size, shuffle=False,
        num_workers=opt.num_workers, pin_memory=False, sampler=train_sampler)
    return train_loader, train_loader2

def set_model(opt, logger):
    if opt.method in ['SupCon', 'SimCLR']:
        model = SupConResNet(name=opt.model)
    elif opt.method == 'CE':
        model = SupCEResNet(name=opt.model, num_classes=opt.n_cls, verbose=True)

    if opt.method == 'SupCon' or opt.method == 'SimCLR':
        criterion = SupConLoss(temperature=opt.temp)
    elif opt.method == 'CE':
        print("Setting creterion to Cross_entropy")
        criterion = torch.nn.CrossEntropyLoss()

    # enable synchronized Batch Normalization
    if opt.syncBN:
        model = apex.parallel.convert_syncbn_model(model)

    if torch.cuda.is_available():
        if torch.cuda.device_count() > 1:
            print("Used devices: {}".format(torch.cuda.device_count()))
            model.encoder = torch.nn.DataParallel(model.encoder)
        model = model.cuda()
        criterion = criterion.cuda()
        cudnn.benchmark = True

        if opt.resume_model_path:
            # get pre ssl epoch
            ckpt = torch.load(opt.resume_model_path, map_location='cpu')
            state_dict = ckpt['model']
            new_state_dict = {}
            for k, v in state_dict.items():
                if torch.cuda.device_count() > 1:
                    pass
                    # print(k)
                    #if k.split(".")[0] != 'head':
                    #    k = ".".join([k.split(".")[0], "module"] + k.split(".")[1:])
                else:
                    k = k.replace("module.", "")
                new_state_dict[k] = v
            state_dict = new_state_dict
            model.load_state_dict(state_dict)

            logger.logger.info("Model loaded! Pretrained from epoch {}".format(opt.pre_ssl_epoch))

    return model, criterion

def train(train_loader, model, criterion, optimizer, epoch, opt, cluster_result=None):
    """one epoch training"""
    model.train()

    batch_time = AverageMeter('BT')
    data_time = AverageMeter('DT')
    losses = AverageMeter('Loss', ':6.2f')
    display = [batch_time, data_time, losses]

    progress = ProgressMeter(
        len(train_loader),
        display,
        prefix="Epoch: [{}]".format(epoch))

    end = time.time()

    for idx, (images, labels, index) in tqdm(enumerate(train_loader), total=len(train_loader)):
        data_time.update(time.time() - end)
        if opt.method in ['SupCon', 'SimCLR']:
            images = torch.cat([images[0], images[1]], dim=0)
        images = images.cuda()
        labels = labels.cuda()
        bsz = labels.shape[0]

        # compute loss
        features = model(images)
        f1, f2 = torch.split(features, [bsz, bsz], dim=0)
        features = torch.cat([f1.unsqueeze(1), f2.unsqueeze(1)], dim=1) # [bz, 2, feature_dim]

        if cluster_result is None:
            if opt.method in ['SupCon', 'CE']:
                loss = criterion(features, labels)
            elif opt.method == 'SimCLR':
                loss = criterion(features)
            else:
                raise ValueError('contrastive method not supported: {}'.
                                format(opt.method))
        else:
            # switch to clustering labels
            assert len(cluster_result['im2cluster']) == 1, f"cluster_result['im2cluster'] has {len(cluster_result['im2cluster'])} items, but 1 is required."
            cluster_i = cluster_result['im2cluster'][0]
            cluster_labels = cluster_i[index]
            loss = criterion(features, cluster_labels) # no actual weak label involved

        # update metric
        losses.update(loss.item(), bsz)

        # SGD
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()
        if idx % opt.print_freq == 0:
            progress.display(idx)

    return losses.avg

# ============== clustering ==================
def compute_features(eval_loader, model, args):
    print('Computing features...')
    model.eval()
    # free GPU for that batch
    torch.cuda.empty_cache()
    features = torch.zeros(len(eval_loader.dataset),args.low_dim).cuda()
    for i, (images, labels, index) in enumerate(tqdm(eval_loader)):
        with torch.no_grad():
            images = images.cuda(non_blocking=True)
            feat = model(images)
            features[index] = feat
    return features.cpu()

def perform_clustering(args, epoch, eval_loader, model):

    if epoch>=args.warmup_epoch:
        if (epoch+1) % args.perform_cluster_epoch == 0:
            # compute momentum features for center-cropped images
            features = compute_features(eval_loader, model, args)

            # placeholder for clustering result
            cluster_result = {'im2cluster':[],'centroids':[],'density':[]}
            for num_cluster in args.num_cluster:
                cluster_result['im2cluster'].append(torch.zeros(len(eval_loader.dataset),dtype=torch.long).cuda())
                cluster_result['centroids'].append(torch.zeros(int(num_cluster),args.low_dim).cuda())
                cluster_result['density'].append(torch.zeros(int(num_cluster)).cuda())

            features[torch.norm(features,dim=1)>1.5] /= 2 #account for the few samples that are computed twice
            features = features.numpy()
            cluster_result = run_kmeans(features,args)  #run kmeans clustering on master node
            # save the clustering result
            if (epoch+1) % args.save_freq == 0:
                print("\nSaving cluster results...\n")
                torch.save(cluster_result, os.path.join(args.save_folder, 'clusters_%d'%epoch))

            return cluster_result

def run_kmeans(x, args):
    """
    Args:
        x: data to be clustered
    """

    print('performing kmeans clustering')
    results = {'im2cluster':[],'centroids':[],'density':[]}

    for seed, num_cluster in enumerate(args.num_cluster):
        # intialize faiss clustering parameters
        d = x.shape[1]
        k = int(num_cluster)
        clus = faiss.Clustering(d, k)
        clus.verbose = True
        clus.niter = 20
        clus.nredo = 5
        clus.seed = seed
        clus.max_points_per_centroid = 1000
        clus.min_points_per_centroid = 10

        res = faiss.StandardGpuResources()
        cfg = faiss.GpuIndexFlatConfig()
        cfg.useFloat16 = False
        cfg.device = 0
        index = faiss.GpuIndexFlatL2(res, d, cfg)

        clus.train(x, index)

        D, I = index.search(x, 1) # for each sample, find cluster distance and assignments
        im2cluster = [int(n[0]) for n in I]

        # get cluster centroids
        centroids = faiss.vector_to_array(clus.centroids).reshape(k,d)

        # sample-to-centroid distances for each cluster
        Dcluster = [[] for c in range(k)]
        for im,i in enumerate(im2cluster):
            Dcluster[i].append(D[im][0])

        # concentration estimation (phi)
        density = np.zeros(k)
        for i,dist in enumerate(Dcluster):
            if len(dist)>1:
                d = (np.asarray(dist)**0.5).mean()/np.log(len(dist)+10)
                density[i] = d

        #if cluster only has one point, use the max to estimate its concentration
        dmax = density.max()
        for i,dist in enumerate(Dcluster):
            if len(dist)<=1:
                density[i] = dmax

        density = density.clip(np.percentile(density,10),np.percentile(density,90)) #clamp extreme values for stability
        density = args.temp*density/density.mean()  #scale the mean to temperature

        # convert to cuda Tensors for broadcast
        centroids = torch.Tensor(centroids).cuda()
        centroids = nn.functional.normalize(centroids, p=2, dim=1)

        im2cluster = torch.LongTensor(im2cluster).cuda()
        density = torch.Tensor(density).cuda()

        results['centroids'].append(centroids)
        results['density'].append(density)
        results['im2cluster'].append(im2cluster)

    return results

def main(opt):

    # tensorboard and logger
    tf_logger = tb_logger.Logger(logdir=opt.tb_folder, flush_secs=2)
    scalar_logger = txt_logger(opt.save_folder, opt, 'python ' + ' '.join(sys.argv))

    # build data loader
    train_loader,eval_loader = set_loader(opt)

    # build model and criterion
    model, criterion = set_model(opt, scalar_logger)

    # build optimizer
    optimizer = set_optimizer(opt, model)

    # training routine
    # resume model path
    if opt.resume_model_path:
        start = opt.pre_ssl_epoch
    else:
        start = 0

    cluster_result = None
    for epoch in range(start + 1, opt.epochs + 1):
        # clustering
        c_result = perform_clustering(opt, epoch, eval_loader, model)
        if c_result:
            cluster_result = c_result

        # adjust lr
        adjust_learning_rate(opt, optimizer, epoch)

        # train for one epoch
        start_time_epoch = time.time()
        loss = train(train_loader, model, criterion, optimizer, epoch, opt, cluster_result=cluster_result)
        end_time_epoch = time.time()

        # tensorboard logger
        tf_logger.log_value('loss', loss, epoch)
        tf_logger.log_value('learning_rate', optimizer.param_groups[0]['lr'], epoch)

        # file logger
        scalar_logger.log_value(epoch, ('loss', loss),
                                    ('learning_rate', optimizer.param_groups[0]['lr']),
                                    ('time: ', "{} s".format(end_time_epoch - start_time_epoch)))

        if epoch % opt.save_freq == 0:
            save_file = os.path.join(
                opt.save_folder, 'ckpt_epoch_{epoch}.pth'.format(epoch=epoch))
            save_model(model, optimizer, opt, epoch, save_file)

    # save the last model
    if epoch % opt.save_freq != 0:
        save_file = os.path.join(
            opt.save_folder, 'ckpt_epoch_{epoch}.pth'.format(epoch=epoch))
        save_model(model, optimizer, opt, opt.epochs, save_file)

    return opt.save_folder

@suppress_std
def suppress_std_main(opt):
    save_folder = main(opt)
    return save_folder

if __name__ == '__main__':
    """
    Command:
    - Super Con:
        - utzap50k:
        $ python main_supcon.py --batch_size 128 --utzap_gran_lvl 3 --learning_rate 0.5 --temp 0.1 --cosine --save_freq 10
        - deepfashion
        $ python main_supcon.py --batch_size 128 --deepfashion_gran_lvl 0 --dataset deepfashion --data_root_name img --learning_rate 0.5 --temp 0.1 --cosine --data_folder /home/tianqinl/Code/WeakSupervisionSSL/data_unzip/deepfashion
        - CUB (attribute loss)
        $ python main_supcon.py --batch_size 128 --dataset CUB --data_root_name images --data_folder ../data_unzip/CUB_200_2011 --method WeakSupCon --learning_rate 0.5 --temp 0.1
        - CUB (original resnet50)
        $ python main_supcon.py --batch_size 20 --dataset CUB --data_root_name images --data_folder ../data_unzip/CUB_200_2011 --method WeakSupCon --learning_rate 0.5 --temp 0.1 --img_size 224 --model resnet50_original
    - SimCLR:
        - utzap50k:
        $ python main_supcon.py --batch_size 128 --method SimCLR --learning_rate 0.5 --temp 0.5 --cosine
        - deepfashion
        $ python main_supcon.py --batch_size 128 --method SimCLR --dataset deepfashion --data_root_name img --learning_rate 0.5 --temp 0.5 --cosine --data_folder /home/tianqinl/Code/WeakSupervisionSSL/data_unzip/deepfashion

    - Adding loss of weak label
        - utzap50k:
        $ python main_supcon.py --meta_file_train meta_data_lossadd_train.csv --method WeakSupCon --learning_rate 0.5 --temp 0.1 --cosine --batch_size 128 --used_attributes 0 1 2 3 4 5 6 7 8 9 --utzap_gran_lvl 10 --data_folder /projects/rsalakhugroup/tianqinl/ut-zap50K-processed

    """
    # get options
    opt = parse_option()

    # decide suppress std or not
    if opt.pipeline:
        save_folder = suppress_std_main(opt)
        #save_folder = main(opt)
        sys.stdout.write(save_folder + "\n")
    else:
        save_folder = main(opt)
        print("ckpt#{}".format(save_folder))