import subprocess
from argparse import ArgumentParser
import shlex
import sys
import os
from cyvcf2 import VCF
import matplotlib.pyplot as plt
from matplotlib_venn import venn2
import tempfile
import numpy as np

def main():
    args = get_args()
    if args.ignore_type:
        ignore_type = "-1"
    else:
        ignore_type = "1"
    combined_vcf = survivor(samples=[normalize_vcf(s) for s in [args.truth, args.test]],
                            distance=args.distance,
                            ignore_type=ignore_type,
                            minlength=args.minlength)
    truth_set, test_set = make_venn(combined_vcf)
    tp = len(truth_set & test_set)
    precision = tp / len(test_set)
    recall = tp / len(truth_set)
    fmeasure = 2*(precision*recall)/(precision + recall)
    print(f"Precision: {round(precision, ndigits=4)}%")
    print(f"Recall: {round(recall, ndigits=4)}%")
    print(f"F-measure: {round(fmeasure, ndigits=4)}")
    if args.bar:
        bar_chart(combined_vcf)

def normalize_vcf(vcff):
    handle, name = tempfile.mkstemp()
    out = open(name, 'w')
    for line in open(vcff):
        out.write(line.replace('DUP', 'INS'))
    os.close(handle)
    return name

def survivor(samples, distance, ignore_type, minlength):
    """
    Executes SURVIVOR merge, with parameters:
    -samples.fofn (truth and test)
    -distance between calls (args.distance)
    -number of callers to support call (1)
    -require variants to have sampe type (args.ignore_type)
    -require variants to be on same strand (no)
    -estimate distance between calls (no)
    -specify minimal size of SV event (args.minlength)
    """
    fhf, fofn_f = tempfile.mkstemp()
    fhv, vcf_out = tempfile.mkstemp()
    with open(fofn_f, 'w') as fofn:
        for s in samples:
            fofn.write(s + "\n")
    survivor_cmd = f"SURVIVOR merge {fofn_f} {distance} 1 {ignore_type} -1 -1 {minlength} {vcf_out}"
    sys.stderr.write("Executing SURVIVOR...\n")
    subprocess.call(shlex.split(survivor_cmd), stdout=subprocess.DEVNULL)
    os.close(fhf)
    os.close(fhv)
    return vcf_out

def is_variant(call):
    """Check if a variant position qualifies as a variant

    0,1,2,3==HOM_REF, HET, UNKNOWN, HOM_ALT"""
    if call == 1 or call == 3:
        return True
    else:
        return False

def make_venn(vcf, outname="venn.png"):
    positions = [[], []]
    for v in VCF(vcf):
        for index, call in enumerate(v.gt_types):
            if is_variant(call):
                positions[index].append("{}:{}-{}".format(v.CHROM, v.start, v.INFO.get('SVTYPE')))
    identifier_sets = [set(i) for i in positions]
    venn2(identifier_sets, set_labels=('Truth', 'Test'))
    plt.savefig(outname)
    plt.close()
    return identifier_sets

def bar_chart(vcf, outname="stacked_bar.png"):
    """
    Make a stacked bar chart for length of the SV split by validation status
    This ignores zygosity.
    """
    len_dict = {"True": [], "False": [], "Missed": []}
    for v in VCF(vcf):
        if not v.INFO.get('SVTYPE') == 'TRA' and abs(v.INFO.get('AVGLEN')) >= 50:
            calls = [is_variant(call) for call in v.gt_types]
            if calls == [True, True]:
                len_dict['True'].append(v.INFO.get('AVGLEN'))
            elif calls == [False, True]:
                len_dict['False'].append(v.INFO.get('AVGLEN'))
            elif calls == [True, False]:
                len_dict['Missed'].append(v.INFO.get('AVGLEN'))
    plt.subplot(2, 1, 1)
    plt.hist(x=np.array(list(len_dict.values())),
             bins=[i for i in range(0, 2000, 10)],
             stacked=True,
             histtype='bar',
             label=list(len_dict.keys()))
    plt.xlabel('Length of structural variant')
    plt.ylabel('Number of variants')
    plt.legend(frameon=False,
               fontsize="small")
    plt.subplot(2, 1, 2)
    plt.hist(x=np.array(list(len_dict.values())),
             bins=[i for i in range(0, 20000, 100)],
             stacked=True,
             histtype='bar',
             label=list(len_dict.keys()),
             log=True)
    plt.xlabel('Length of structural variant')
    plt.ylabel('Number of variants')
    plt.legend(frameon=False,
               fontsize="small")
    plt.tight_layout()
    plt.savefig(outname)
    plt.close()

def get_args():
    parser = ArgumentParser(description="Calculate precision-recall metrics from 2 vcf files")
    parser.add_argument("--truth", help="vcf containing truth set", required=True)
    parser.add_argument("--test", help="vcf containing test set", required=True)
    parser.add_argument("-d", "--distance",
                        help="maximum distance between test and truth call",
                        default=500)
    parser.add_argument("--minlength",
                        help="Minimum length of SVs to be taken into account",
                        default=50)
    parser.add_argument("-i", "--ignore_type",
                        help="Ignore the type of the structural variant",
                        action="store_true")
    parser.add_argument("--bar",
                        help="Make stacked bar chart of SV lengths coloured by validation status",
                        action="store_true")
    return parser.parse_args()

if __name__ == '__main__':
    main()


import sys
import numpy as np
import pandas as pd
import stock.utils.symbol_util
from stock.globalvar import *
from stock.marketdata.storefactory import get_store
from config import store_type
from pandas.tseries.offsets import BDay
import tushare as ts
from stock.utils.calc_price import get_zt_price

def get_last_trading_date(today):
    yest = today - BDay(1)
    folder = TICK_DIR["daily"]
    while True:
        yest_str = yest.strftime("%Y-%m-%d")
        filepath = os.path.join(folder, yest_str + ".csv")
        if os.path.isfile(filepath):
            break
        yest = yest - BDay(1)
    return yest

def get_industry():
    df_industry = stock.utils.symbol_util.load_industry()
    df_res = df_industry.groupby("exsymbol")["industry"].agg(industry=lambda x: ",".join(x))
    return df_res

def get_concept():
    df = stock.utils.symbol_util.load_concept()
    df_res = df.groupby("exsymbol")["concept"].agg(concept=lambda x: ",".join(x))
    return df_res

def filter_by_history(date, exsymbols):
    store = get_store(store_type)
    result = []
    df = pd.DataFrame(columns=["increase60", "increase5", "future"])
    for exsymbol in exsymbols:
        if not store.has(exsymbol):
            continue
        df_stock = store.get(exsymbol)
        df_past = df_stock.loc[:date].copy()
        if len(df_past) == 0:
            continue
        close_min60 = np.min(df_past.iloc[-60:].close)
        close_min5 = np.min(df_past.iloc[-5:].close)
        increase60 = df_past.iloc[-1].close/close_min60-1
        increase5 = df_past.iloc[-1].close/close_min5-1
        future = df_stock.iloc[-1].close/df_past.iloc[-1].close-1
        df.loc[exsymbol] = [increase60, increase5, future]
    return df

def get_lianban(exsymbol, date):
    store = get_store(store_type)
    lianban = -1
    xingu = None
    try:
        df_stock = store.get(exsymbol)
        idx = df_stock.index.get_loc(date)
        df_stock = df_stock.iloc[:idx+1]
        df_stock.loc[:, "zt_price"] = df_stock.close.shift(1).apply(lambda x: get_zt_price(exsymbol[2:], x))
        df_stock.loc[:, "is_zhangting"] = np.absolute(df_stock["zt_price"]-df_stock["close"])<1e-8
        df_nozt = df_stock[df_stock.is_zhangting==False]
        lianban = 0
        if len(df_nozt) == 0:
            lianban = len(df_stock)
        else:
            idx_start = df_stock.index.get_loc(df_nozt.index[-1])
            idx_end = df_stock.index.get_loc(df_stock.index[-1])
            lianban = idx_end - idx_start
        xingu = lianban == len(df_stock)-1
    except Exception as e:
        print(str(e))
        pass

    return (lianban, xingu)

def get_zhangting(today):
    today_str = today.strftime("%Y-%m-%d")
    df_today = stock.utils.symbol_util.get_realtime_by_date(today_str)
    df_today["zt_price"] = df_today.apply(lambda x: get_zt_price(x.name[2:], x["yest_close"]), axis=1)
    df_today.loc[:, "is_zhangting"] = np.absolute(df_today["zt_price"]-df_today["close"])<1e-8
    df_zt = df_today[(df_today.is_zhangting==True) & (df_today.lt_mcap>0) & (df_today.volume>0)].copy()
    df_zt.loc[:, "turnover"] = df_zt["volume"]/(df_zt["lt_mcap"]/df_zt["close"]*1e6)
    df_zt.loc[:, "fengdan"] = df_zt["b1_v"] * df_zt["b1_p"] *100 / df_zt["lt_mcap"] / 1e8
    df_zt.loc[:, "fengdan_money"] = df_zt["b1_v"]*df_zt["b1_p"]/1e6
    for exsymbol in df_zt.index.tolist():
        (lianban, xingu) = get_lianban(exsymbol, today_str)
        df_zt.loc[exsymbol, "lianban"] = lianban
        df_zt.loc[exsymbol, "xingu"] = xingu

    df_zt = df_zt[df_zt.xingu==False]
    df_tick = stock.utils.symbol_util.get_tick_by_date(today_str)
    df_res = df_zt.merge(df_tick[["zhangting_sell", "zhangting_min", "cost"]], how="inner", left_index=True, right_index=True)
    df_res["costoverflow"] = df_res["cost"] / df_res["close"]
    df_industry = get_industry()
    df_concept = get_concept()
    df_res = df_res.merge(df_industry, how="left", left_index=True, right_index=True)
    columns = ["fengdan", "fengdan_money", "lt_mcap", "zhangting_sell", "lianban", "industry"]
    print("========================== zhangting ==========================")
    print(df_res[columns].sort_values(["lianban"], ascending=False))

def get_duanban(today):
    today_str = today.strftime("%Y-%m-%d")
    yest = get_last_trading_date(today)
    yest_str = yest.strftime("%Y-%m-%d")
    df_yest = stock.utils.symbol_util.get_realtime_by_date(yest_str)
    df_yest["zt_price"] = df_yest.apply(lambda x: get_zt_price(x.name[2:], x["yest_close"]), axis=1)
    df_yest.loc[:, "is_zhangting"] = np.absolute(df_yest["zt_price"]-df_yest["close"])<1e-8
    df_yest_zt = df_yest[(df_yest.is_zhangting==True) & (df_yest.lt_mcap>0) & (df_yest.volume>0)].copy()
    for exsymbol in df_yest_zt.index.tolist():
        (lianban, xingu) = get_lianban(exsymbol, yest_str)
        df_yest_zt.loc[exsymbol, "lianban"] = lianban
        df_yest_zt.loc[exsymbol, "xingu"] = xingu

    df_yest_zt = df_yest_zt[df_yest_zt.xingu==False]
    df_yest_zt.loc[:, "fengdan"] = df_yest_zt["b1_v"] * df_yest_zt["b1_p"] *100 / df_yest_zt["lt_mcap"] / 1e8
    df_yest_zt.loc[:, "fengdan_money"] = df_yest_zt["b1_v"]*df_yest_zt["b1_p"]/1e6

    df_today = stock.utils.symbol_util.get_realtime_by_date(today_str)
    df_today["zt_price"] = df_today.apply(lambda x: get_zt_price(x.name[2:], x["yest_close"]), axis=1)
    df_today.loc[:, "is_zhangting"] = np.absolute(df_today["zt_price"]-df_today["close"])<1e-8
    df_today_nozt = df_today[(df_today.is_zhangting==False) & (df_today.lt_mcap>0) & (df_today.volume>0)].copy()

    df_res = df_yest_zt.merge(df_today_nozt[["chgperc"]], how="inner", left_index=True, right_index=True)

    df_industry = get_industry()
    df_res = df_res.merge(df_industry, how="left", left_index=True, right_index=True)
    columns = ["fengdan", "fengdan_money", "chgperc_y", "lt_mcap", "lianban", "industry"]
    print("========================== duanban ==========================")
    print(df_res[columns].sort_values(["lianban"], ascending=False))

if __name__ == "__main__":
    pd.set_option('display.max_rows', None)
    pd.set_option('display.max_columns', None)
    today = None
    if len(sys.argv) == 1:
        today = pd.datetime.today()
    else:
        today = pd.datetime.strptime(sys.argv[1], "%Y-%m-%d")

    get_zhangting(today)
    get_duanban(today)


from typing import Any, Dict, List, Optional, Sequence, Tuple

import numpy as np

from ..argument_utility import (
    ActionScalerArg,
    EncoderArg,
    QFuncArg,
    RewardScalerArg,
    ScalerArg,
    UseGPUArg,
    check_encoder,
    check_q_func,
    check_use_gpu,
)
from ..constants import IMPL_NOT_INITIALIZED_ERROR, ActionSpace
from ..dataset import Transition, TransitionMiniBatch
from ..dynamics import DynamicsBase
from ..gpu import Device
from ..models.encoders import EncoderFactory
from ..models.optimizers import AdamFactory, OptimizerFactory
from ..models.q_functions import QFunctionFactory
from .base import AlgoBase
from .torch.sac_impl import SACImpl
from .utility import ModelBaseMixin

class MOPO(ModelBaseMixin, AlgoBase):
    r"""Model-based Offline Policy Optimization.

    MOPO is a model-based RL approach for offline policy optimization.
    MOPO leverages the probablistic ensemble dynamics model to generate
    new dynamics data with uncertainty penalties.
    The ensemble dynamics model consists of :math:`N` probablistic models
    :math:`\{T_{\theta_i}\}_{i=1}^N`.
    At each epoch, new transitions are generated via randomly picked dynamics
    model :math:`T_\theta`.

    .. math::
        s_{t+1}, r_{t+1} \sim T_\theta(s_t, a_t)

    where :math:`s_t \sim D` for the first step, otherwise :math:`s_t` is the
    previous generated observation, and :math:`a_t \sim \pi(\cdot|s_t)`.
    The generated :math:`r_{t+1}` would be far from the ground truth if the
    actions sampled from the policy function is out-of-distribution.
    Thus, the uncertainty penalty reguralizes this bias.

    .. math::
        \tilde{r_{t+1}} = r_{t+1} - \lambda \max_{i=1}^N
            || \Sigma_i (s_t, a_t) ||

    where :math:`\Sigma(s_t, a_t)` is the estimated variance.
    Finally, the generated transitions
    :math:`(s_t, a_t, \tilde{r_{t+1}}, s_{t+1})` are appended to dataset
    :math:`D`.
    This generation process starts with randomly sampled
    ``n_initial_transitions`` transitions till ``horizon`` steps.

    Note:
        Currently, MOPO only supports vector observations.

    References:
        * `Yu et al., MOPO: Model-based Offline Policy Optimization.
          <https://arxiv.org/abs/2005.13239>`_

    Args:
        actor_learning_rate (float): learning rate for policy function.
        critic_learning_rate (float): learning rate for Q functions.
        temp_learning_rate (float): learning rate for temperature parameter.
        actor_optim_factory (d3rlpy.models.optimizers.OptimizerFactory):
            optimizer factory for the actor.
        critic_optim_factory (d3rlpy.models.optimizers.OptimizerFactory):
            optimizer factory for the critic.
        temp_optim_factory (d3rlpy.models.optimizers.OptimizerFactory):
            optimizer factory for the temperature.
        actor_encoder_factory (d3rlpy.models.encoders.EncoderFactory or str):
            encoder factory for the actor.
        critic_encoder_factory (d3rlpy.models.encoders.EncoderFactory or str):
            encoder factory for the critic.
        q_func_factory (d3rlpy.models.q_functions.QFunctionFactory or str):
            Q function factory.
        batch_size (int): mini-batch size.
        n_frames (int): the number of frames to stack for image observation.
        n_steps (int): N-step TD calculation.
        gamma (float): discount factor.
        tau (float): target network synchronization coefficiency.
        n_critics (int): the number of Q functions for ensemble.
        target_reduction_type (str): ensemble reduction method at target value
            estimation. The available options are
            ``['min', 'max', 'mean', 'mix', 'none']``.
        update_actor_interval (int): interval to update policy function.
        initial_temperature (float): initial temperature value.
        dynamics (d3rlpy.dynamics.DynamicsBase): dynamics object.
        rollout_interval (int): the number of steps before rollout.
        rollout_horizon (int): the rollout step length.
        rollout_batch_size (int): the number of initial transitions for
            rollout.
        lam (float): :math:`\lambda` for uncertainty penalties.
        real_ratio (float): the real of dataset samples in a mini-batch.
        generated_maxlen (int): the maximum number of generated samples.
        use_gpu (bool, int or d3rlpy.gpu.Device):
            flag to use GPU, device ID or device.
        scaler (d3rlpy.preprocessing.Scaler or str): preprocessor.
            The available options are `['pixel', 'min_max', 'standard']`.
        action_scaler (d3rlpy.preprocessing.ActionScaler or str):
            action preprocessor. The available options are ``['min_max']``.
        reward_scaler (d3rlpy.preprocessing.RewardScaler or str):
            reward preprocessor. The available options are
            ``['clip', 'min_max', 'standard']``.
        impl (d3rlpy.algos.torch.sac_impl.SACImpl): algorithm implementation.

    """

    _actor_learning_rate: float
    _critic_learning_rate: float
    _temp_learning_rate: float
    _actor_optim_factory: OptimizerFactory
    _critic_optim_factory: OptimizerFactory
    _temp_optim_factory: OptimizerFactory
    _actor_encoder_factory: EncoderFactory
    _critic_encoder_factory: EncoderFactory
    _q_func_factory: QFunctionFactory
    _tau: float
    _n_critics: int
    _target_reduction_type: str
    _update_actor_interval: int
    _initial_temperature: float
    _dynamics: Optional[DynamicsBase]
    _rollout_interval: int
    _rollout_horizon: int
    _rollout_batch_size: int
    _lam: float
    _use_gpu: Optional[Device]
    _impl: Optional[SACImpl]

    def __init__(
        self,
        *,
        actor_learning_rate: float = 3e-4,
        critic_learning_rate: float = 3e-4,
        temp_learning_rate: float = 3e-4,
        actor_optim_factory: OptimizerFactory = AdamFactory(),
        critic_optim_factory: OptimizerFactory = AdamFactory(),
        temp_optim_factory: OptimizerFactory = AdamFactory(),
        actor_encoder_factory: EncoderArg = "default",
        critic_encoder_factory: EncoderArg = "default",
        q_func_factory: QFuncArg = "mean",
        batch_size: int = 100,
        n_frames: int = 1,
        n_steps: int = 1,
        gamma: float = 0.99,
        tau: float = 0.005,
        n_critics: int = 2,
        target_reduction_type: str = "min",
        update_actor_interval: int = 1,
        initial_temperature: float = 1.0,
        dynamics: Optional[DynamicsBase] = None,
        rollout_interval: int = 1000,
        rollout_horizon: int = 5,
        rollout_batch_size: int = 50000,
        lam: float = 1.0,
        real_ratio: float = 0.05,
        generated_maxlen: int = 50000 * 5 * 5,
        use_gpu: UseGPUArg = False,
        scaler: ScalerArg = None,
        action_scaler: ActionScalerArg = None,
        reward_scaler: RewardScalerArg = None,
        impl: Optional[SACImpl] = None,
        **kwargs: Any
    ):
        super().__init__(
            batch_size=batch_size,
            n_frames=n_frames,
            n_steps=n_steps,
            gamma=gamma,
            scaler=scaler,
            action_scaler=action_scaler,
            reward_scaler=reward_scaler,
            real_ratio=real_ratio,
            generated_maxlen=generated_maxlen,
            kwargs=kwargs,
        )
        self._actor_learning_rate = actor_learning_rate
        self._critic_learning_rate = critic_learning_rate
        self._temp_learning_rate = temp_learning_rate
        self._actor_optim_factory = actor_optim_factory
        self._critic_optim_factory = critic_optim_factory
        self._temp_optim_factory = temp_optim_factory
        self._actor_encoder_factory = check_encoder(actor_encoder_factory)
        self._critic_encoder_factory = check_encoder(critic_encoder_factory)
        self._q_func_factory = check_q_func(q_func_factory)
        self._tau = tau
        self._n_critics = n_critics
        self._target_reduction_type = target_reduction_type
        self._update_actor_interval = update_actor_interval
        self._initial_temperature = initial_temperature
        self._dynamics = dynamics
        self._rollout_interval = rollout_interval
        self._rollout_horizon = rollout_horizon
        self._rollout_batch_size = rollout_batch_size
        self._lam = lam
        self._use_gpu = check_use_gpu(use_gpu)
        self._impl = impl

    def _create_impl(
        self, observation_shape: Sequence[int], action_size: int
    ) -> None:
        self._impl = SACImpl(
            observation_shape=observation_shape,
            action_size=action_size,
            actor_learning_rate=self._actor_learning_rate,
            critic_learning_rate=self._critic_learning_rate,
            temp_learning_rate=self._temp_learning_rate,
            actor_optim_factory=self._actor_optim_factory,
            critic_optim_factory=self._critic_optim_factory,
            temp_optim_factory=self._temp_optim_factory,
            actor_encoder_factory=self._actor_encoder_factory,
            critic_encoder_factory=self._critic_encoder_factory,
            q_func_factory=self._q_func_factory,
            gamma=self._gamma,
            tau=self._tau,
            n_critics=self._n_critics,
            target_reduction_type=self._target_reduction_type,
            initial_temperature=self._initial_temperature,
            use_gpu=self._use_gpu,
            scaler=self._scaler,
            action_scaler=self._action_scaler,
            reward_scaler=self._reward_scaler,
        )
        self._impl.build()

    def _update(self, batch: TransitionMiniBatch) -> Dict[str, float]:
        assert self._impl is not None, IMPL_NOT_INITIALIZED_ERROR
        metrics = {}

        critic_loss = self._impl.update_critic(batch)
        metrics.update({"critic_loss": critic_loss})

        # delayed policy update
        if self._grad_step % self._update_actor_interval == 0:
            actor_loss = self._impl.update_actor(batch)
            metrics.update({"actor_loss": actor_loss})

            # lagrangian parameter update for SAC temperature
            if self._temp_learning_rate > 0:
                temp_loss, temp = self._impl.update_temp(batch)
                metrics.update({"temp_loss": temp_loss, "temp": temp})

            self._impl.update_critic_target()
            self._impl.update_actor_target()

        return metrics

    def get_action_type(self) -> ActionSpace:
        return ActionSpace.CONTINUOUS

    def _is_generating_new_data(self) -> bool:
        return self._grad_step % self._rollout_interval == 0

    def _sample_initial_transitions(
        self, transitions: List[Transition]
    ) -> List[Transition]:
        # uniformly sample transitions
        n_transitions = self._rollout_batch_size
        indices = np.random.randint(len(transitions), size=n_transitions)
        return [transitions[i] for i in indices]

    def _get_rollout_horizon(self) -> int:
        return self._rollout_horizon

    def _mutate_transition(
        self,
        observations: np.ndarray,
        rewards: np.ndarray,
        variances: np.ndarray,
    ) -> Tuple[np.ndarray, np.ndarray]:
        # regularize by uncertainty
        rewards -= self._lam * variances
        return observations, rewards


import numpy as np

class Chromosome:
    def __init__(self, parameter_dict, t_init, v_init, t_end_blocked=0, active_rep=1, length=0, sites=0, blocked=True, rida=2):
        self.parameter_dict = parameter_dict
        # changing chromosome variables
        self.t_init = t_init
        self.v_init = v_init
        self.division_time_RDA = t_init + self.calculate_next_division_time_RDA(self.v_init)
        print('initiated, division time set:', self.division_time_RDA, ' current volume:', self.v_init)
        self.t_end_blocked = t_end_blocked
        self.time = np.array([self.t_init])
        self.length = np.array([length])
        if self.parameter_dict.homogeneous_dist_sites == 0:
            self.sites = np.array([self.parameter_dict.n_c_max_0])  # if distribution is inhomogeneous, all sites are created instantaneously after the origin has been replicated
        else:
            self.sites = np.array([sites])  # if distribution is homogeneous, the number of sites depends on how the chromosome is initiated, either fully replicated or being actively replicated
        self.active_replication = active_rep
        self.division_time = self.parameter_dict.t_max
        if self.parameter_dict.period_blocked > 0:
            self.blocked = blocked
        else:
            self.blocked = False

        # changing ultrasensitivity parameters
        self.n_datA = 1
        self.n_rida = rida
        if self.active_replication ==0:
            self.n_dars2 = 1
            self.n_dars1 = 1
        else:
            self.n_dars2 = 0
            self.n_dars1 = 0
        self.n_dars2_high_activity = 0
        self.n_datA_high_activity = 0

    def replicate(self, time):
        if self.active_replication == 1:
            if self.length[-1] < 1:
                # print(' length before was ', self.length[-1])
                new_length = self.length[-1] + self.parameter_dict.rate_rep * self.parameter_dict.time_step
                self.length = np.append(self.length, new_length)
                if self.parameter_dict.homogeneous_dist_sites == 1:
                    new_sites = self.sites[-1] + self.parameter_dict.rate_synth_sites * self.parameter_dict.time_step
                    self.sites = np.append(self.sites, new_sites)
                    # print('replicated new sites',  self.rate_synth_sites * self.dt)
                else:
                    self.sites = np.append(self.sites, self.parameter_dict.n_c_max_0)
                self.time = np.append(self.time, time)
                # print('new sites', self.rate_synth_sites* self.dt)
                # print(' length after replication was ', self.length[-1])
                if self.length[-1] > 1 or self.sites[-1] > self.parameter_dict.n_c_max_0:
                    self.length[-1] = 1
                    self.sites[-1] = self.parameter_dict.n_c_max_0
                # Check whether dars1 and dars2 are being doubled
                if self.length[-1] >= self.parameter_dict.site_dars2: #as soon as we reach position of site of dars2, the number of dars2 sites goes from 0 to 1
                    self.n_dars2 = 1
                if self.length[-1] >= self.parameter_dict.site_dars1: #as soon as we reach position of site of dars2, the number of dars2 sites goes from 0 to 1
                    self.n_dars1 = 1
                # Check whether datA and dars2 are in high or low activity state
                if self.length[-1] >= self.parameter_dict.relative_chromosome_position_onset_dars2: #as soon as we reach position where dars2 becomes more active, the number of high activity dars2 goes up to two
                    self.n_dars2_high_activity = 2
                if self.length[-1] >= self.parameter_dict.relative_chromosome_position_onset_datA: #as soon as we reach position where datA becomes more active, the number of high activity datA goes up to two
                    self.n_datA_high_activity = 2
                if self.length[-1] >= self.parameter_dict.relative_chromosome_position_offset_dars2: #as soon as we reach position of site of dars2, the number of dars2 sites goes from 0 to 1
                    # print('dars was doubled')
                    self.n_dars2_high_activity = 0
                if self.length[-1] >= self.parameter_dict.relative_chromosome_position_offset_datA: #as soon as we reach position of site of dars2, the number of dars2 sites goes from 0 to 1
                    # print('dars was doubled')
                    self.n_datA_high_activity = 0

            else:
                # replication is not active anymore and therefore rida is switched off
                self.active_replication = 0
                self.n_rida = 0
                if self.parameter_dict.version_of_coupled_division_regulation == 'RDA':
                    print('replication finished, set now division time')
                    print('Division time RDA: ',self.division_time_RDA, 'time at the moment:', time)
                    if time < self.division_time_RDA:
                        self.division_time = self.division_time_RDA
                    else:
                        self.division_time = time
                        print('divide now')
                else:
                    self.division_time = time + self.parameter_dict.t_D
                print('replication was finished, division time was set')

    #  checks whether blocked, if condition correct unblocks, returns 0 if not and 1 if it is blocked
    def check_blocked(self, time):
        if time >= self.t_end_blocked:
            self.blocked = False
            # print('unblocked')
            return self.blocked
        else:
            self.blocked = True
            return self.blocked

    def store_everything(self):
        return self.time, self.length, self.sites

    def set_t_end_blocked(self, t_end_blocked):
        self.t_end_blocked = t_end_blocked

    def calculate_next_division_time_RDA(self, volume):
        average_added_volume = self.parameter_dict.v_init_th * (np.exp(self.parameter_dict.rate_growth * self.parameter_dict.t_CD)-1)
        noisy_added_volume = average_added_volume + np.random.normal(0, self.parameter_dict.cv_added_volume) * average_added_volume
        noisy_division_volume = noisy_added_volume + volume
        return np.log(noisy_division_volume/volume)/self.parameter_dict.rate_growth


from xgboost import XGBClassifier
import flask
import config
import time
import pickle
import numpy as np
from flask import Flask, request
from utils import extract_features
import sys

#app = Flask(__name__)

MODEL = None

def audio_prediction(audio_file):
    feats = extract_features(
        audio_file, mel=True, mfcc=True, chroma=True, contrast=True)
    scaler = pickle.load(open(config.SCALAR_PATH, 'rb'))
    X = scaler.transform(feats.reshape(1, -1))
    pred = MODEL.predict_proba(X)
    return pred[0][1]

"""
@app.route("/predict")
def predict():
    audio_file = request.args.get("audio_file")
    start_time = time.time()
    male_prediction = audio_prediction(audio_file)
    female_prediction = 1 - male_prediction
    response = {}
    response["response"] = {
        "male": str(male_prediction),
        "female": str(female_prediction),
        "audio_file": str(audio_file),
        "time_taken": str(time.time() - start_time),
    }
    return flask.jsonify(response)
"""

if __name__ == "__main__":

    MODEL = XGBClassifier()
    MODEL.load_model(config.MODEL_PATH)

    file = sys.argv[1]
    p_male = audio_prediction(file)
    p_female = 1 - p_male
    print('file:', file)
    print('male:', p_male)
    print('female:', p_female)


# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import h5py
import pickle
import argparse
import numpy as np
import torch
import torch.nn as nn
import torch.utils.data
import torch.multiprocessing
from tqdm import tqdm
from easydict import EasyDict as edict
import pandas

import _init_paths
from core.config import config
from core.config import update_config, update_dir
from utils.utils import create_logger
from multiviews.pictorial import rpsm, RpsmFunc
# from multiviews.body import HumanBody
from multiviews.cameras import camera_to_world_frame
# from dataset.heatmap_dataset import HeatmapDataset, no_mix_collate_fn
import dataset
import models
from multiviews.triangulate import triangulate_poses
from core.inference import get_max_preds, get_final_preds

def parse_args():
    parser = argparse.ArgumentParser(description='Evaluate 3D Pose Estimation')
    parser.add_argument(
        '--cfg', help='configuration file name', required=True, type=str)
    args, rest = parser.parse_known_args()
    update_config(args.cfg)

    parser.add_argument(
        '--modelDir', help='model directory', type=str, default='')
    parser.add_argument('--logDir', help='log directory', type=str, default='')
    parser.add_argument(
        '--dataDir', help='data directory', type=str, default='')
    parser.add_argument(
        '--withIMU', help='use bone orientation in 3d',
        type=lambda x: (str(x).lower() in ['true', '1', 'yes']), default=False)

    args = parser.parse_args()
    update_dir(args.modelDir, args.logDir, args.dataDir)
    return args

def main():
    args = parse_args()

    logger, final_output_dir, tb_log_dir = create_logger(
        config, args.cfg, 'test3d')

    prediction_path = os.path.join(final_output_dir,
                                   config.TEST.HEATMAP_LOCATION_FILE)
    # prediction_path = os.path.join(final_output_dir, 'image_only_heatmaps.h5')
    logger.info(prediction_path)
    test_dataset = eval('dataset.' + config.DATASET.TEST_DATASET)(
        config, config.DATASET.TEST_SUBSET, False)

    if config.DATASET.TRAIN_DATASET == 'multiview_h36m':
        from multiviews.h36m_body import HumanBody
        from dataset.heatmap_dataset_h36m import HeatmapDataset, no_mix_collate_fn
    elif config.DATASET.TRAIN_DATASET == 'totalcapture':
        from multiviews.totalcapture_body import HumanBody
        from dataset.heatmap_dataset import HeatmapDataset, no_mix_collate_fn

    all_heatmaps = h5py.File(prediction_path)['heatmaps']
    all_heatmaps = all_heatmaps[()]  # load all heatmaps into ram to avoid a h5 multi-threading bug
    # with open('/data/extra/zhe/projects/multiview-pose-github/all_heatmaps.pkl', 'rb') as f:  # todo
    #     all_heatmaps = pickle.load(f)

    # pairwise_file = os.path.join(config.DATA_DIR, config.PICT_STRUCT.PAIRWISE_FILE)
    # with open(pairwise_file, 'rb') as f:
    #     pairwise = pickle.load(f)['pairwise_constrain']

    # mp = torch.multiprocessing.get_context('spawn')
    # os.environ['CUDA_VISIBLE_DEVICES'] = '1'
    # gpus = [0]  # todo write in config rather than hard code
    # do_bone_vectors = args.withIMU
    # logger.info('Whether use IMU Bone Orientation: {}'.format(str(do_bone_vectors)))
    # dev = torch.device('cuda:{}'.format(gpus[0]))

    grouping = test_dataset.grouping
    db = test_dataset.db
    mpjpes = []
    body = HumanBody()
    body_joints = []
    for j in body.skeleton:
        body_joints.append(j['name'])
    heatmap_dataset = HeatmapDataset(all_heatmaps, db, grouping, body)
    heatmap_loader = torch.utils.data.DataLoader(heatmap_dataset,
                                                 batch_size=1,
                                                 shuffle=False,
                                                 num_workers=1,
                                                 pin_memory=True,
                                                 collate_fn=no_mix_collate_fn)

    # modify df definition if add extra metrics to report
    results_df = pandas.DataFrame(columns=['imgid','subject', 'action', 'subaction', 'mpjpe'] + body_joints)
    # for i, items in tqdm(enumerate(heatmap_loader)):
    for i, items in enumerate(heatmap_loader):
        input_params_all_devices = []
        for item in items:
            # item = items[0]
            heatmaps = item['heatmaps']
            datum = item['datum']
            boxes = item['boxes']
            poses = item['poses']
            cameras = item['cameras']
            limb_length = item['limb_length']
            bone_vectors = item['bone_vectors']

        # preds = []
        # maxvs = []
        nview = heatmaps.shape[0]
        njoints = heatmaps.shape[1]
        # for idv in range(nview):  # nview
        #     hm = heatmaps[idv]
        #     center = boxes[idv]['center']
        #     scale = boxes[idv]['scale']
        #     pred, maxv = get_final_preds(config, hm, center, scale)
        #     preds.append(pred)
        #     maxvs.append(maxv)

        centers = [boxes[i]['center'] for i in range(nview)]
        scales = [boxes[i]['scale'] for i in range(nview)]
        preds, maxvs = get_final_preds(config, heatmaps, centers, scales)

        # obtain joint vis from maxvs by a threshold
        vis_thresh = 0.3
        joints_vis = np.greater(maxvs, vis_thresh)

        # if not np.all(joints_vis):  # for debug
        #     print(maxvs)

        # check if at least two views available for each joints
        valid_views = np.swapaxes(joints_vis, 0, 1).sum(axis=1).reshape(-1)
        # print(valid_views)
        if np.any(valid_views < 2):
            # print(maxvs)
            maxvs_t = np.swapaxes(maxvs, 0, 1).reshape(njoints, nview)  # (njoints, nview)
            sorted_index = np.argsort(maxvs_t, axis=1)
            top2_index = sorted_index[:, ::-1][:, :2]  # large to fewer, select top 2
            top2_vis = np.zeros((njoints, nview), dtype=np.bool)
            for j in range(njoints):
                for ind_view in top2_index[j]:
                    top2_vis[j, ind_view] = True
            top2_vis_reshape = np.transpose(top2_vis).reshape(nview, njoints, 1)
            joints_vis = np.logical_or(joints_vis, top2_vis_reshape)
            logger.info('idx_{:0>6d} sub_{} act_{} subact_{} has some joints whose valid view < 2'.format(
                datum['image_id'], datum['subject'], datum['action'], datum['subaction']))

        poses2ds = np.array(preds)
        pose3d = np.squeeze(triangulate_poses(cameras, poses2ds, joints_vis))

        # for idx_datum, prediction in enumerate(outputs_cat):
        datum = items[0]['datum']
        # gt_poses = datum['joints_gt']
        # mpjpe = np.mean(np.sqrt(np.sum((prediction - gt_poses) ** 2, axis=1)))

        metric = get_one_grouping_metric(datum, pose3d, results_df)
        mpjpe = metric['mpjpe']
        mpjpes.append(mpjpe)

        logger.info('idx_{:0>6d} sub_{} act_{} subact_{} mpjpe is {}'.format(
            datum['image_id'], datum['subject'], datum['action'], datum['subaction'], mpjpe))
            # logger.info(prediction[0], )

        # prediction = rpsm_nn(**input_params)

    logger.info('avg mpjpes on {} val samples is: {}'.format(len(grouping), np.mean(mpjpes)))
    # flag_orient = 'with' if do_bone_vectors else 'without'
    backbone_time = os.path.split(tb_log_dir)[1]
    results_df_save_path = os.path.join(final_output_dir, 'estimate3d_triangulate_{}.csv'.format(backbone_time))
    results_df.to_csv(results_df_save_path)

def get_one_grouping_metric(datum, pred, results_df):
    gt_poses = datum['joints_gt']
    # gt_poses = torch.as_tensor(gt_poses, device=torch.device('cuda'))
    # pred_cuda = torch.as_tensor(pred, device=torch.device('cuda'))
    # pjpe = torch.sqrt(torch.sum(torch.pow((pred_cuda - gt_poses), 2), dim=1))
    # mpjpe = torch.mean(pjpe)
    pjpe = np.sqrt(np.sum((pred - gt_poses) ** 2, axis=1))
    mpjpe = np.mean(pjpe)
    results_df.loc[len(results_df)] = [datum['image_id'], datum['subject'], datum['action'], datum['subaction'], mpjpe] + pjpe.tolist()
    return results_df.loc[len(results_df)-1]

if __name__ == '__main__':
    main()


import sys
import numpy as np
import tensorflow as tf
from tensorflow.python.training import training_util
from .. import evaluator, metrics
from ..configuration import *
from .doc2vec_train_doc_prediction import doc2vec_prediction_model
from .doc2vec_train_doc_prediction import DocPredictionDataset

class DocPredictionEval(evaluator.Evaluator):
    def __init__(self, dataset, log_dir=DIR_D2V_DOC_LOGDIR):
        config = tf.ConfigProto()
        config.gpu_options.allow_growth = True
        super(DocPredictionEval, self).__init__(checkpoints_dir=log_dir,
                                                output_path=os.path.join(log_dir,
                                                                         dataset.type),
                                                dataset=dataset,
                                                singular_monitored_session_config=config,
                                                infinite_loop=True)
        self.best_loss = -1

    def model(self, input_vectors, input_gene, input_variation, output_label, batch_size,
              embedding_size=EMBEDDINGS_SIZE,
              output_classes=9):
        logits, targets = doc2vec_prediction_model(input_vectors, input_gene, input_variation,
                                                   output_label, batch_size,
                                                   is_training=False, embedding_size=embedding_size,
                                                   output_classes=output_classes)

        loss = tf.nn.softmax_cross_entropy_with_logits(labels=targets, logits=logits)
        self.global_step = training_util.get_or_create_global_step()
        global_step_increase = tf.assign_add(self.global_step, 1)
        self.accumulated_loss = tf.Variable(0.0, dtype=tf.float32, name='accumulated_loss',
                                            trainable=False)
        self.accumulated_loss = tf.assign_add(self.accumulated_loss, tf.reduce_sum(loss))
        self.prediction = tf.nn.softmax(logits)
        self.metrics = metrics.single_label(self.prediction, targets, moving_average=False)
        steps = tf.cast(global_step_increase, dtype=tf.float32)
        tf.summary.scalar('loss', self.accumulated_loss / (steps * batch_size))
        return None

    def create_graph(self, dataset_tensor, batch_size):
        input_vectors, input_gene, input_variation, output_label = dataset_tensor
        self.batch_size = batch_size
        return self.model(input_vectors, input_gene, input_variation, output_label, batch_size)

    def step(self, session, graph_data, summary_op):
        self.num_steps, self.final_metrics, self.final_loss, summary = \
            session.run([self.global_step, self.metrics, self.accumulated_loss, summary_op])
        return summary

    def after_create_session(self, session, coord):
        super(DocPredictionEval, self).after_create_session(session, coord)

    def end(self, session):
        super(DocPredictionEval, self).end(session)
        cm = self.final_metrics['confusion_matrix']
        data_size = self.num_steps * self.batch_size
        loss = self.final_loss / data_size
        print('Loss: {}'.format(loss))
        print('Confusion matrix:')
        for r in cm:
            print('\t'.join([str(x) for x in r]))
        if self.best_loss < 0 or loss < self.best_loss:
            self.best_loss = loss
            self.copy_checkpoint_as_best()

class DocPredictionInference(evaluator.Evaluator):
    def __init__(self, dataset, log_dir=DIR_D2V_DOC_LOGDIR):
        config = tf.ConfigProto()
        config.gpu_options.allow_growth = True
        super(DocPredictionInference, self).__init__(checkpoints_dir=log_dir,
                                                     output_path=os.path.join(log_dir,
                                                                              dataset.type),
                                                     dataset=dataset,
                                                     singular_monitored_session_config=config,
                                                     infinite_loop=False)

    def model(self, input_vectors, input_gene, input_variation, batch_size,
              embedding_size=EMBEDDINGS_SIZE, output_classes=9):
        self.global_step = training_util.get_or_create_global_step()

        logits, _ = doc2vec_prediction_model(input_vectors, input_gene, input_variation,
                                             None, batch_size,
                                             is_training=False, embedding_size=embedding_size,
                                             output_classes=output_classes)
        global_step_increase = tf.assign_add(self.global_step, 1)
        with tf.control_dependencies([global_step_increase]):
            self.prediction = tf.nn.softmax(logits)
        return self.prediction

    def end(self, session):
        pass

    def create_graph(self, dataset_tensor, batch_size):
        input_vectors, input_gene, input_variation, _ = dataset_tensor
        return self.model(input_vectors, input_gene, input_variation, batch_size)

    def after_create_session(self, session, coord):
        super(DocPredictionInference, self).after_create_session(session, coord)
        print('ID,class1,class2,class3,class4,class5,class6,class7,class8,class9')

    def step(self, session, graph_data, summary_op):
        step, predictions = session.run([self.global_step, self.prediction])
        predictions = predictions[0]
        predictions = [p + 0.01 for p in predictions]  # penalize less the mistakes
        sum = np.sum(predictions)
        predictions = [p / sum for p in predictions]
        print('{},{}'.format(step, ','.join(['{:.3f}'.format(x) for x in predictions])))
        return None

if __name__ == '__main__':
    import logging

    logging.getLogger().setLevel(logging.INFO)
    if len(sys.argv) > 1 and sys.argv[1] == 'val':
        # get validation error
        evaluator = DocPredictionEval(dataset=DocPredictionDataset(type='val'),
                                      log_dir=os.path.join(DIR_D2V_DOC_LOGDIR))
        evaluator.run()
    elif len(sys.argv) > 1 and sys.argv[1] == 'test':
        # get validation error
        evaluator = DocPredictionInference(dataset=DocPredictionDataset(type='stage2_test'),
                                           log_dir=os.path.join(DIR_D2V_DOC_LOGDIR, 'val'))
        evaluator.run()
    elif len(sys.argv) > 1 and sys.argv[1] == 'train':
        # get validation error
        evaluator = DocPredictionEval(dataset=DocPredictionDataset(type='train'),
                                      log_dir=os.path.join(DIR_D2V_DOC_LOGDIR))
        evaluator.run()


import sys
sys.path.append('pddlstream/')
from ss.algorithms.dual_focused import dual_focused
import numpy as np

import kitchen2d.kitchen_stuff as ks
from kitchen_tasks.kitchen_utils import print_plan, execute_plan, step_plan, kitchen_creation_methods, initialize_objects_in_kitchen, update_sizes

from kitchen_tasks.kitchen_problem import create_problem
from kitchen_tasks.kitchen_predicates import *
##################################################

DEFAULT_TASK = {
    'kitchen_params': {
        'do_gui': False,
        'sink_w': 10.,
        'sink_h': 5.,
        'sink_d': 1.,
        'sink_pos_x': 20.,
        'left_table_width': 100.,
        'right_table_width': 100.,
        'faucet_h': 15.,
        'faucet_w': 5.,
        'faucet_d': 0.5,
        'obstacles': [],
        'liquid_name': 'coffee',
        'overclock': 100 # Set overclock higher if run visualization faster
    },
    'initial_poses': {
        'gripper': (0., 15., 0.),
        'cup': (7.5, 0., 0.),
        'sugar_cup': (-10., 0., 0.),
        'cream_cup': (15., 0, 0),
        'spoon': (0.5, 0.5, 0),
        'stirrer': (20, 0.5, 0),
        'block': (-20., 0, 0),
    },
    'object_sizes': {
        'spoon': (0.2, 3, 1.),
        'stirrer': (0.2, 5., 0.5),
        'block': (5, 1, 0),
        'holder': (1., 2., 0.5), # default size of a holder
    },
    'expid_pour': 0,
    'expid_scoop': 0,
}

'''
We use a holder for each spoon and stirrer object to
ensure that these objects can be grasped.
The size and location of each holder depends on the
spoon or stirrer it is holding. See function make_kitchen
in kitchen_tasks/kitchen_utils.py for how these holders
are created.
'''

##################################################

def sample_task_instance():
    '''
    Randomly sample a task instance, which has the same
    format as DEFAULT_TASK.
    '''
    task = DEFAULT_TASK.copy()
    expid_pour = task['expid_pour']
    expid_scoop = task['expid_scoop']
    sizes = task['object_sizes']
    sizes = update_sizes(expid_pour, expid_scoop, sizes)

    faucet_w = task['kitchen_params']['faucet_w']
    holder_d = task['object_sizes']['holder'][2]
    delta_dist = 2.
    grasp_dist = 3.
    left = -ks.SCREEN_WIDTH/2 + delta_dist + grasp_dist
    right = ks.SCREEN_WIDTH/2 - delta_dist - grasp_dist

    def sample_kitchen():
        while True:
            sink_pos_x = np.random.uniform(left, right)
            sink_w = np.random.uniform(faucet_w+delta_dist, faucet_w*2)
            if sink_pos_x + sink_w < right:
                return sink_pos_x, sink_w
    sink_pos_x, sink_w = sample_kitchen()
    task['kitchen_params']['sink_pos_x'] = sink_pos_x
    task['kitchen_params']['sink_w'] = sink_w

    x_intervals = [[left, sink_pos_x - grasp_dist],
                   [sink_pos_x + sink_w + grasp_dist, right]]

    gripper_x = np.random.uniform(left, right)

    names = ['cup', 'sugar_cup', 'cream_cup', 'spoon', 'stirrer', 'block']

    def sample_from_region(region):
        weight = [np.prod(a[1]-a[0]) for a in region]
        weight = weight / np.sum(weight)
        rid = np.random.choice(len(weight), p=weight)
        return np.random.uniform(region[rid][0], region[rid][1])

    def test_feasible(x_poses, debug=False):
        '''
        We use the following constraints:
        1, the distance between two objects is at least delta_dist.
        The distance between two objects is measured by the Euclidean
        distance from the right side of the left object to the left side
        of the right object.
        2, The distance from the center of any object to any side of other
        objects is at least grasp_dist.
        3, The objects are put on the table (table y axis is 0), but not in
        the sink (sink y axis is -sink_h).
        4, block is to the left of everything.
        '''
        for i in range(len(names)):
            if names[i] == 'block':
                if x_poses[i] > sink_pos_x:
                    return False
            for j in range(len(names)):
                if names[i] == 'block':
                    if x_poses[i] > x_poses[j]:
                        return False
                if i <= j:
                    continue
                if (i == 1 or j == 1) and debug:
                    import pdb; pdb.set_trace()
                if x_poses[i] > x_poses[j]:
                    # ensure object i is to the left of j
                    i_ = j; j_ = i
                else:
                    i_ = i; j_ = j
                r_i = x_poses[i_] + max(sizes[names[i_]][0], sizes[names[i_]][2])
                l_j = x_poses[j_] - max(sizes[names[j_]][0], sizes[names[i_]][2])
                if l_j - r_i < delta_dist:
                    return False
                if x_poses[j_] - r_i < grasp_dist:
                    return False
                if l_j - x_poses[i_] < grasp_dist:
                    return False

        return True

    while True:
        x_poses = [sample_from_region(x_intervals) for i in range(len(names))]
        if test_feasible(x_poses):
            break
    test_feasible(x_poses, False)
    cup_x, sugar_x, cream_x, spoon_x, stir_x, block_x = x_poses
    initial_poses = {
        'gripper': (gripper_x, 20., 0.),
        'cup': (cup_x, 0., 0.),
        'sugar_cup': (sugar_x, 0., 0.),
        'cream_cup': (cream_x, 0, 0),
        # Notice that spoons and stirrers will be put in a holder cup
        'spoon': (spoon_x, holder_d, 0),
        'stirrer': (stir_x, holder_d, 0),
        'block': (block_x, 0, 0),
    }
    task['initial_poses'] = initial_poses

    return task

##################################################

def solve_task(kitchen_params, initial_poses, object_sizes, expid_pour, expid_scoop):
    '''
    Solve for a plan of a task and visualize the results.
    Args:
        kitchen_params: init parameters of a Kitchen2D object. See Kitchen2D in
            kitchen2d/kitchen_stuff.py
        initial_poses: a dictionary mapping from the object names to their poses.
        object_sizes: a dictionary mapping from the object names to their sizes.
        expid_pour: experiment ID for learning to pour.
        expid_scoop: experiment ID for learning to scoop.
    '''
    make_kitchen, make_body = kitchen_creation_methods(expid_pour, expid_scoop,
        kitchen_params, initial_poses, object_sizes)

    block_goal = (-ks.SCREEN_WIDTH/2 + 5, 0, 0)

    '''
    Initial state of the system.
    Poses of the objects will be added
    later in create_problem.
    '''
    initial_atoms = [
        # Static atoms
        IsPose('block', block_goal),
        # Fluent atoms
        Empty('gripper'),
        CanMove('gripper'),
        HasSugar('sugar_cup'),
        HasCream('cream_cup'),
        IsPourable('cream_cup'),
        Stackable('cup', 'block'),
        Clear('block')
    ]

    # Goal literals (list of conjunctive literals)
    goal_literals = [
        AtPose('block', block_goal),
        On('cup', 'block'),
        HasCoffee('cup'),
        HasCream('cup'),
        HasSugar('cup'),
        Mixed('cup'),
        Empty('gripper'),
    ]

    # Construct the simulator
    kitchen = make_kitchen(planning=False)
    body_from_name = initialize_objects_in_kitchen(kitchen, make_body, initial_poses, initial_atoms)
    kitchen.enable_gui()

    # Create a problem for STRIPStream
    problem = create_problem(initial_poses, initial_atoms, goal_literals, make_kitchen, make_body,
                   expid_pour, expid_scoop, do_motion=False, do_collisions=True, do_grasp=True)

    # Solve the problem in STRIPStream
    raw_input('Start running STRIPStream solver?')
    plan, evaluations = dual_focused(problem, planner='ff-wastar3', max_time=600, verbose=True,
        max_planner_time=100, bind=True, use_context=True)

    print_plan(plan, evaluations)
    if plan is None:
        return

    if raw_input('Show plan? [y/n]') in ['y', 'Y']:
        # Show plan simulation
        step_plan(plan, kitchen, body_from_name)
    if raw_input('Show plan execution? [y/n]') in ['y', 'Y']:
        # Show execution simulation
        kitchen = make_kitchen(planning=False)
        body_from_name = initialize_objects_in_kitchen(kitchen, make_body, initial_poses, initial_atoms)
        kitchen.enable_gui()
        execute_plan(plan, kitchen, body_from_name)

    raw_input('Done!')

if __name__ == '__main__':
    print 'Sampling task instances...'
    task = sample_task_instance()
    solve_task(**task)


import numpy as np
import torch
from torch.utils.data import TensorDataset, random_split
from torch.utils.data import DataLoader, RandomSampler, SequentialSampler
from transformers import (BertForSequenceClassification,
                          BertTokenizer,
                          AdamW,
                          BertConfig,
                          get_linear_schedule_with_warmup,
                         )
from typing import Optional

def NER(data):

    # Load a trained model and vocabulary that you have fine-tuned
    output_dir = '../fake_news/model_save/'
    model = BertForSequenceClassification.from_pretrained(output_dir)
    tokenizer = BertTokenizer.from_pretrained(output_dir)

    target_names = ['fake','real']

    # Tokenize all of the sequences and map the tokens to thier IDs.
    input_ids_new = []
    attention_masks_new = []

    encoded_dict = tokenizer.encode_plus(
                        data,                             # Sequence to encode
                        add_special_tokens = True,       # Add '[CLS]' and '[SEP]'
                        max_length = 128,
                        padding = 'max_length',          # Pad and truncate
                        truncation=True,                 #Truncate the seq
                        return_attention_mask = True,    # Construct attn. masks
                        return_tensors = 'pt',           # Return pytorch tensors
                    )

    # Add the encoded sequences to the list
    input_ids_new.append(encoded_dict['input_ids'])

    # And its attention mask
    attention_masks_new.append(encoded_dict['attention_mask'])

    input_ids_new = torch.cat(input_ids_new, dim=0)
    attention_masks_new = torch.cat(attention_masks_new, dim=0)

    # Prediction on test set
    model.eval()

    with torch.no_grad():
        outputs = model(input_ids_new, token_type_ids=None, attention_mask=attention_masks_new)
        logits = outputs[0]
        logits = logits.detach().cpu().numpy()
        predictions = logits[0].tolist()

    category_name = target_names[predictions.index(max(predictions))]
    print("The predicted category is:")
    print(target_names[predictions.index(max(predictions))])

    return(category_name)


import os
import sys
import cv2
import numpy as np
import traceback
import time
import threading
import pickle
import math

from PIL import ImageGrab
from PIL import Image
from PIL.ImageQt import ImageQt
from PyQt5.QtCore import *
from PyQt5.QtGui import *
from PyQt5.QtWidgets import *
from PyQt5 import uic

import likeyoubot_win
import likeyoubot_resource
import likeyoubot_configure

form_class = uic.loadUiType('./dogfootermaker/DogFooterMaker/dogfootermaker.ui')[0]
dogfootermaker_title = 'DogFooter Maker'

def resource_path(relative):
    return os.path.join(
    os.environ.get(
        "_MEIPASS2",
        os.path.abspath(".")
    ),
    relative
    )

class DFMConst():
    def __init__(self):
        pass

    SCREEN_W = likeyoubot_win.LYBWin.WIDTH
    SCREEN_H = likeyoubot_win.LYBWin.HEIGHT
    PREVIEW_W = 256
    PREVIEW_H = 256

class PlayThread(QThread):
    def __init__(self, ui):
        QThread.__init__(self)
        self.ui = ui
        self.win = self.ui.win

    def __del__(self):
        self.wait()

    def run(self):
        while True:
            appPlayer_name = self.ui.playerComboBox.currentText()
            if len(appPlayer_name) > 0:
                self.hwnd = self.ui.hwnd_dic[appPlayer_name]
                (anchor_x, anchor_y, end_x, end_y) = self.win.get_window_location(self.hwnd)
                adj_x, adj_y = self.win.get_player_adjust(self.hwnd)

                img = self.win.get_window_screenshot(self.hwnd, 2)
                crop_area = (adj_x + 2, adj_y + 30, adj_x + 2 + DFMConst.SCREEN_W, adj_y + 30 + DFMConst.SCREEN_H)
                self.ui.mainImg = img.crop(crop_area)
                img = self.ui.mainImg.convert("RGB")
                data = img.tobytes("raw", "RGB")
                qim = QImage(data, img.size[0], img.size[1], QImage.Format_RGB888)
                pixmap = QPixmap.fromImage(qim)

                self.ui.screenLabel.setPixmap(pixmap)

                # img = img.convert("RGB")
                # data = img.tobytes("raw", "RGB")
                # qim = QImage(data, img.size[0], img.size[1], QImage.Format_RGB888)
                # pixmap = QPixmap.fromImage(qim)
                # self.ui.screenLabel.setPixmap(pixmap)

                # self.ui.mainImg = Image.open(imagePath)
                # pixmap = QPixmap(imagePath)
                # # img = img.convert("RGB")
                # # data = img.tobytes("raw", "RGB")
                # # qim = QImage(data, img.size[0], img.size[1], QImage.Format_RGB888)
                # # pixmap = QPixmap.fromImage(qim)
                # self.screenLabel.setPixmap(pixmap)

                if self.ui.current_pixelbox_label != None:
                    x = self.ui.current_pixelbox_label.x()
                    y = self.ui.current_pixelbox_label.y()
                    w = self.ui.current_pixelbox_label.width()
                    h = self.ui.current_pixelbox_label.height()

                    loc_x = x + int(w*0.5)
                    loc_y = y + int(w*0.5)

                    self.ui.screenLabel.mouseMove(loc_x, loc_y, init=True)

                time.sleep(0.01)

class PixelBoxLabel(QLabel):
    def __init__(self, parent=None):
        super(PixelBoxLabel, self).__init__(parent)
        self.ui = None
        self.pb_name = None
        self.size = None
        self.setMouseTracking(True)

    def setupPixelBox(self, ui, pb_name, size):
        self.ui = ui
        self.pb_name = pb_name
        self.size = size
        self.is_pressed = False
        self.setGeometry(0, 0, size, size)
        self.setMinimumWidth(size)
        self.setMinimumHeight(size)
        # self.setCursor(QCursor(Qt.BlankCursor))
        self.setCurrentFocus(False)
        self.hide()

    def setCurrentFocus(self, focus):
        self.ui.resourceLineEdit.clearFocus()
        if focus == True:
            self.setStyleSheet('PixelBoxLabel { background-color: rgb(255, 255, 255, 50%); }')
        else:
            self.setStyleSheet('PixelBoxLabel { background-color: rgb(255, 0, 0, 50%); }')

    def mouseMove(self, x, y, init=False):
        # print(self.x(), self.y(), self.width(), self.height())

        if init == True or ( self.ui.current_pixelbox_label == self and self.is_pressed == True and self.ui.controlKeyPressed == True):
            self.setGeometry(x - int(self.width()*0.5), y - int(self.height()*0.5), self.width(), self.height())
            return (self.x() + int(self.width()*0.5), self.y() + int(self.height()*0.5))

        return (-1, -1)

    def mouseMoveEvent (self, e):
        if self.ui.mainImg == None:
            return

        self.ui.screenLabel.mouseMove(e.x() + self.x(), e.y() + self.y())

    def mousePressEvent(self, e):
        self.is_pressed = True

        if self.ui.mainImg == None:
            return

        if e.buttons() & Qt.LeftButton:
            print('mousePressEvent LEFT Button')
            if self.ui.current_pixelbox_label != None:
                self.ui.current_pixelbox_label.setCurrentFocus(False)
            self.ui.current_pixelbox_label = self
            self.ui.current_pixelbox_label.setCurrentFocus(True)

            self.ui.screenLabel.mouseMove(e.x() + self.x(), e.y() + self.y())

        elif e.buttons() & Qt.RightButton:
            print('mousePressEvent RIGHT Button')
            if self.ui.current_pixelbox_label == self:
                self.ui.current_pixelbox_label = None

            self.setParent(None)
        else:
            print('mousePressEvent Else Button')

        self.ui.focusCurrentPixelBoxLabel()

    def mouseReleaseEvent(self, e):
        self.is_pressed = False
        if self.ui.mainImg == None:
            return

class MainScreenShotLabel(QLabel):
    def __init__(self, parent=None):
        super(MainScreenShotLabel, self).__init__(parent)
        self.setMouseTracking(True)
        self.ui = None
        self.is_pressed = False
        self.last_x = 0
        self.last_y = 0

    def setUI(self, ui):
        self.ui = ui

    def mouseMoveEvent (self, e):
        if self.ui.mainImg == None:
            return

        # 드래그 전송
        # if self.is_pressed == True:
        #     appPlayer_name = self.ui.playerComboBox.currentText()
        #     if len(appPlayer_name) > 0 and appPlayer_name in self.ui.hwnd_dic:
        #         hwnd = self.ui.hwnd_dic[appPlayer_name]
        #         (x, y, w, h) = self.ui.win.get_player_screen_rect(hwnd)
        #         self.ui.win.mouse_drag(hwnd, self.last_x + x, self.last_y + y, e.x() + x, e.y() + y)

        # print(e.x(), e.y())
        self.mouseMove(e.x(), e.y())

        # QWidget.mouseMoveEvent(self, eventQMouseEvent)

    def mouseMove(self, anchor_x, anchor_y, init=False):
        if self.ui.mainImg == None:
            return

        if self.ui.current_pixelbox_label == None:
            return

        (loc_x, loc_y) = self.ui.current_pixelbox_label.mouseMove(self.x() + anchor_x, self.y() + anchor_y, init=init)
        if loc_x == -1:
            return

        crop_area = self.getCropArea(loc_x - self.x(), loc_y - self.y(), 15)

        img = self.ui.mainImg.crop(crop_area)
        # img.save('test_maker.png')
        # img = img.convert("RGB")
        # data = img.tobytes("raw", "RGB")
        # print(img.size)
        qim = ImageQt(img)
        # qim = QImage(data, img.size[0], img.size[1], QImage.Format_RGB888)
        # qim = QImage(data, 14, 14, QImage.Format_RGB888)
        pixmap = QPixmap.fromImage(qim)
        pixmap = pixmap.scaled(DFMConst.PREVIEW_W * 0.75, DFMConst.PREVIEW_H * 0.75, Qt.KeepAspectRatio)
        self.ui.pixelBoxPreviewLabel.setPixmap(pixmap)

    def getCropArea(self, x, y, size):

        distance = int((size - 1) / 2)

        anchor_x = x
        anchor_y = y

        if anchor_x < distance:
            anchor_x = distance
        elif anchor_x + distance > DFMConst.SCREEN_W:
            anchor_x = DFMConst.SCREEN_W - distance

        if anchor_y < distance:
            anchor_y = distance
        elif anchor_y + distance > DFMConst.SCREEN_H:
            anchor_y = DFMConst.SCREEN_H - distance

        # print((anchor_x - distance, anchor_y - distance, anchor_x + distance, anchor_y + distance + 1))

        return (anchor_x - distance, anchor_y - distance, anchor_x + distance, anchor_y + distance)
        # return (x, y, x + size, y + size)
        # return (x, y, x + 15, y + 15)

    def mousePressEvent(self, e):
        self.is_pressed = True
        self.last_x = e.x()
        self.last_y = e.y()

        if self.ui.mainImg == None:
            return

        if e.buttons() & Qt.LeftButton:
            print('Left mouse button pressed ', e.x(), e.y())
            pb = self.ui.addPixelBoxLabel(e, 'pixel_box_label_' + str(e.x()), 15)
            self.mouseMove(e.x(), e.y(), init=True)
            # self.dragstart = event.pos()
            # self.clicked.emit()

            appPlayer_name = self.ui.playerComboBox.currentText()
            if len(appPlayer_name) > 0 and appPlayer_name in self.ui.hwnd_dic:
                hwnd = self.ui.hwnd_dic[appPlayer_name]
                # (x, y, w, h) = self.ui.win.get_player_screen_rect(hwnd)
                # self.ui.win.mouse_click(hwnd, e.x() + x, e.y() + y)

        elif e.buttons() & Qt.RightButton:
            print('Right mouse button pressed ', e.x(), e.y())

    def mouseReleaseEvent(self, e):
        self.is_pressed = False

        if self.ui.mainImg == None:
            return
        print('mouse released ', e.x(), e.y())
        # self.dragstart = None

    # def mouseMoveEvent(self, e):
    #     if self.ui.mainImg == None:
    #         return
    #     print('mouse move ', e.x(), e.y())
    #     if (self.dragstart is not None and
    # event.buttons() & QtCore.Qt.LeftButton and
    # (event.pos() - self.dragstart).manhattanLength() >
    # QtWidgets.qApp.startDragDistance()):
    # self.dragstart = None
    # drag = QtGui.QDrag(self)
    # drag.setMimeData(QtCore.QMimeData())
    # drag.exec_(QtCore.Qt.LinkAction)

    def dragEnterEvent(self, e):
        if self.ui.mainImg == None:
            return
        e.acceptProposedAction()
        if e.source() is not self:
            print('dragEnterEvent source is not self')
        else:
            print('dragEnterEvent source is self')

    # if event.source() is not self:
    # self.clicked.emit()

class PixelImageLabel(QLabel):
    def __init__(self, parent=None):
        super(PixelImageLabel, self).__init__(parent)
        self.setMouseTracking(True)
        self.ui = None

    def setUI(self, ui):
        self.ui = ui

    def mouseMoveEvent (self, eventQMouseEvent):
        if self.ui.mainImg == None:
            return

        # QWidget.mouseMoveEvent(self, eventQMouseEvent)

class MainWindow(QMainWindow, form_class):
    def __init__(self):
        super().__init__()
        self.setupUi(self)
        self.setWindowTitle(dogfootermaker_title)
        self.setFixedSize(QSize(DFMConst.SCREEN_W + 335, DFMConst.SCREEN_H + 250))
        self.statusBar.setSizeGripEnabled(False)
        self.win = likeyoubot_win.LYBWin(dogfootermaker_title)
        self.hwnd_dic = {}
        self.play_thread = None
        self.controlKeyPressed = False
        self.mainImg = None

        self.connectAction = QAction('Connect', self)
        self.connectAction.triggered.connect(self.connectPlayer)
        self.menuBar.addAction(self.connectAction)

        openAction = QAction('Open Image', self)
        openAction.triggered.connect(self.openImage)
        self.menuBar.addAction(openAction)

        self.screenLabel = MainScreenShotLabel(self.mainImageWidget)
        self.screenLabel.setUI(self)
        self.screenLabel.setGeometry(0, 0, DFMConst.SCREEN_W, DFMConst.SCREEN_H)
        self.screenLabel.setMinimumWidth(DFMConst.SCREEN_W)
        self.screenLabel.setMinimumHeight(DFMConst.SCREEN_H)
        self.screenLabel.setStyleSheet('MainScreenShotLabel { background-color: white; }')
        self.screenLabel.setAlignment(Qt.AlignLeft|Qt.AlignTop)

        self.pixelBoxPreviewLabel = PixelImageLabel(self.pixelImageWidget)
        self.pixelBoxPreviewLabel.setUI(self)
        self.pixelBoxPreviewLabel.setGeometry(0, 0, DFMConst.PREVIEW_W, DFMConst.PREVIEW_H)
        self.pixelBoxPreviewLabel.setMinimumWidth(DFMConst.PREVIEW_W)
        self.pixelBoxPreviewLabel.setMinimumHeight(DFMConst.PREVIEW_H)
        self.pixelBoxPreviewLabel.setStyleSheet('PixelImageLabel { background-color: white; border: 1px solid black; }')
        self.pixelBoxPreviewLabel.setAlignment(Qt.AlignCenter)

        self.pixelbox_label_dic = {}
        self.current_pixelbox_label = None
        self.playerComboBox.setFocusPolicy(Qt.NoFocus)
        self.playerComboBox.currentIndexChanged.connect(self.callbackPlayerComboBox)

        self.resource_manager_dic = {}

        for each_game, each_game_data in likeyoubot_configure.LYBConstant.LYB_GAMES.items():
            self.gameComboBox.addItem(each_game)
            try:
                with open(resource_path(each_game_data + '.lyb'), 'rb') as resource_file:
                    print('resource_file', resource_file)
                    self.resource_manager_dic[each_game] = pickle.load(resource_file)
            except:
                print(traceback.format_exc())
                self.resource_manager_dic[each_game] = likeyoubot_resource.LYBResourceManager(likeyoubot_resource.LYBPixelBoxDic(), likeyoubot_resource.LYBResourceDic())

        self.callbackGameComboBox()
        self.gameComboBox.currentIndexChanged.connect(self.callbackGameComboBox)
        self.gameComboBox.setFocusPolicy(Qt.NoFocus)

        self.callbackResourceComboBox()
        self.resourceComboBox.currentIndexChanged.connect(self.callbackResourceComboBox)
        self.resourceComboBox.setFocusPolicy(Qt.NoFocus)

        # self.pixelBoxListWidget.indexesMoved.connect(self.callbackPixelBoxListWidget)
        self.pixelBoxListWidget.clicked.connect(self.callbackPixelBoxListWidget)

        self.allCheckBox.stateChanged.connect(self.callbackAllCheckBox)
        self.eventCheckBox.stateChanged.connect(self.callbackEventCheckBox)
        self.sceneCheckBox.stateChanged.connect(self.callbackSceneCheckBox)
        self.locCheckBox.stateChanged.connect(self.callbackLocCheckBox)

        self.eventTableWidget.setRowCount(self.getResourceCount('event'))
        self.sceneTableWidget.setRowCount(self.getResourceCount('scene'))
        self.locTableWidget.setRowCount(self.getResourceCount('loc'))

        self.resourceLineEdit.textChanged.connect(self.callbackResourceLineEdit)

        self.updateResourceTable()

        # self.win.find_window_wildcard('')
        # for each_hwnd in self.win.handle_list:
        #     if each_hwnd in self.win.parent_handle_dic:
        #         title = self.win.get_title(self.win.parent_handle_dic[each_hwnd])
        #     else:
        #         title = self.win.get_title(each_hwnd)

        #     self.hwnd_dic[title] = each_hwnd

    def callbackResourceLineEdit(self):
        print('callbackResourceLineEdit')
        self.updateResourceComboBox()
        self.resourceLineEdit.setFocus()

    def getResourceCount(self, postfix):

        count = 0
        for each in self.resource_manager_dic[self.gameComboBox.currentText()].resource_dic.sortedResourceListByName():
            if postfix in each.resource_name[-5:]:
                count += 1

        return count

    def setResourceTableRow(self, widget, postfix):
        widget.setColumnCount(2)
        widget.setHorizontalHeaderLabels([postfix, '%'])
        widget.horizontalHeader().setSectionResizeMode(0, QHeaderView.Stretch)
        widget.horizontalHeader().setStyleSheet('::section{Background-color:#c2b6b6}')
        # widget.horizontalHeader().setSectionResizeMode(1, QHeaderView.Stretch)
        widget.setEditTriggers(QAbstractItemView.NoEditTriggers)
        widget.setColumnWidth(1, 10)
        widget.setSelectionBehavior(QTableView.SelectRows);
        widget.clicked.connect(self.callbackTableWidget)
        widget.setFocusPolicy(Qt.NoFocus)

        row = 0
        for each in self.resource_manager_dic[self.gameComboBox.currentText()].resource_dic.sortedResourceListByName():
            if postfix in each.resource_name[-5:]:
                item = QTableWidgetItem(each.resource_name)
                item.setTextAlignment(Qt.AlignVCenter | Qt.AlignLeft)
                widget.setItem(row, 0, item)
                item = QTableWidgetItem(99)
                item.setTextAlignment(Qt.AlignVCenter | Qt.AlignRight)
                widget.setItem(row, 1, item)
                row += 1

        widget.resizeColumnsToContents()
        widget.resizeRowsToContents()

    def callbackTableWidget(self, QPos=None):
        print(QPos.column(), QPos.row(), QPos.data())
        self.resourceComboBox.setCurrentIndex(self.resourceComboBox.findText(str(QPos.data())))

    def addPixelBoxLabel(self, e, name, size):
        if name in self.pixelbox_label_dic:
            name += '_' + str(len(self.pixelbox_label_dic) - 1)

        self.pixelbox_label_dic[name] = PixelBoxLabel(self.mainImageWidget)
        self.pixelbox_label_dic[name].setupPixelBox(self, name, size)
        self.pixelbox_label_dic[name].mouseMove(e.x(), e.y(), init=True)
        self.pixelbox_label_dic[name].show()
        self.current_pixelbox_label = self.pixelbox_label_dic[name]
        self.focusCurrentPixelBoxLabel()

        return self.pixelbox_label_dic[name]

    def focusCurrentPixelBoxLabel(self):
        if self.current_pixelbox_label == None:
            self.pixelBoxPreviewLabel.clear()
            return

        for key, value in self.pixelbox_label_dic.items():
            print('-- name: ', self.current_pixelbox_label.pb_name, key)
            if self.current_pixelbox_label.pb_name == key:
                self.pixelbox_label_dic[key].setCurrentFocus(True)
            else:
                self.pixelbox_label_dic[key].setCurrentFocus(False)

    def connectPlayer(self):
        # self.play_thread =  PlayThread(self)
        # self.play_thread.start()
        # try:
        #     with open(resource_path('lyb.cfg'), 'rb') as dat_file:
        #         self.configure = pickle.load(dat_file)

        #     self.configure.path = resource_path('lyb.cfg')
        #     dogfootermacro_logger.debug('configure.path [' + self.configure.path + ']')

        # except FileNotFoundError:
        #     exitMessage = QMessageBox.warning(self, ' ', resource_path('lyb2.cfg') + ' 파일이 없습니다.')
        #     sys.exit(1)

        if self.connectAction.text() == 'Connect':
            self.connectAction.setText('Disconnect')
            self.playerComboBox.clear()
            self.hwnd_dic = {}

            try:
                hwnd_dic = {}
                self.win.find_window_wildcard('')
                for each_hwnd in self.win.handle_list:
                    print(str(each_hwnd) + ' ' + self.win.get_title(each_hwnd))
                    if each_hwnd in self.win.parent_handle_dic:
                        title = self.win.get_title(self.win.parent_handle_dic[each_hwnd])
                    else:
                        title = self.win.get_title(each_hwnd)

                    self.playerComboBox.addItem(title)
                    hwnd_dic[title] = each_hwnd

                self.hwnd_dic = hwnd_dic
                print(self.playerComboBox.count())

                self.callbackPlayerComboBox()

            except:
                self.win = likeyoubot_win.LYBWin(dogfootermaker_title)
                self.hwnd_dic = {}
                print(traceback.format_exc())

        else:
            self.connectAction.setText('Connect')
            self.playerComboBox.clear()
            self.hwnd_dic = {}
            self.play_thread.terminate()
            self.play_thread = None
            print('Thread count ', threading.activeCount())
            self.screenLabel.clear()

    def openImage(self):
        print("openImage")
        option = QFileDialog.Options()
        option |= QFileDialog.DontUseNativeDialog
        imagePath, _ = QFileDialog.getOpenFileName(self, "QFileDialog.getOpenFileName()", "", "All Files(*);;Image Files (*.png)", options=option)
        if len(imagePath) > 0:
            self.mainImg = Image.open(imagePath)
            pixmap = QPixmap(imagePath)
            # img = img.convert("RGB")
            # data = img.tobytes("raw", "RGB")
            # qim = QImage(data, img.size[0], img.size[1], QImage.Format_RGB888)
            # pixmap = QPixmap.fromImage(qim)
            self.screenLabel.setPixmap(pixmap)
            self.clearPixelBoxOnScreen()

        # self.screenLabel.setStyleSheet("QLabel { background-color : red; color:blue;}")
        # pixmap = QPixmap('d:/Workspace/dogfooter/screenshot/elite_quest_reward_180901_145709.png')
        # dogfootermacro_logger.debug(pixmap)
        # dogfootermacro_logger.debug(self.mainImageLabel)
        # self.mainImageLabel.setPixmap(pixmap)

    def keyReleaseEvent(self, e):
        if self.controlKeyPressed == True:
            if e.key() == Qt.Key_Control:
                self.controlKeyPressed = False

    def keyPressEvent(self, e):
        def isPrintable(key):
            printable = [ Qt.Key_Space, Qt.Key_Exclam, Qt.Key_QuoteDbl, Qt.Key_NumberSign,
            Qt.Key_Dollar, Qt.Key_Percent, Qt.Key_Ampersand, Qt.Key_Apostrophe, Qt.Key_ParenLeft,
            Qt.Key_ParenRight, Qt.Key_Asterisk, Qt.Key_Plus, Qt.Key_Comma, Qt.Key_Minus, Qt.Key_Period,
            Qt.Key_Slash, Qt.Key_0, Qt.Key_1, Qt.Key_2, Qt.Key_3, Qt.Key_4, Qt.Key_5, Qt.Key_6, Qt.Key_7, Qt.Key_8, Qt.Key_9,
            Qt.Key_Colon, Qt.Key_Semicolon, Qt.Key_Less, Qt.Key_Equal, Qt.Key_Greater, Qt.Key_Question, Qt.Key_At,
            Qt.Key_A, Qt.Key_B, Qt.Key_C, Qt.Key_D, Qt.Key_E, Qt.Key_F, Qt.Key_G, Qt.Key_H, Qt.Key_I, Qt.Key_J, Qt.Key_K, Qt.Key_L,
            Qt.Key_M, Qt.Key_N, Qt.Key_O, Qt.Key_P, Qt.Key_Q, Qt.Key_R, Qt.Key_S, Qt.Key_T, Qt.Key_U, Qt.Key_V, Qt.Key_W, Qt.Key_X,
            Qt.Key_Y, Qt.Key_Z, Qt.Key_BracketLeft, Qt.Key_Backslash, Qt.Key_BracketRight, Qt.Key_AsciiCircum, Qt.Key_Underscore,
            Qt.Key_QuoteLeft, Qt.Key_BraceLeft, Qt.Key_Bar, Qt.Key_BraceRight, Qt.Key_AsciiTilde, ]

            if key in printable:
                return True
            else:
                return False

        if e.modifiers() & Qt.ControlModifier:
            print('Control')
            self.controlKeyPressed = True

        if e.modifiers() & Qt.ShiftModifier:
            print('Shift')

        if e.modifiers() & Qt.AltModifier:
            print('Alt')

        if e.key() == Qt.Key_Delete:
            print('Delete')
        elif e.key() == Qt.Key_Backspace:
            print('Backspace')
        elif e.key() in [Qt.Key_Return, Qt.Key_Enter]:
            print('Enter')
        elif e.key() == Qt.Key_Escape:
            print('Escape')
        elif e.key() == Qt.Key_Q and self.controlKeyPressed == True:
            print('Exit')
            sys.exit(1)
        elif (  e.key() == Qt.Key_Right or
                e.key() == Qt.Key_Left or
                e.key() == Qt.Key_Up or
                e.key() == Qt.Key_Down ):
            if self.current_pixelbox_label == None:
                return

            x = self.current_pixelbox_label.x()
            y = self.current_pixelbox_label.y()
            w = self.current_pixelbox_label.width()
            h = self.current_pixelbox_label.height()
            print(x, y, w, h)
            if e.key() == Qt.Key_Right:
                print('Right')
                loc_x = x + int(w*0.5) + 1
                loc_y = y + int(w*0.5)
            elif e.key() == Qt.Key_Left:
                print('Left')
                loc_x = x + int(w*0.5) - 1
                loc_y = y + int(w*0.5)
            elif e.key() == Qt.Key_Up:
                print('Up')
                loc_x = x + int(w*0.5)
                loc_y = y + int(w*0.5) - 1
            elif e.key() == Qt.Key_Down:
                print('Down')
                loc_x = x + int(w*0.5)
                loc_y = y + int(w*0.5) + 1

            if loc_x < int(w*0.5):
                loc_x = int(w*0.5)
            elif loc_x + w - int(w*0.5) > DFMConst.SCREEN_W:
                loc_x = DFMConst.SCREEN_W - w + int(w*0.5)

            if loc_y < int(h*0.5):
                loc_y = int(h*0.5)
            elif loc_y + h - int(h*0.5) > DFMConst.SCREEN_H:
                loc_y = DFMConst.SCREEN_H - h + int(h*0.5)

            self.screenLabel.mouseMove(loc_x, loc_y, init=True)

        if not self.controlKeyPressed and isPrintable(e.key()):
            print(e.text())

    def callbackPlayerComboBox(self):

        if len(self.hwnd_dic) == 0:
            return

        self.clearPixelBoxOnScreen()

        if self.play_thread == None:
            self.play_thread = PlayThread(self)
            self.play_thread.start()

    def callbackGameComboBox(self):
        self.updateResourceTable()
        self.allCheckBox.setChecked(True)
        self.callbackAllCheckBox()
        self.updateResourceComboBox()

    def updateResourceTable(self):
        self.setResourceTableRow(self.eventTableWidget, 'event')
        self.setResourceTableRow(self.sceneTableWidget, 'scene')
        self.setResourceTableRow(self.locTableWidget, 'loc')

    def updateResourceComboBox(self):
        self.resourceComboBox.clear()
        for each in self.resource_manager_dic[self.gameComboBox.currentText()].resource_dic.sortedResourceListByName():
            if len(self.resourceLineEdit.text()) > 0 and self.resourceLineEdit.text() not in each.resource_name:
                continue

            if self.allCheckBox.isChecked() == True:
                self.resourceComboBox.addItem(each.resource_name)
            else:
                if self.eventCheckBox.isChecked() == True and 'event' in each.resource_name[-5:]:
                    self.resourceComboBox.addItem(each.resource_name)
                elif self.sceneCheckBox.isChecked() == True and 'scene' in each.resource_name[-5:]:
                    self.resourceComboBox.addItem(each.resource_name)
                elif self.locCheckBox.isChecked() == True and 'loc' in each.resource_name[-5:]:
                    self.resourceComboBox.addItem(each.resource_name)
                else:
                    if (    'event' not in each.resource_name[-5:] and
                            'scene' not in each.resource_name[-5:] and
                            'loc' not in each.resource_name[-5:]
                            ):
                        self.resourceComboBox.addItem(each.resource_name)

        self.totalCountLabel.setText(str(len(self.resourceComboBox)))

    # 리소스를 선택할 때마다 호출

    def callbackResourceComboBox(self):
        print('+', self.resourceComboBox.currentText())
        self.pixelBoxListWidget.clear()
        self.clearPixelBoxOnScreen()
        if len(self.resourceComboBox.currentText()) > 0:
            for each in self.resource_manager_dic[self.gameComboBox.currentText()].resource_dic[self.resourceComboBox.currentText()]:
                self.pixelBoxListWidget.addItem(each)

        for i in range(self.pixelBoxListWidget.count()):
            print('-', self.pixelBoxListWidget.item(i).text())
            try:
                print('--', self.get_center_pixel_info(self.pixelBoxListWidget.item(i).text()))
            except:
                print(traceback.format_exc())

        if self.pixelBoxListWidget.count() > 0:
            self.pixelBoxListWidget.item(0).setSelected(True)
            self.pixelBoxListWidget.setFocus()
            self.pixelBoxListWidget.setCurrentRow(0)
            self.callbackPixelBoxListWidget()

    # 픽셀 박스 선택할때마다 호출
    # 1. 개별적인 픽셀박스 보여주기

    def callbackPixelBoxListWidget(self):
        print('DEBUG: PixelBox clicked')
        print('DEBUG: -->', self.pixelBoxListWidget.currentRow())
        if self.pixelBoxListWidget.currentRow() < 0:
            return

        print('DEBUG: --->', self.pixelBoxListWidget.currentItem().text(), self.gameComboBox.currentText())
        img = self.getImagePixelBox(self.pixelBoxListWidget.currentItem().text())
        qim = ImageQt(img)
        pixmap = QPixmap.fromImage(qim)
        pixmap = pixmap.scaled(DFMConst.PREVIEW_W * 0.75, DFMConst.PREVIEW_H * 0.75, Qt.KeepAspectRatio)
        self.pixelBoxPreviewLabel.setPixmap(pixmap)

    def callbackAllCheckBox(self):
        self.eventCheckBox.setChecked(self.allCheckBox.isChecked())
        self.sceneCheckBox.setChecked(self.allCheckBox.isChecked())
        self.locCheckBox.setChecked(self.allCheckBox.isChecked())
        self.updateResourceComboBox()

    def callbackEventCheckBox(self):
        self.updateResourceComboBox()

    def callbackSceneCheckBox(self):
        self.updateResourceComboBox()

    def callbackLocCheckBox(self):
        self.updateResourceComboBox()

    def clearPixelBoxOnScreen(self):
        for key, value in self.pixelbox_label_dic.items():
            self.pixelbox_label_dic[key].setParent(None)

        self.pixelbox_label_dic.clear()
        self.current_pixelbox_label = None

    def getImagePixelBox(self, pixel_box_name):

        try:
            pixel_box = self.resource_manager_dic[self.gameComboBox.currentText()].pixel_box_dic[pixel_box_name]
        except:
            print('Not found pixel box data:' + str(pixel_box_name))
            return None

        template_image = np.ndarray(shape=(pixel_box.height, pixel_box.width, 3), dtype='uint8')

        i = 0
        j = 0
        k = 0

        for elem in np.nditer(template_image, op_flags=['readwrite']):
            elem[...] = pixel_box[i + j][1][k]

            k += 1
            if k >= 3:

                i += pixel_box.height
                k = 0
                if i >= pixel_box.height * pixel_box.width:
                    j += 1
                    i = 0
                    rgb, w, h = template_image.shape[::-1]

        return Image.fromarray(template_image, 'RGB')

    def get_center_pixel_info(self, object_name):
        pixel_box = self.resource_manager_dic[self.gameComboBox.currentText()].pixel_box_dic[object_name]

        width = int(math.sqrt(len(pixel_box)))
        height = int(math.sqrt(len(pixel_box)))
        half = int(width * 0.5)

        middle_index = width * half + half

        return pixel_box[middle_index]

if __name__ == '__main__':

    app = QApplication(sys.argv)

    mainWindow = MainWindow()
    mainWindow.show()
    app.exec_()


import copy
import numpy as np
from torch import nn

from components.episode_buffer import EpisodeBatch
import torch as th
from torch.optim import Adam, RMSprop

from utils.torch_utils import huber
from modules.mixers import REGISTRY

def try_detach(maybe_tensor):
    if isinstance(maybe_tensor, int) or isinstance(maybe_tensor, float):
        return maybe_tensor
    else:
        return maybe_tensor.detach().cpu().numpy()

class RMIXLearner:
    """The learner of RMIX"""
    def __init__(self, mac, scheme, logger, args):
        self.args = args
        self.mac = mac
        self.logger = logger

        self.params = list(mac.parameters())

        self.last_target_update_episode = 0

        self.mixer = None
        if args.mixer is not None:
            self.mixer = REGISTRY[args.mixer](args)

            self.params += list(self.mixer.parameters())
            self.target_mixer = copy.deepcopy(self.mixer)

        self.cvar_optimiser = Adam(params=self.params, lr=args.lr)
        # a little wasteful to deepcopy (e.g. duplicates action selector), but should work for any MAC
        self.target_mac = copy.deepcopy(mac)

        self.log_stats_t = -self.args.learner_log_interval - 1

        self.cumulative_density = th.tensor(
            (2 * np.arange(self.args.num_atoms) + 1) / (2.0 * self.args.num_atoms), dtype=th.float32).view(1, -1).to(self.args.device)

        self._temp_qr_loss = th.tensor(0, device=self.args.device)
        self._temp_qr_td = th.tensor(0, device=self.args.device)
        # use to update the distribution without update the CVaR loss
        # not all params need to be updated
        self.last_qr_update_t = 0
        self.dist_optimiser = Adam(params=self.params, lr=args.qr_lr, eps=0.01/32)

    def train(self, batch: EpisodeBatch, t_env: int, episode_num: int):
        # Get the relevant quantities
        rewards = batch["reward"][:, :-1]
        actions = batch["actions"][:, :-1]
        terminated = batch["terminated"][:, :-1].float()
        mask = batch["filled"][:, :-1].float()
        mask[:, 1:] = mask[:, 1:] * (1 - terminated[:, :-1])
        avail_actions = batch["avail_actions"]

        # Calculate estimated Q-Values
        mac_out, mac_logits_out, masks, mac_risk_levels = [], [], [], []
        self.mac.init_hidden(batch.batch_size)
        for t in range(batch.max_seq_length):
            agent_outs, logits = self.mac.forward(batch, t=t)
            masks.append(self.mac.mask.view(self.args.batch_size, self.args.n_agents, -1))
            mac_out.append(agent_outs)
            mac_risk_levels.append(self.mac.current_risk_level)

            if self.update_qr(t_env):
                mac_logits_out.append(logits)

        mac_out = th.stack(mac_out, dim=1)  # Concat over time
        if not self.args.alpha_risk_static:
            mac_risk_levels = th.stack(mac_risk_levels, dim=1)  # Concat over time
        if self.update_qr(t_env):
            mac_logits_out = th.stack(mac_logits_out, dim=1)  # Concat over time
        masks = th.stack(masks, dim=1)

        _chosen_action_qdist = None
        # Pick the Q-Values and distribution for the actions taken by each agent
        _chosen_action_qvals = th.gather(mac_out[:, :-1], dim=3, index=actions).squeeze(3)  # Remove the last dim
        if self.update_qr(t_env):
            _actions = actions[..., None].repeat(*([1]*len(actions.size())), self.args.num_atoms)
            _chosen_action_qdist = th.gather(mac_logits_out[:, :-1], dim=3, index=_actions).squeeze(3)  # Remove the last dim

        # Calculate the Q-Values necessary for the target
        target_mac_out, target_logits_out, target_mask, target_mac_risk_levels = [], [], [], []
        self.target_mac.init_hidden(batch.batch_size)
        for t in range(batch.max_seq_length):
            target_agent_outs, target_logits = self.target_mac.forward(batch, t=t)
            target_mac_out.append(target_agent_outs)
            target_mac_risk_levels.append(self.target_mac.current_risk_level)
            if self.update_qr(t_env):
                target_logits_out.append(target_logits)
            target_mask.append(self.target_mac.mask)

        # We don't need the first (when n_step=1) timesteps Q-Value estimate for calculating targets
        target_mac_out = th.stack(target_mac_out[self.args.n_step:], dim=1)  # Concat across time
        if not self.args.alpha_risk_static:
            target_mac_risk_levels = th.stack(target_mac_risk_levels[self.args.n_step:], dim=1)   # Concat across time
        if self.update_qr(t_env):
            target_logits_out = th.stack(target_logits_out[self.args.n_step:], dim=1)  # Concat across time
        target_masks = th.stack(target_mask, dim=1)

        # Mask out unavailable actions
        target_mac_out[avail_actions[:, 1:] == 0] = -9999999
        if self.update_qr(t_env):
            _avail_actions = avail_actions[..., None].repeat(*([1]*len(avail_actions.size())), self.args.num_atoms)
            target_logits_out[_avail_actions[:, 1:] == 0] = -9999999

        # Max over target Q-Values
        _target_max_qdist = None
        if self.args.double_q:
            # Get actions that maximise live Q (for double q-learning)
            mac_out_detach = mac_out.clone().detach()
            mac_out_detach[avail_actions == 0] = -9999999
            cur_max_actions = mac_out_detach[:, 1:].max(dim=3, keepdim=True)[1]
            _target_max_qvals = th.gather(target_mac_out, 3, cur_max_actions).squeeze(3)

            if self.update_qr(t_env):
                _cur_max_actions = cur_max_actions[..., None].repeat(*([1]*len(cur_max_actions.size())), self.args.num_atoms)
                _target_max_qdist = th.gather(target_logits_out, 3, _cur_max_actions).squeeze(3)
        else:
            _target_max_qvals = target_mac_out.max(dim=3)[0]
            cur_max_actions = target_mac_out[:, 1:].max(dim=3, keepdim=True)[1]

            if self.update_qr(t_env):
                _cur_max_actions = cur_max_actions[..., None].repeat(*([1]*len(cur_max_actions.size())), self.args.num_atoms)
                _target_max_qdist = th.gather(target_logits_out, 3, _cur_max_actions).squeeze(3)

        # Mix
        if self.mixer is not None:
            chosen_action_qvals, chosen_action_qdist = self.mixer(_chosen_action_qvals,
                                                                  _chosen_action_qdist,
                                                                  batch["state"][:, :-1],
                                                                  masks[:, :-1],
                                                                  rewards)
            target_max_qvals, target_max_qdist = self.target_mixer(_target_max_qvals,
                                                                   _target_max_qdist,
                                                                   batch["state"][:, 1:],
                                                                   target_masks[:, :-1],
                                                                   rewards)
        else:
            chosen_action_qvals, chosen_action_qdist = _chosen_action_qvals, _chosen_action_qdist
            target_max_qvals, target_max_qdist = _target_max_qvals, _target_max_qdist

        # CVaR loss
        # Calculate 1-step Q-Learning targets
        cvar_targets = rewards + self.args.gamma * (1 - terminated) * target_max_qvals

        # CVaR Td-error
        cvar_td_error = (chosen_action_qvals - cvar_targets.detach())

        cvar_mask = mask.expand_as(cvar_td_error)

        # 0-out the targets that came from padded data
        masked_cvar_td_error = cvar_td_error * cvar_mask

        # Normal L2 loss, take mean over actual data
        cvar_loss = (masked_cvar_td_error ** 2).sum() / cvar_mask.sum()

        if self.update_qr(t_env):
            # Quantile target here
            # Calculate 1-step Q-Learning targets
            qr_targets = _chosen_action_qvals.detach()[..., None] + self.args.gamma * (1 - terminated[..., None]) * _target_max_qdist

            # QR Td-error
            qr_td_error = (_chosen_action_qdist[..., None] - qr_targets[:, :, :, None, :].detach())
            # 0-out the targets that came from padded data
            masked_qr_td_error = qr_td_error * mask[..., None, None]

            # Quantile Huber Loss
            qr_loss = huber(masked_qr_td_error) * (self.cumulative_density - (masked_qr_td_error.detach() < 0).float()).abs()

            qr_loss = qr_loss.sum(-1).mean(3).sum() / mask.repeat(1, 1, 3).sum()

            loss = cvar_loss
            self._temp_qr_loss = qr_loss
            self._temp_qr_td = masked_qr_td_error
        else:
            qr_loss = self._temp_qr_loss
            masked_qr_td_error = self._temp_qr_td

            loss = cvar_loss

        if self.update_qr(t_env):
            self.dist_optimiser.zero_grad()
            qr_loss.backward()
            grad_norm = th.nn.utils.clip_grad_norm_(self.params, self.args.grad_norm_clip)
            self.dist_optimiser.step()
        else:
            # Optimise cvar loss
            self.cvar_optimiser.zero_grad()
            loss.backward()
            grad_norm = th.nn.utils.clip_grad_norm_(self.params, self.args.grad_norm_clip)
            self.cvar_optimiser.step()

        if (episode_num - self.last_target_update_episode) / self.args.target_update_interval >= 1.0:
            self._update_targets()
            self.last_target_update_episode = episode_num

        if t_env - self.log_stats_t >= self.args.learner_log_interval:
            self.logger.log_stat("loss", loss.item(), t_env)
            self.logger.log_stat("cvar_loss", cvar_loss.item(), t_env)
            self.logger.log_stat("qr_loss", qr_loss.item(), t_env)
            self.logger.log_stat("grad_norm", grad_norm.item(), t_env)
            mask_elems = mask.sum().item()
            self.logger.log_stat("cvar_td_error_abs", (masked_cvar_td_error.abs().sum().item()/mask_elems), t_env)
            self.logger.log_stat("qr_td_error_abs", (masked_qr_td_error.abs().sum().item()/mask_elems), t_env)
            self.logger.log_stat("q_taken_mean", (chosen_action_qvals * mask).sum().item()/(mask_elems * self.args.n_agents), t_env)
            self.logger.log_stat("target_mean", (cvar_targets * mask).sum().item()/(mask_elems * self.args.n_agents), t_env)
            self.logger.log_stat("batch_rewards_mean", (rewards * mask).sum().item()/mask_elems, t_env)
            self.logger.log_stat("batch_rewards", rewards.cpu().numpy(), t_env)

            self.logger.log_stat("cvar_chosen_action_qvals", chosen_action_qvals.detach().cpu().numpy(), t_env)
            self.logger.log_stat("cvar_target_max_qvals", target_max_qvals.detach().cpu().numpy(), t_env)
            # QR distributions
            if self.update_qr(t_env):
                self.logger.log_stat("qr_orig_target_max_qdist", _target_max_qdist.detach().cpu().numpy(), t_env)
                self.logger.log_stat("qr_mixer_target_max_qdist", target_max_qdist.detach().cpu().numpy(), t_env)
                self.logger.log_stat("qr_orig_chosen_action_qdist", _chosen_action_qdist.detach().cpu().numpy(), t_env)
                self.logger.log_stat("qr_mixer_chosen_action_qdist", chosen_action_qdist.detach().cpu().numpy(), t_env)

            # save distribution and CVaR and Risk level
            if self.args.version in {'v1', 'v2'}:
                if self.args.save_risk_level:
                    self.logger.log_stat("risk_level", try_detach(self.mac.agent.current_risk_level), t_env)

            self.log_stats_t = t_env

        # fianally update the self.last_qr_update_t
        if self.update_qr(t_env):
            self.last_qr_update_t = t_env

    def update_qr(self, t_env):
        if (t_env - self.last_qr_update_t) / self.args.qr_update_interval >= 1.0 and t_env > self.args.start_to_update_qr:
            return True
        else:
            return False

    def _update_targets(self):
        self.target_mac.load_state(self.mac)
        if self.mixer is not None:
            self.target_mixer.load_state_dict(self.mixer.state_dict())
        self.logger.console_logger.info("Updated target network")

    def cuda(self):
        self.mac.cuda()
        self.target_mac.cuda()
        if self.mixer is not None:
            self.mixer.cuda()
            self.target_mixer.cuda()

    def save_models(self, path):
        self.mac.save_models(path)
        if self.mixer is not None:
            th.save(self.mixer.state_dict(), "{}/mixer.th".format(path))
        th.save(self.cvar_optimiser.state_dict(), "{}/cvar_opt.th".format(path))
        th.save(self.dist_optimiser.state_dict(), "{}/dist_opt.th".format(path))

    def load_models(self, path):
        self.mac.load_models(path)
        # Not quite right but I don't want to save target networks
        self.target_mac.load_models(path)
        if self.mixer is not None:
            self.mixer.load_state_dict(th.load("{}/mixer.th".format(path), map_location=lambda storage, loc: storage))
        self.cvar_optimiser.load_state_dict(th.load("{}/cvar_opt.th".format(path), map_location=lambda storage, loc: storage))
        self.dist_optimiser.load_state_dict(th.load("{}/dist_opt.th".format(path), map_location=lambda storage, loc: storage))