import pybullet as p
import numpy as np
from scipy.spatial.transform import Rotation
import xml.dom.minidom as xml
import os

from .active_obj import ActiveObject
from .end_effector import Gripper,Suction

class Robot(ActiveObject):
    def __init__(self,scene,**kwargs):
        if 'reset_joint_poses' not in kwargs: kwargs['reset_joint_poses'] = []
        if 'joint_damping' not in kwargs: kwargs['joint_damping'] = []
        if 'end_effector' not in kwargs: kwargs['end_effector'] = ''
        if 'speed' not in kwargs:  kwargs['speed'] = 1.0
        self.speed = kwargs['speed']
        self.current_joint_poses = self.reset_joint_poses = kwargs['reset_joint_poses']
        self.joint_damping = kwargs['joint_damping']
        self.end_effector = kwargs['end_effector']

        super().__init__(scene, **kwargs)
        self.pickup = False
        pass

    def properties(self):
        properties = super().properties()

        x,y,z,rx,ry,rz = self.get_end_effector_pose()
        properties.update(kind='Robot',speed=self.speed,
                          end_effector=self.end_effector,
                          joint_damping=self.joint_damping,
                          current_joint_poses=self.current_joint_poses,
                          reset_joint_poses=self.reset_joint_poses,
                          end_effector_pos=[x,y,z],end_effector_rot=[rx,ry,rz])
        return properties

    def update(self,dt):
        super().update(dt)
        if self.end_effector_obj:
            self.end_effector_obj.update(dt)
            if not self.end_effector_obj.idle: self.actions.append((lambda *args:None,()))
        pass

    def restore(self):
        super().restore()
        p.setJointMotorControlArray(self.id, self.used_joint_indexes, p.POSITION_CONTROL, self.reset_joint_poses)
        self.end_effector_obj.do(False)

    def set_base(self,base):
        base_temp = self.base = base

        ee_kind = ""
        mimic_name = ''
        gears = []
        with xml.parse(os.path.join(self.data_dir,base)) as doc_robot:
            node_robot = doc_robot.getElementsByTagName('robot')[0]
            for mesh in node_robot.getElementsByTagName('mesh'):
                filename = mesh.getAttribute('filename')
                mesh.setAttribute('filename',os.path.join(self.data_dir,os.path.dirname(base),filename))

            num_joints = len(node_robot.getElementsByTagName('link'))

            if self.end_effector:
                link_ee = node_robot.getElementsByTagName('link')[-1]
                link_ee_name = link_ee.getAttribute('name')
                with xml.parse(os.path.join(self.data_dir,self.end_effector)) as doc_ee:
                    node_ee = doc_ee.getElementsByTagName('robot')[0]
                    ee_kind = node_ee.getAttribute("kind")
                    for mesh in node_ee.getElementsByTagName('mesh'):
                        filename = mesh.getAttribute('filename')
                        mesh.setAttribute('filename',os.path.join(self.data_dir,os.path.dirname(self.end_effector),filename))
                    for i,link in enumerate(node_ee.getElementsByTagName('link')):
                        if i == 0: link.setAttribute('name',link_ee.getAttribute('name'))
                        node_robot.appendChild(link)
                    for i,joint in enumerate(node_ee.getElementsByTagName('joint')):
                        if i == 0: joint.getElementsByTagName('parent')[0].setAttribute('link',link_ee_name)
                        node_robot.appendChild(joint)
                        node_mimics = joint.getElementsByTagName('mimic')
                        if not node_mimics: continue
                        node_mimic = node_mimics[0]
                        mimic_name = node_mimic.getAttribute('joint')
                        joint_name = joint.getAttribute('name')
                        multiplier = int(node_mimic.getAttribute('multiplier')) if node_mimic.hasAttribute('multiplier') else 1
                        offset = 0#int(node_mimic.getAttribute('offset'))
                        gears.append((mimic_name,joint_name,multiplier))
                node_robot.removeChild(link_ee).unlink()

            import tempfile
            f = tempfile.NamedTemporaryFile("w",delete=False)
            base_temp = f.name
            f.write(node_robot.toxml())
            f.close()

        if 'id' in vars(self): p.removeBody(self.id)
        print('robot urdf',base_temp,flush=True)
        self.id = p.loadURDF(base_temp, self.pos, p.getQuaternionFromEuler(self.rot),useFixedBase=True)
        if ee_kind == 'Gripper': self.end_effector_obj = Gripper(self.id,gears)
        elif ee_kind == 'Suction': self.end_effector_obj = Suction(self.id)
        else:
            class EndEffector:
                def __init__(self): self.idle = True
                def update(self,dt):pass
                def do(self,pickup):pass
                def get_properties(self): return dict(king='EndEffector')
            self.end_effector_obj = EndEffector()
            pass

        self.used_joint_indexes = []
        self.active_joint_indexes = []
        self.joint_velocities = []
        for i in range(p.getNumJoints(self.id)):
            ji = p.getJointInfo(self.id, i)
            jointName,jointType,jointVelocity = ji[1],ji[2],ji[11]
            if (jointType != p.JOINT_REVOLUTE and jointType != p.JOINT_PRISMATIC and jointType != p.JOINT_SPHERICAL): continue
            self.active_joint_indexes.append(i)
            self.joint_velocities.append(jointVelocity)
            if i < num_joints: self.used_joint_indexes.append(i)

        if len(self.active_joint_indexes) > len(self.joint_damping): self.joint_damping.extend([0]*(len(self.active_joint_indexes)-len(self.joint_damping)))
        else: self.joint_damping = self.joint_damping[:len(self.active_joint_indexes)]

        if len(self.used_joint_indexes) > len(self.reset_joint_poses): self.reset_joint_poses.extend([0]*(len(self.used_joint_indexes)-len(self.reset_joint_poses)))
        else: self.reset_joint_poses = self.reset_joint_poses[:len(self.used_joint_indexes)]

        self.set_joints(self.reset_joint_poses)

        if 'lines' in vars(self):
            for line in self.lines:
                p.removeUserDebugItem(line)

        self.lines = list()

        pos,orn,_,_,_,_ = p.getLinkState(self.id,p.getNumJoints(self.id)-1)
        axis_x = Rotation.from_quat(orn).apply(np.array([0.05,0,0])) + pos
        axis_y = Rotation.from_quat(orn).apply(np.array([0,0.05,0])) + pos
        axis_z = Rotation.from_quat(orn).apply(np.array([0,0,0.05])) + pos
        self.lines.extend([p.addUserDebugLine(pos,axis_x,[1,0,0],2),
                          p.addUserDebugLine(pos,axis_y,[0,1,0],2),
                          p.addUserDebugLine(pos,axis_z,[0,0,1],2)])

        self.home_pos,self.home_rot = pos,p.getEulerFromQuaternion(orn)

    def set_speed(self,value):
        if round(value) == 0.000: return
        self.speed = value

    def get_joints(self):
        return self.current_joint_poses[:len(self.used_joint_indexes)]

    def set_joints(self,joint_poses):
        joint_poses = joint_poses[:len(self.used_joint_indexes)]
        for i in self.used_joint_indexes: p.resetJointState(self.id, i, joint_poses[i])
        p.setJointMotorControlArray(self.id, self.used_joint_indexes, p.POSITION_CONTROL, joint_poses)
        self.current_joint_poses = joint_poses[:len(self.used_joint_indexes)]

    def set_end_effector_pose(self,pose):
        num_joints = p.getNumJoints(self.id)
        ee_index = num_joints-1
        joint_poses = p.calculateInverseKinematics(self.id, ee_index, pose[0:3], p.getQuaternionFromEuler(pose[3:6]),restPoses=self.reset_joint_poses,maxNumIterations=200)
        self.set_joints(list(joint_poses)[:len(self.used_joint_indexes)])

    def get_end_effector_pose(self):
        num_joints = p.getNumJoints(self.id)
        ee_pos,ee_orn,_,_,_,_ = p.getLinkState(self.id,num_joints-1)
        ee_rot = p.getEulerFromQuaternion(ee_orn)
        return [ee_pos[0],ee_pos[1],ee_pos[2],ee_rot[0],ee_rot[1],ee_rot[2]]

    def set_end_effector(self,base):
        if 'id' not in vars(self): return
        self.end_effector = base
        self.set_base(self.base)

    def end_pose_from_joint_poses(self,joint_poses):
        pos = np.array([0,0,0])
        rot = np.array([0,0,0])
        i = p.getNumJoints(self.id)-1

        while i > -1:
            _,_,l_pos,l_orn,_,_,_,_ = p.getLinkState(self.id,i,True)
            pos = pos + l_pos
            rot = (Rotation.from_quat(l_orn) * Rotation.from_euler('xyz',rot)).as_euler('xyz')
            _,_,_,_,_,_,_,_,_,_,_,_,_,j_axis,j_pos,j_orn,parent_link_index = p.getJointInfo(self.id, i)
            R = Rotation.from_quat(j_orn)
            if i in self.used_joint_indexes:
                R = R * Rotation.from_rotvec(np.array(j_axis) * joint_poses[self.used_joint_indexes.index(i)])

            pos = R.apply(pos) + j_pos
            rot = (R * Rotation.from_euler('xyz',rot)).as_euler('xyz')
            i = parent_link_index

        pos = Rotation.from_euler('xyz',self.rot).apply(pos) + self.pos
        rot = (Rotation.from_euler('xyz',self.rot) * Rotation.from_euler('xyz',rot)).as_euler('xyz')

        return [pos[0],pos[1],pos[2],rot[0],rot[1],rot[2]]

    def plan(self,target_pos,target_rot,s=1,mode='joint'):
        target_orn = p.getQuaternionFromEuler(target_rot)
        num_joints = p.getNumJoints(self.id)
        ee_index = num_joints-1

        ee_pos,ee_orn,_,_,_,_ = p.getLinkState(self.id,num_joints-1)
        ee_pos = np.array(ee_pos)
        ee_rot = p.getEulerFromQuaternion(ee_orn)

        route_poses = list()

        if mode=='joint':
            end_poses = p.calculateInverseKinematics(self.id, ee_index, target_pos, target_orn, restPoses=self.reset_joint_poses, maxNumIterations=200)
            end_poses = list(end_poses)[:len(self.used_joint_indexes)]
            num = round(1 / self.scene.timestep)
            for n in range(num):
                poses = []
                t = n / num
                for i in range(len(self.current_joint_poses)):
                    poses.append(np.interp(t,[0,1],[self.current_joint_poses[i],end_poses[i]]))
                route_poses.append(poses[:len(self.used_joint_indexes)])
        else:
            num = round(1 / self.scene.timestep)
            for n in range(num):
                poses = []
                t = n / num
                x = np.interp(t,[0,1],[ee_pos[0],target_pos[0]])
                y = np.interp(t,[0,1],[ee_pos[1],target_pos[1]])
                z = np.interp(t,[0,1],[ee_pos[2],target_pos[2]])
                pos = [x,y,z]

                q1 = np.array(ee_orn)
                q2 = np.array(target_orn)
                q_dot = q1.dot(q2)

                if (np.abs(q_dot) < 1e-5):
                    orn = q1
                else:
                    if q_dot < 0:
                        q2 = -q2

                    orn = (1 - t) * q1 + t * q2
                    orn = orn / np.linalg.norm(orn)

                poses = p.calculateInverseKinematics(self.id, ee_index, pos, orn, restPoses=self.reset_joint_poses, maxNumIterations=200)
                route_poses.append(list(poses)[:len(self.used_joint_indexes)])

        # points = []
        # for poses in route_poses:
        #     point = self.end_pose_from_joint_poses(poses)
        #     points.append(point)

        # def task(i):
        #     self.lines.append(p.addUserDebugLine(points[i][0:3],points[i+1][0:3],[0,1,0],1,lifeTime=0))

        # for i in range(len(points)-1):
        #     self.actions.append((task,(i,)))
        return route_poses

    def signal_plan_move(self,*args,**kwargs):
        objs, = args
        def near_obj(o):
            e_pos,e_orn,_,_,_,_ = p.getLinkState(self.id,p.getNumJoints(self.id)-1)
            o_pos,e_pos = np.array(o[0]),np.array(e_pos)
            distance = abs(np.linalg.norm(o_pos - e_pos))
            return o_pos[2] * 100 - distance

        o_pos,o_rot,o_mesh = max(objs,key=near_obj)
        o_pos = np.array(o_pos)
        o_r = Rotation.from_euler("xyz",o_rot)

        def perfect_pick_pose(pick_pose):
            pick_p = np.array(pick_pose['pos'])
            pick_r = Rotation.from_euler("xyz",pick_pose['rot'])
            o = o_pos + o_r.apply(pick_p)
            pick_x = o_pos + o_r.apply(pick_p + pick_r.apply([0.1,0,0]))
            pick_y = o_pos + o_r.apply(pick_p + pick_r.apply([0,0.1,0]))
            pick_z = o_pos + o_r.apply(pick_p + pick_r.apply([0,0,0.1]))
            # p.addUserDebugLine(o,pick_z,[0,0,1],5,lifeTime=5)

            axis_up = [0,0,1]
            pick_axis_up = (o_r * pick_r).apply([0,0,1])
            same = np.dot(axis_up,pick_axis_up) / np.linalg.norm(axis_up) * np.linalg.norm(pick_axis_up)
            return same

        pose = max(kwargs['pick_poses'],key=perfect_pick_pose)
        pick_pos,pick_rot = pose['pos'],pose['rot']
        pick_pos = o_pos + o_r.apply(pick_pos)
        pick_r = o_r * Rotation.from_euler("xyz",pick_rot)
        pick_rot = pick_r.as_euler('xyz')

        num_joints = p.getNumJoints(self.id)
        ee_index = num_joints-1

        route_poses = self.plan(pick_pos, pick_rot)
        def task(*poses):
            p.setJointMotorControlArray(self.id, self.used_joint_indexes, p.POSITION_CONTROL, poses)
            self.current_joint_poses = poses

        for poses in route_poses: self.actions.append((task,poses))
        def output(last_pos):
            pos,_,_,_,_,_,_,_ = p.getLinkState(self.id,ee_index,True)
            pos = np.linalg.norm(pos)
            if round(last_pos - pos,6) != 0.000000:
                self.actions.append((output,(pos,)))
                return

            if 'pickup' in kwargs: self.end_effector_obj.do(kwargs['pickup'])
            self.result = None,

        self.actions.append((output,(0,)))

    def signal_pick_move(self,*args,**kwargs):
        if not args:
            active_plugins = kwargs['plugins']
            picklight = active_plugins['PickLight']
            _,pickpose = picklight.result
            args = pickpose,

        num_joints = p.getNumJoints(self.id)
        ee_index = num_joints-1
        ee_pos,ee_orn,_,_,_,_ = p.getLinkState(self.id,ee_index)
        ee_pos = np.array(ee_pos)
        ee_rot = p.getEulerFromQuaternion(ee_orn)

        if 'speed' not in kwargs: kwargs['speed'] = 1
        speed = kwargs['speed']
        if speed == 0:
            if 'pickup' in kwargs: self.end_effector_obj.do(kwargs['pickup'])
            self.result = None,[ee_pos[0],ee_pos[1],ee_pos[2],ee_rot[0],ee_rot[1],ee_rot[2]]
            return

        if 'mode' not in kwargs: kwargs['mode'] = 'joint'
        mode = kwargs['mode']

        pick_pose = args[0]
        pick_pos,pick_rot = pick_pose[0:3],pick_pose[3:6]
        pick_orn = p.getQuaternionFromEuler(pick_rot)
        route_poses = self.plan( pick_pos, pick_rot, speed, mode)

        axis_x = Rotation.from_euler('xyz',pick_rot).apply([0.05,0,0]) + pick_pos
        axis_y = Rotation.from_euler('xyz',pick_rot).apply([0,0.05,0]) + pick_pos
        axis_z = Rotation.from_euler('xyz',pick_rot).apply([0,0,0.05]) + pick_pos

        self.lines.extend([p.addUserDebugLine(pick_pos,axis_x,[1,0,0],2),
                           p.addUserDebugLine(pick_pos,axis_y,[0,1,0],2),
                           p.addUserDebugLine(pick_pos,axis_z,[0,0,1],2)])

        def task(*poses):
            p.setJointMotorControlArray(self.id, self.used_joint_indexes, p.POSITION_CONTROL, poses)
            self.current_joint_poses = poses

        for poses in route_poses: self.actions.append((task,poses))
        def output(last_pos):
            pos,_,_,_,_,_,_,_ = p.getLinkState(self.id,ee_index,True)
            pos = np.linalg.norm(pos)
            if round(last_pos - pos,6) != 0.000000:
                self.actions.append((output,(pos,)))
                return

            if 'pickup' in kwargs: self.end_effector_obj.do(kwargs['pickup'])
            self.result = None,

        self.actions.append((output,(0,)))
        pass

    def signal_move(self,*args,**kwargs):
        if 'home' in kwargs and kwargs['home']:
            self.signal_home()
            return

        num_joints = p.getNumJoints(self.id)
        ee_index = num_joints-1
        ee_pos,ee_orn,_,_,_,_ = p.getLinkState(self.id,ee_index)
        ee_pos = np.array(ee_pos)
        ee_rot = p.getEulerFromQuaternion(ee_orn)

        if 'speed' not in kwargs: kwargs['speed'] = 1
        speed = kwargs['speed']
        if speed == 0:
            if 'pickup' in kwargs: self.end_effector_obj.do(kwargs['pickup'])
            self.result = None,[ee_pos[0],ee_pos[1],ee_pos[2],ee_rot[0],ee_rot[1],ee_rot[2]]
            return

        t_pos = ee_pos
        t_rot = ee_rot
        poses = []

        def task(poses):
            p.setJointMotorControlArray(self.id, self.used_joint_indexes, p.POSITION_CONTROL, poses)
            self.current_joint_poses = poses

        if 'mode' not in kwargs: kwargs['mode'] = 'joint'
        mode = kwargs['mode']

        if 'point' in kwargs:
            t_pos = kwargs['point'][0:3]
            t_rot = kwargs['point'][3:6]
            route_poses = self.plan(t_pos,t_rot,speed,mode)
            for poses in route_poses: self.actions.append((task,(poses,)))
        else:
            end_poses = kwargs['joints'][:len(self.current_joint_poses)]
            num = int(1 / speed / self.scene.timestep)
            route_poses = []
            for n in range(num):
                poses = []
                t = n / num
                for i in range(len(self.current_joint_poses)):
                    poses.append(np.interp(t,[0,1],[self.current_joint_poses[i],end_poses[i]]))
                self.actions.append((task,(poses,)))

        def output(last_pos):
            pos,_,_,_,_,_ = p.getLinkState(self.id,ee_index)
            pos = np.linalg.norm(pos)
            if round(last_pos - pos,6) != 0.000000:
                self.actions.append((output,(pos,)))
                return

            if 'pickup' in kwargs: self.end_effector_obj.do(kwargs['pickup'])
            self.result = None,

        self.actions.append((output,(0,)))
        pass

    def signal_move_relatively(self,*args,**kwargs):
        num_joints = p.getNumJoints(self.id)
        ee_index = num_joints-1
        ee_pos,ee_orn,_,_,_,_ = p.getLinkState(self.id,ee_index)
        ee_rot = p.getEulerFromQuaternion(ee_orn)

        if args:
            pick_point = args[0]
            ee_pos = pick_point[0:3]
            ee_rot = pick_point[3:6]
            ee_orn = p.getQuaternionFromEuler(ee_rot)

        if 'target' not in kwargs: kwargs['target'] = 'task_current'
        target = kwargs['target']

        if target == 'task_next':
            declare = kwargs['declare']
            cursor = kwargs['cursor']
            active_plugins = kwargs['plugins']

            act = declare[cursor]
            next = act['next']
            act_next = declare[next]

            if act_next['fun'] == 'move_relatively':
                self.result = 'link error.',
                return
            elif act_next['fun'] == 'move':
                move_args = act_next['args']

                if 'point' in move_args:
                    pick_pose = move_args['point']
                elif 'joints' in move_args:
                    joints = move_args['joints']
                    n = len(self.used_joint_indexes) - len(joints)
                    if n > 0: joints.extend([0] * n)
                    pick_pose = self.end_pose_from_joint_poses(joints)
                else:
                    self.result = 'parameters error.',
                    return
            elif act_next['fun'] == 'pick_move':
                picklight = active_plugins['PickLight']
                _,pick_pose = picklight.result

                if not pick_pose:
                    self.result = 'failed',
                    return

            ee_pos = pick_pose[0:3]
            ee_rot = pick_pose[3:6]
            ee_orn = p.getQuaternionFromEuler(ee_rot)

        if 'speed' not in kwargs: kwargs['speed'] = 1
        speed = kwargs['speed']
        if speed == 0:
            if 'pickup' in kwargs: self.end_effector_obj.do(kwargs['pickup'])
            self.result = None,[ee_pos[0],ee_pos[1],ee_pos[2],ee_rot[0],ee_rot[1],ee_rot[2]]
            return

        if 'mode' not in kwargs: kwargs['mode'] = 'joint'
        mode = kwargs['mode']

        t_pos = ee_pos
        t_rot = ee_rot

        if 'point' in kwargs:
            t_pos = kwargs['point'][0:3] + np.array(ee_pos)
            t_rot = kwargs['point'][3:6] + np.array(ee_rot)

        route_poses = self.plan(t_pos,t_rot,speed,mode)

        def task(poses):
            p.setJointMotorControlArray(self.id, self.used_joint_indexes, p.POSITION_CONTROL, poses)
            self.current_joint_poses = poses

        for poses in route_poses: self.actions.append((task,(poses,)))

        def output(last_pos):
            pos,_,_,_,_,_ = p.getLinkState(self.id,ee_index)
            pos = np.linalg.norm(pos)
            if round(last_pos - pos,6) != 0.000000:
                self.actions.append((output,(pos,)))
                return

            if 'pickup' in kwargs: self.end_effector_obj.do(kwargs['pickup'])
            self.result = None,

        self.actions.append((output,(0,)))
        pass

    def signal_do(self,*args,**kwargs):
        self.end_effector_obj.do(kwargs['pickup'])
        pass

    def signal_pick(self,*args,**kwargs):
        self.end_effector_obj.do(True)
        pass

    def signal_place(self,*args,**kwargs):
        self.end_effector_obj.do(False)
        pass

    def signal_home(self):
        if 'lines' in vars(self):
            for line in self.lines:
                p.removeUserDebugItem(line)

        self.lines = list()

        num_joints = p.getNumJoints(self.id)
        ee_index = num_joints-1
        ee_pos,ee_orn,_,_,_,_ = p.getLinkState(self.id,ee_index)
        ee_pos = np.array(ee_pos)
        ee_rot = p.getEulerFromQuaternion(ee_orn)

        t_pos = self.home_pos
        t_rot = self.home_rot

        self.reset_joint_poses
        joint_poses = []
        for joint in self.used_joint_indexes:
            position,velocity,reaction_force,motor_torque = p.getJointState(self.id,joint)
            joint_poses.append(position)

        def task(poses):
            p.setJointMotorControlArray(self.id, self.used_joint_indexes, p.POSITION_CONTROL, poses)
            self.current_joint_poses = poses

        num = int(1 / self.speed / self.scene.timestep) + 1
        for n in range(num):
            poses = []
            t = n / num
            for i in range(len(self.used_joint_indexes)):
                poses.append(np.interp(t,[0,1],[self.current_joint_poses[i],self.reset_joint_poses[i]]))
            self.actions.append((task,(poses,)))

        def output(): self.result = None,
        self.actions.append((output,()))
        pass


# coding=utf-8

import base64

if __name__=="__main__":
    result_list = []

    with open('in.txt', 'rb') as in_file:
        with open('result.txt', 'w') as out_file:
            for item in in_file:
                out_file.write(base64.b64encode(item).decode() + '\n')


import os
import cv2
import torch
import numpy as np
from PIL import Image
from advt.defence.defence import Defend

class TV(Defend):
    """
    TV: Total Variance
    """
    def __init__(self, model, device, test_dir, niters=30):
        """
        Initialize the TV class.

        Args:
            model: torch model
            device: torch.device
            test_dir: directory to be handled
            niters: int, iteration number
        """
        self.model = model
        self.device = device
        self.niters = niters
        self.test_dir = test_dir

    def defend_lbls(self):
        """
        Apply tv algorithm to the imgs of all labels in test_dir.

        """
        for file in os.listdir(self.test_dir):
            lbl_dir = os.path.join(self.test_dir, file)
            self.denoise_dir_imgs(lbl_dir)
            # self.denoise_img(lbl_dir)

    def defend_imgs(self):
        """
        Apply tv algorithm to the imgs in test_dir.

        """
        cnt = 0
        for file in os.listdir(self.test_dir):
            lbl_dir = os.path.join(self.test_dir, file)
            # self.denoise_dir_imgs(lbl_dir)
            self.denoise_img(lbl_dir)
            cnt += 1
            print('processed', cnt)

    def denoise_img(self, f_path):
        im = cv2.imread(f_path)
        cv2.cvtColor(src=im, code=cv2.COLOR_BGR2RGB, dst=im)
        im = im / 255
        de_im = self._tv(im) * 255

        # warning: image write back
        de_im = de_im.astype(np.uint8)
        de_im = Image.fromarray(de_im)
        de_im.save(f_path)

    def denoise_dir_imgs(self, path):
        """
        Denoise all images in specific path.

        Args:
            path: str, directory of images

        Returns:
            None

        """
        for f in os.listdir(path):
            f_path = os.path.join(path, f)
            im = cv2.imread(f_path)
            im = im / 255
            de_im = self._tv(im) * 255

            # warning: image write back
            de_im = de_im.astype(np.uint8)
            de_im = Image.fromarray(de_im)
            de_im.save(f_path)

    def _tv(self, image, weight=0.1, eps=2.e-4, n_iter_max=200):
        """
        Perform total variance minimization on a n-dim image.

        Args:
            image: ndarray, input image, shape [h, w, c], pixel value range [0, 1]
            weight: float
            eps: float
            n_iter_max: int

        Returns:
            out: ndarray, output image, shape [h, w, c], pixel value range [0, 1]

        """
        if image.dtype == np.uint8:
            image = image / 255
        ndim = image.ndim
        p = np.zeros((image.ndim,) + image.shape, dtype=image.dtype)
        g = np.zeros_like(p)
        d = np.zeros_like(image)
        i = 0
        while i < n_iter_max:
            if i > 0:
                # d will be the (negative) divergence of p
                d = -p.sum(0)
                slices_d = [slice(None), ] * ndim
                slices_p = [slice(None), ] * (ndim + 1)
                for ax in range(ndim):
                    slices_d[ax] = slice(1, None)
                    slices_p[ax + 1] = slice(0, -1)
                    slices_p[0] = ax
                    d[tuple(slices_d)] += p[tuple(slices_p)]
                    slices_d[ax] = slice(None)
                    slices_p[ax + 1] = slice(None)
                out = image + d
            else:
                out = image
            E = (d ** 2).sum()

            # g stores the gradients of out along each axis
            # e.g. g[0] is the first order finite difference along axis 0
            slices_g = [slice(None), ] * (ndim + 1)
            for ax in range(ndim):
                slices_g[ax + 1] = slice(0, -1)
                slices_g[0] = ax
                g[tuple(slices_g)] = np.diff(out, axis=ax)
                slices_g[ax + 1] = slice(None)

            norm = np.sqrt((g ** 2).sum(axis=0))[np.newaxis, ...]
            E += weight * norm.sum()
            tau = 1. / (2. * ndim)
            norm *= tau / weight
            norm += 1.
            p -= tau * g
            p /= norm
            E /= float(image.size)
            if i == 0:
                E_init = E
                E_previous = E
            else:
                if np.abs(E_previous - E) < eps * E_init:
                    break
                else:
                    E_previous = E
            i += 1
        return out


import pytest
import numpy as np

from model.elements import make_chromosome, make_individual, make_population, Individual

@pytest.fixture
def lists():
    first_list = np.array([0, 0, 0, 0, 1, 1, 1, 1])
    second_list = np.array([1, 1, 1, 1, 0, 0, 0, 0])
    third_list = np.array([1, 0, 1, 0, 1, 0, 1, 0])
    ones_list = np.array([1 for _ in range(8)])
    zeros_list = np.array([0 for _ in range(8)])
    return first_list, second_list, third_list, ones_list, zeros_list

@pytest.fixture
def chromosomes(lists):
    first_array, second_array, third_array, ones_array, zeros_array = lists
    first = make_chromosome(first_array)
    second = make_chromosome(second_array)
    third = make_chromosome(third_array)
    ones = make_chromosome(ones_array)
    zeros = make_chromosome(zeros_array)
    return first, second, third, ones, zeros

@pytest.fixture
def dicts(chromosomes):
    first, second, third, ones, zeros = chromosomes
    first_dict = {"x": first, "y": second, "z": third}
    second_dict = {"x": first, "y": second}
    third_dict = {"z": third}
    ones_dict = {"ones": ones}
    zeros_dict = {"zeros": zeros}
    return first_dict, second_dict, third_dict, ones_dict, zeros_dict

@pytest.fixture
def individuals(dicts):
    first_dict, second_dict, third_dict, ones_dict, zeros_dict = dicts
    first = make_individual(first_dict)
    second = make_individual(second_dict)
    third = make_individual(third_dict)
    ones = make_individual(ones_dict)
    zeros = make_individual(zeros_dict)
    return first, second, third, ones, zeros

@pytest.fixture
def individuals_lists(individuals):
    first, second, third, ones, zeros = individuals
    first_list = [first, second]
    second_list = [second, third]
    third_list = [third, ones, zeros]
    ones_list = [ones]
    zeros_list = [zeros]
    return first_list, second_list, third_list, ones_list, zeros_list

@pytest.fixture
def populations(individuals_lists):
    first_list, second_list, third_list, ones_list, zeros_list = individuals_lists
    first = make_population(*first_list)
    second = make_population(*second_list)
    third = make_population(*third_list)
    ones = make_population(*ones_list)
    zeros = make_population(*zeros_list)
    return first, second, third, ones, zeros

@pytest.fixture
def big_populations(individuals):
    first, second, third, ones, zeros = individuals

    def _make_population(individual):
        return make_population(*([individual for _ in range(25)] + [ones for _ in range(25)] + [zeros for _ in range(25)]))

    first_population = _make_population(first)
    second_population = _make_population(second)
    third_population = _make_population(third)
    ones_population = make_population(*([ones for _ in range(50)] + [first for _ in range(25)]))
    zeros_population = make_population(*([zeros for _ in range(50)] + [first for _ in range(25)]))

    return first_population, second_population, third_population, ones_population, zeros_population

@pytest.fixture
def small_population_and_individuals():
    first, second, ones = Individual(), Individual(), Individual()
    first["x"] = make_chromosome(np.asarray([1, 0, 1, 0, 1, 0, 1, 0]))
    second["x"] = make_chromosome(np.asarray([0, 1, 0, 1, 0, 1, 0, 0]))
    ones["x"] = make_chromosome(np.asarray([1, 1, 1, 1, 1, 1, 1, 1]))
    population = make_population(first, second)
    return population, first, second, ones


#
# sparkTest.py -- word count example to test installation
#

from pyspark import SparkContext

sc = SparkContext(appName="PythonWordCount")
#lines = sc.textFile("file:///home/code/climatology-gaussianInterp/README.md", 8)
lines = sc.textFile("file:///usr/share/dict/words", 128)
counts = lines.flatMap(lambda x: x.split(' ')) \
              .map(lambda x: (x, 1)) \
              .reduceByKey(lambda x,y: x + y)

output = counts.collect()
#for (word, count) in output:
#    print("%s: %i" % (word, count))


import aiohttp

from fastapi import APIRouter, Body
from ..schemas import STTResponse, STTRequest, TranslateRequest, TranslateResponse, \
    ObjectDetectionRequest, ObjectDetectionResponse

from .helpers import extract_text, google_speech_to_text, google_object_detection, google_translate

router = APIRouter()

@router.post("/object_detection", response_model=ObjectDetectionResponse)
async def detect_object_in_image(req_data: ObjectDetectionRequest = Body(...)):
    return await google_object_detection(req_data.file_url, req_data.headers)

@router.post("/text_detection", response_model=dict)
async def detect_text_in_image(req_data: ObjectDetectionRequest = Body(...)):
    async with aiohttp.ClientSession(headers=req_data.headers) as session:
        response = await session.get(req_data.file_url, allow_redirects=True)
        response.raise_for_status()
        content = await response.read()
    detected_text = extract_text(content)
    return {
        "detected_text": detected_text
    }

@router.post("/get_image_details", response_model=dict)
async def detect_text_and_object(req_data: ObjectDetectionRequest = Body(...)):
    async with aiohttp.ClientSession(headers=req_data.headers) as session:
        response = await session.get(req_data.file_url, allow_redirects=True)
        if response.status == 200:
            content = await response.read()
        else:
            return {
                "error": "Error occurred while retrieving image"
            }
    detected_text = extract_text(content)
    if not detected_text:
        detected_text = await google_object_detection(req_data.file_url, req_data.headers).get("detected_objects", [])
    return {
        "result": detected_text
    }

@router.post("/speech_to_text", response_model=STTResponse)
async def get_text_from_audio(req_data: STTRequest = Body(...)):
    result = await google_speech_to_text(req_data.file_url, req_data.headers, req_data.language_code)
    return {
        "detected_text": result
    }

@router.post("/translate", response_model=TranslateResponse)
def translate_text(req_data: TranslateRequest = Body(...)):
    data = {
        'contents': req_data.source_text,
        'target_language_code': req_data.target_language[0],
    }
    return google_translate(data)


# -*- coding: utf-8 -*-
# test/get_temp.py

def main():
    import serial
    import sys
    import datetime
    argvs = sys.argv

    if len(argvs) < 2:
        print('argument must be 2')
        port = '/dev/ttyUSB0'
    else:
        port = argvs[1]
    ser = serial.Serial(port, 38400, timeout=10)  # 10s-timeout baudrate=38400
    ser.write(b'<RM,>??\r\n')  # send command (CR+LF)
    line = ser.readline()
    line = str(datetime.datetime.now()) + str(",") + line.decode('utf-8')
    line = line.split(',')
    print(line[0])
    # print(line[1])
    print('velocity:'+line[2])
    velocity = float(line[2])
    # print(line[3])
    # print(line[4])
    print('temperature:'+line[5])
    temp = float(line[5])

    from mymodule import get_air_parameter, volumetric_flow_rate
    print(get_air_parameter(temp))
    print('specific_heat:')
    sh = get_air_parameter(temp, 'sh')
    print(sh)
    print('density:')
    d = get_air_parameter(temp, 'd')
    print(d)
    hc = get_air_parameter(temp, 'hc')
    print('heat_conductivity: ' + str(hc))

    # print(volumetric_flow_rate(v_wind=velocity,hole_area=0.001125))
    volume = volumetric_flow_rate(v_wind=velocity, hole_area=0.001125)*1
    print('volume is ' + str(volume))

    from mymodule import get_heat_air
    heat = get_heat_air(specific_heat=sh, density=d, volume=volume, temp_target=70, temp_now=temp)
    print('generated heat is ' + str(heat))

if __name__ == '__main__':
    main()


from __future__ import print_function

from torch.utils.data import Dataset
import os
import numpy as np
from tqdm import tqdm
import pickle
import time
import cv2
import random

class Preprocess(object):

    def __init__(self, format_size=112):
        super(Preprocess, self).__init__()
        self.format_size = format_size

    def _format(self, img_cv2):
        org_h, org_w = img_cv2.shape[0:2]
        rescale_ratio = self.format_size / max(org_h, org_w)
        h, w = int(org_h * rescale_ratio), int(org_w * rescale_ratio)
        img_rescaled = cv2.resize(img_cv2, (w, h))
        paste_pos = [int((self.format_size - w)/2), int((self.format_size - h)/2)]
        img_format = np.zeros((self.format_size, self.format_size, 3), dtype=np.uint8)
        img_format[paste_pos[1]:paste_pos[1]+h, paste_pos[0]:paste_pos[0]+w] = img_rescaled
        return img_format

    def _normalize(self, img_cv2):
        img_cv2 = cv2.cvtColor(img_cv2, cv2.COLOR_BGR2RGB)
        # mean = [123.675, 116.28, 103.53]
        # std = [58.395, 57.12, 57.375]
        mean = 127.5
        std = 127.5
        img_data = np.asarray(img_cv2, dtype=np.float32)
        img_data = img_data - mean
        img_data = img_data / std
        img_data = img_data.astype(np.float32)
        return img_data

    def __call__(self, img_cv2):
        img_format = self._format(img_cv2)
        img_data = self._normalize(img_format)
        img_data = np.transpose(img_data, axes=[2, 0, 1])
        return img_data

class TripletDataset(Dataset):

    def __init__(self,
                 finetune_dir='/data/FaceDetect/results/hegui_faces',
                 base_dir='/data/data/emore/released_data',
                 n_finetune_selected=0,
                 n_base_selected=20000,
                 finetune_ratio=0.0,
                 n_triplets=1000000):
        super(TripletDataset, self).__init__()
        self.finetune_dir = finetune_dir
        self.base_dir = base_dir
        self.n_finetune_seletected = n_finetune_selected
        self.n_base_selected = n_base_selected
        self.finetune_ratio = finetune_ratio
        self.n_triplets = n_triplets
        self.pickle_fpath = os.path.join(self.base_dir,
                                         'finetune_preload_{}_{}.pkl'.format(
                                             self.base_dir.split('/')[-1],
                                             n_base_selected,
                                         ))

        self.base_idx2fpathes = self._load_data_info()
        self.finetune_idx2fpathes = self._load_hegui_data_info()
        self.merge_idx2fpathes = {**self.base_idx2fpathes, **self.finetune_idx2fpathes}
        self.n_cls = len(self.merge_idx2fpathes)
        print('n_base:{}, n_finetune:{}, n_merge:{}'.format(self.n_base_selected, self.n_finetune_seletected, self.n_cls))
        self.triplet_groups = self._build_triplet_groups(self.merge_idx2fpathes)
        print('n_triplet_groups:{}'.format(len(self.triplet_groups)))
        self.preprocess_func = Preprocess()

    def __len__(self):
        return self.n_triplets

    def get_n_cls(self):
        return self.n_cls

    def _load_data_info(self):
        if os.path.exists(self.pickle_fpath):
            print('pickle file for loading emore data already exists, load it.')
            start_time = time.time()
            with open(self.pickle_fpath, 'rb') as f:
                idx2fpathes = pickle.load(f)
            end_time = time.time()
            time_spend = end_time - start_time
            print('load_success, time spend: ', '%.3f' % time_spend, ', n_identities: ', len(idx2fpathes))
        else:
            print('pickle file for loading emore data not exist, start reading original data...')
            idx2fpathes = self._preload_data(self.n_base_selected)
            if len(idx2fpathes) < 10:
                print('### ERROR: fail to load valid dataset !')
                exit(-1)
            # save load info
            with open(self.pickle_fpath, 'wb') as f:
                pickle.dump(idx2fpathes, f)
                print('pickle file for loading emore data has been generated.')
        return idx2fpathes

    def _preload_data(self, n_identities):
        ch_dir_list = [ch_dir
                       for ch_dir in os.listdir(self.base_dir)
                       if os.path.isdir(os.path.join(self.base_dir, ch_dir))]
        selected_dir_list = ch_dir_list[0: n_identities]
        print('n_chdir: {}'.format(len(selected_dir_list)))
        idx2fpathes = {}
        pbar = tqdm(total=len(selected_dir_list), desc='finetune_emore_preload')
        for ch_idx, ch_dir in enumerate(selected_dir_list):
            ch_dir_path = os.path.join(self.base_dir, ch_dir)
            sub_img_fpath_list = [os.path.join(ch_dir_path, img_fn)
                                  for img_fn in os.listdir(ch_dir_path)
                                  if img_fn.endswith('.jpg')]
            idx2fpathes[ch_idx] = sub_img_fpath_list
            pbar.update(1)
        pbar.close()
        return idx2fpathes

    def _load_hegui_data_info(self):
        register_dir = os.path.join(self.finetune_dir, 'register')
        test_dir = os.path.join(self.finetune_dir, 'test')
        register_img_fn_list = os.listdir(register_dir)[0: self.n_finetune_seletected]
        test_img_fn_list = os.listdir(test_dir)
        idx2fpathes = {}
        for idx, r_img_fn in enumerate(register_img_fn_list):
            t_img_fn = r_img_fn.replace('0000', '0001')
            if t_img_fn not in test_img_fn_list:
                print('{} not found.'.format(t_img_fn))
                exit(1)
            register = os.path.join(register_dir, r_img_fn)
            test = os.path.join(test_dir, t_img_fn)
            idx2fpathes[idx + self.n_base_selected] = [register, test]

        print('n_finetune_identities: {}'.format(len(idx2fpathes)))
        return idx2fpathes

    def _build_triplet_groups(self, idx2fpathes):
        triplet_groups = []
        for _ in range(self.n_triplets):
            if random.uniform(0.0, 1.0) < self.finetune_ratio:
                c1 = np.random.randint(self.n_base_selected, self.n_cls)
            else:
                c1 = np.random.randint(0, self.n_base_selected)
            c2 = np.random.randint(0, self.n_cls)
            while len(idx2fpathes[c1]) < 2:
                c1 = np.random.randint(0, self.n_cls - 1)
            while c1 == c2:
                c2 = np.random.randint(0, self.n_cls - 1)
            if len(idx2fpathes[c1]) == 2:  # hack to speed up process
                s1, s2 = 0, 1
            else:
                s1 = np.random.randint(0, len(idx2fpathes[c1]))
                s2 = np.random.randint(0, len(idx2fpathes[c1]))
                while s1 == s2:
                    s2 = np.random.randint(0, len(idx2fpathes[c1]))
            if len(idx2fpathes[c2]) == 1:
                s3 = 0
            else:
                s3 = np.random.randint(0, len(idx2fpathes[c2]))
            triplet_groups.append(
                [idx2fpathes[c1][s1], idx2fpathes[c1][s2], idx2fpathes[c2][s3], c1, c2])
        return triplet_groups

    def __getitem__(self, idx):
        anchor, positive, negtive, anc_cls_idx, neg_cls_idx = self.triplet_groups[idx]
        anc_img = self.preprocess_func(cv2.imread(anchor))
        pos_img = self.preprocess_func(cv2.imread(positive))
        neg_img = self.preprocess_func(cv2.imread(negtive))
        return anc_img, pos_img, neg_img, anc_cls_idx, neg_cls_idx

if __name__ == '__main__':

    dataset = TripletDataset()
    print(len(dataset))
    for sample in dataset:
        print(sample[3], sample[4])


import json
import multiprocessing as mp
import numpy as np
import scipy as sp
import os
import pickle
import pandas as pd
import pkg_resources
import time
import copy
from sklearn.model_selection import train_test_split
import tensorly as tl
import signal
from scipy.stats import mode
from contextlib import contextmanager

from . import convex_opt
from .linalg import pca, approx_matrix_rank, pivot_columns, impute, approx_tensor_rank, impute_with_coefficients
from .util import extract_columns, check_dataframes, tucker_on_error_tensor, generate_settings, error
from .pipeline import PipelineObject
from .model import Model
from .ensemble import Ensemble, Model_collection
from . import experiment_design as ED

class AutoLearner:
    """An object that automatically selects pipelines by greedy D-optimal design.

    Basic attributes:
        p_type (str):                  Problem type. One of {'classification', 'regression'}.
        algorithms (list):             A list of algorithm types to be considered, in strings. (e.g. ['KNN', 'lSVM']).
        hyperparameters (dict):        A nested dict of hyperparameters to be considered.
        n_folds (int):                 Number of cross-validation folds.
        verbose (bool):                Whether or not to generate print statements that showcase the progress.
        n_cores (int):                 Maximum number of cores over which to parallelize (None means no limit).
        runtime_limit (int):           Maximum training time for AutoLearner, in seconds.
        fit_ensemble_despite_timeout (bool): Whether to still fit the ensemble despite a fitting timeout.

    Advanced attributes:
        mode (str):                    The running mode of TensorOboe.
        new_row (np.ndarray):          Predicted row of matricized error tensor, corresponding to the new dataset. Default None.
        dataset_ratio_threshold(float):The threshold of dataset ratio (number of points / number of features) for dataset subsampling, if the training set is tall and skinny (number of data points much larger than number of features).
        runtime_predictor_algorithm (str):
                                       Model for runtime prediction. One of {'LinearRegression', 'KNeighborsRegressor'}.
        ensemble_method (str):         Ensemble method. One of {'greedy', 'stacking', 'best_several'}.
        load_default_tensor (bool):    Whether to use the default error and runtime tensors.
        selection_method (str):        Method of selecting entries of new row to sample. One of {'qr', 'min_variance', 'random'}.
        scalarization (str):           (Only in Oboe) Scalarization of the covariance matrix for mininum variance selection. One of {'D', 'A', 'E'}.

    Attributes to be added in the future:
        X, Y (np.ndarray):             PCA decomposition of error matrix.
        sampled_indices (set):         Indices of new row that have been sampled.
        sampled_models (list):         List of models that have been sampled (i.e. k-fold fitted).
        fitted_indices (set):          Indices of new row that have been fitted (i.e. included in ensemble)
        fitted_models (list):          List of models that have been fitted.
    """

    def __init__(self,
                 p_type='classification',
                 method='Oboe',
                 mode='warm',
                 algorithms=None,
                 hyperparameters=None,
                 verbose=False,
                 n_cores=1,
                 n_folds=5,
                 runtime_limit=512,
                 dataset_ratio_threshold=100,
                 new_row=None,
                 load_defaults=True,
                 customized_defaults_path='',
                 load_imputed_error_tensor=True,
                 original_error_tensor_dir=None,
                 path_to_imputed_error_tensor='default',
                 save_imputed_error_tensor=True,
                 selection_method='ED', scalarization='D',
                 build_ensemble=True,
                 load_saved_latent_factors=True,
                 save_latent_factors=True,
                 ensemble_method='best_several',
                 ensemble_max_size=5,
                 runtime_predictor_algorithm='LinearRegression',
                 load_saved_runtime_predictors=True,
                 save_fitted_runtime_predictors=False,
                 random_state=0,
                 fit_ensemble_despite_timeout=True,
                 **stacking_hyperparams):

        method = method.lower()
        assert method in {'oboe', 'tensoroboe'}, "The method must be one of {'Oboe', 'TensorOboe'}."
        self.method = method

        if method == 'tensoroboe':
            if mode == 'initialize':  # impute error tensor, compute low rank factors, fit runtime predictors
                load_imputed_error_tensor=False
                save_imputed_error_tensor=True
                load_saved_latent_factors=False
                load_saved_runtime_predictors=False
                save_fitted_runtime_predictors=True
            elif mode == 'impute':  # impute error tensor, compute low rank factors, use fitted runtime predictors
                load_imputed_error_tensor=False
                save_imputed_error_tensor=True
                load_saved_latent_factors=False
                load_saved_runtime_predictors=True
                save_fitted_runtime_predictors=False
            elif mode == 'factorize':  # used imputed error tensor, compute low rank factors, use fitted runtime predictors
                load_imputed_error_tensor=True
                save_imputed_error_tensor=False
                load_saved_latent_factors=False
                load_saved_runtime_predictors=True
                save_fitted_runtime_predictors=False
            elif mode == 'warm':  # used imputed error tensor, use low rank factors, use fitted runtime predictors
                load_imputed_error_tensor=True
                save_imputed_error_tensor=False
                load_saved_latent_factors=True
                load_saved_runtime_predictors=True
                save_fitted_runtime_predictors=False

        self.verbose = verbose
        self.random_state = random_state
        self.selection_method = selection_method
        self.p_type = p_type.lower()
        self.fit_ensemble_despite_timeout = fit_ensemble_despite_timeout

        if self.method == 'oboe':
            DEFAULTS = pkg_resources.resource_filename(__name__, 'defaults/Oboe')

            # TODO: check if arguments to constructor are valid; set to defaults if not specified
            assert selection_method in {'qr', 'ED', 'random'}, "The method to select entries to sample must be " \
                "either qr (QR decomposition), ED (minimize variance with time constraints), " \
                "or random (time-constrained random selection, i.e., set the models that should be able to finish within time budget)."

            # model configurations
            with open(os.path.join(DEFAULTS, p_type + '.json')) as file:
                defaults = json.load(file)

            # attributes of ML problem
            self.p_type = p_type.lower()
            self.algorithms = algorithms or defaults['algorithms']
            self.hyperparameters = hyperparameters or defaults['hyperparameters']
            self.verbose = verbose

            # computational considerations
            self.n_cores = n_cores
            self.runtime_limit = runtime_limit

            # configurations for column selection in cold-start
            self.selection_method = selection_method
            self.scalarization = scalarization

            ERROR_MATRIX = pd.read_csv(os.path.join(DEFAULTS, 'error_matrix.csv'), index_col=0, header=0)
            RUNTIME_MATRIX = pd.read_csv(os.path.join(DEFAULTS, 'runtime_matrix.csv'), index_col=0, header=0)

            # error matrix factorization
            # TODO: determine whether to generate new error matrix or use default/subset of default
            self.error_matrix = extract_columns(ERROR_MATRIX, self.algorithms, self.hyperparameters)
            self.runtime_matrix = extract_columns(RUNTIME_MATRIX, self.algorithms, self.hyperparameters)
            assert check_dataframes(self.error_matrix, self.runtime_matrix)
            self.column_headings = np.array([eval(heading) for heading in list(self.error_matrix)])
            self.X, self.Y, _ = pca(self.error_matrix.values, rank=min(self.error_matrix.shape)-1)

            # sampled & fitted models
            self.new_row = new_row or np.zeros((1, self.error_matrix.shape[1]))
            self.sampled_indices = set()
            self.sampled_models = [None] * self.error_matrix.shape[1]
            self.fitted_indices = set()
            self.fitted_models = [None] * self.error_matrix.shape[1]

            # ensemble attributes
            self.build_ensemble = build_ensemble
            self.ensemble_method = ensemble_method
            self.stacking_hyperparams = stacking_hyperparams
            if self.build_ensemble:
                self.ensemble = Ensemble(p_type=self.p_type, algorithm=self.ensemble_method, max_size=ensemble_max_size,
                                         hyperparameters=self.stacking_hyperparams, verbose=self.verbose, method='oboe')
            else:
                self.ensemble = Model_collection(self.p_type, method='oboe')

            # runtime predictor
            self.runtime_predictor_algorithm = runtime_predictor_algorithm

            self.dataset_ratio_threshold = dataset_ratio_threshold

        elif self.method == 'tensoroboe':
            if load_defaults:
                DEFAULTS = pkg_resources.resource_filename(__name__, 'defaults/TensorOboe')
            else:
                DEFAULTS = customized_defaults_path

            if not os.path.exists(os.path.join(DEFAULTS, 'tmp')):
                os.makedirs(os.path.join(DEFAULTS, 'tmp'))
                if self.verbose:
                    print("created a tmp directory in DEFAULTS ...")

            if not load_imputed_error_tensor:
                assert original_error_tensor_dir != None
                if self.verbose:
                    print("loading Float16 not-imputed error tensor ...")
                    # print("loading original Float64 error tensor ...")

                ERROR_TENSOR = np.float64(np.load(os.path.join(original_error_tensor_dir, 'error_tensor_f16_compressed.npz'))['a'])
                # ERROR_TENSOR = np.vstack((np.load(os.path.join(original_error_tensor_dir, 'error_tensor_part_1.npy')),
                #     np.load(os.path.join(original_error_tensor_dir, 'error_tensor_part_2.npy'))))

            RUNTIME_TENSOR = np.float64(np.load(os.path.join(DEFAULTS, 'runtime_tensor_f16_compressed.npz'))['a'])

            with open(os.path.join(DEFAULTS, 'training_index.pkl'), 'rb') as handle:
                TRAINING_INDEX = pickle.load(handle)

            with open(os.path.join(DEFAULTS, self.p_type + '.json')) as file:
                configs_unzipped = json.load(file)

            # attributes of ML problem
            self.algorithms = algorithms or configs_unzipped['estimator']['algorithms']
            self.hyperparameters = hyperparameters or configs_unzipped['estimator']['hyperparameters']

            # computational considerations
            if n_cores is None:
                n_cores = mp.cpu_count()
            self.n_cores = n_cores
            self.runtime_limit = runtime_limit
            self.n_folds = n_folds

            # error tensor completion
            rank_for_imputation = (20, 4, 2, 2, 8, 20)
            if verbose:
                print("rank for EM-Tucker imputation: {}".format(rank_for_imputation))
            rank_tuple = '-'.join([str(item) for item in rank_for_imputation])
            if load_imputed_error_tensor:
                try:
                    if path_to_imputed_error_tensor == 'default':
                        error_tensor_imputed = np.float64(np.load(os.path.join(DEFAULTS, 'error_tensor_imputed_20-4-2-2-8-20_f16_compressed.npz'))['a'])
                    else:
                        if self.verbose:
                            print("loading customized tensor at {} ...".format(path_to_imputed_error_tensor))
                        error_tensor_imputed = np.load(path_to_imputed_error_tensor)
                except:
                    print("Error loading imputed error tensor!")
            else:
                _, _, error_tensor_imputed, _ = tucker_on_error_tensor(ERROR_TENSOR, rank_for_imputation, save_results=False, verbose=self.verbose)
                if save_imputed_error_tensor:
                    imputed_error_tensor_save_path = os.path.join(DEFAULTS, 'tmp', 'error_tensor_imputed_{}.npy'.format(rank_tuple))
                    np.save(imputed_error_tensor_save_path, error_tensor_imputed)
                    if self.verbose:
                        print("saved imputed error tensor to {}".format(imputed_error_tensor_save_path))

            if verbose:
                print("shape of the error tensor: {}".format(error_tensor_imputed.shape))

            self.error_tensor_imputed = error_tensor_imputed

            # error tensor factorization
            # TODO: determine whether to generate new error matrix or use default/subset of default

            # the maximum ranks for dataset and estimator dimensions
            k_dataset_for_factorization = 30
            k_estimator_for_factorization = 30

            if load_saved_latent_factors:
                try:
                    U_t = np.load(os.path.join(DEFAULTS, 'tmp', 'error_tensor_U_t.npy'))
                    S_t = np.load(os.path.join(DEFAULTS, 'tmp', 'error_tensor_S_t.npy'))
                    Vt_t = np.load(os.path.join(DEFAULTS, 'tmp', 'error_tensor_Vt_t.npy'))
                    factorize_error_tensor = False
                except:
                    factorize_error_tensor = True
                    if self.verbose:
                        print("No saved latent factors. Factorizing the error tensor now ...")
            else:
                factorize_error_tensor = True

            if factorize_error_tensor:
                if self.verbose:
                    print("Factorizing the error matrix to get latent factors ...")
                core_tr, factors_tr = tl.decomposition.tucker(error_tensor_imputed, rank=(k_dataset_for_factorization, 4, 2, 2, 8, k_estimator_for_factorization))
                pipeline_latent_factors = tl.unfold(tl.tenalg.multi_mode_dot(core_tr, factors_tr[1:], modes=[1, 2, 3, 4, 5]), mode=0)
                U_t, S_t, Vt_t = sp.linalg.svd(pipeline_latent_factors, full_matrices=False)
                if save_latent_factors:
                    np.save(os.path.join(DEFAULTS, 'tmp', 'error_tensor_U_t.npy'), U_t)
                    np.save(os.path.join(DEFAULTS, 'tmp', 'error_tensor_S_t.npy'), S_t)
                    np.save(os.path.join(DEFAULTS, 'tmp', 'error_tensor_Vt_t.npy'), Vt_t)
                    if self.verbose:
                        print("latent factors saved to {}".format(os.path.join(DEFAULTS, 'tmp')))
            else:
                if self.verbose:
                    print("Loading latent factors from storage ...")

            # not yet implemented the part of selecting specific algorithms and hyperparameters
            self.error_matrix = tl.unfold(error_tensor_imputed, mode=0)
            self.runtime_matrix = tl.unfold(RUNTIME_TENSOR, mode=0)
            self.Y = Vt_t

            self.training_index = TRAINING_INDEX

            self.pipeline_settings = generate_settings(configs_unzipped)

             # sampled & fitted models
            self.new_row = new_row or np.full((1, len(self.pipeline_settings)), np.nan)
            self.new_row_pred = new_row or np.full((1, len(self.pipeline_settings)), np.nan)
            self.sampled_indices = set()
            self.sampled_pipelines = [None] * len(self.pipeline_settings)
            self.fitted_indices = set()
            self.fitted_pipelines = [None] * len(self.pipeline_settings)

            # ensemble attributes
            self.build_ensemble = build_ensemble
            self.ensemble_method = ensemble_method
            self.stacking_hyperparams = stacking_hyperparams
            if self.build_ensemble:
                self.ensemble = Ensemble(p_type=self.p_type, algorithm=self.ensemble_method, max_size=ensemble_max_size, hyperparameters=self.stacking_hyperparams, verbose=self.verbose, method='tensoroboe')
            else:
                self.ensemble = Model_collection(self.p_type, method='tensoroboe')

            # find the best-on-average pipeline
            self.best_on_average_pipeline_index = np.argmin(np.nanmean(tl.unfold(self.error_tensor_imputed, mode=0), axis=0))

            # runtime predictor
            # TODO: let it customizable
            self.runtime_predictor_algorithm = runtime_predictor_algorithm
            self.runtime_predictor = convex_opt.initialize_runtime_predictor(runtime_matrix=self.runtime_matrix, runtimes_index=self.training_index, model_name=self.runtime_predictor_algorithm, load=load_saved_runtime_predictors, save=save_fitted_runtime_predictors, method=self.method, verbose=self.verbose)
            self.dataset_ratio_threshold = dataset_ratio_threshold

    def _fit(self, x_train, y_train, categorical, t_predicted, ranks=None, runtime_limit=None, remaining_global=None):
        """This private method is a single round of the doubling process. It fits an AutoLearner object on a new dataset.
        This will sample the performance of several algorithms on the new dataset, predict performance on the rest, then construct an optimal ensemble model.

        Args:
            x_train (np.ndarray):  Features of the training dataset.
            y_train (np.ndarray):  Labels of the training dataset.
            t_predicted (np.ndarray): Predicted running time.
            ranks (int):            Rank of error tensor factorization.
            runtime_limit (float): Maximum time to allocate to AutoLearner fitting.
            remaining_global (float): The remaining runtime
        """

        if self.method == 'oboe':
            if self.verbose:
                print("\nSingle round runtime target: {}".format(runtime_limit))

            # set to defaults if not provided
            rank = ranks or approx_matrix_rank(self.error_matrix, threshold=0.01)
            runtime_limit = runtime_limit or self.runtime_limit

            if self.verbose:
                print('Fitting AutoLearner with max runtime {}s'.format(runtime_limit))

            # cold-start: pick the initial set of models to fit on the new dataset
            if self.selection_method == 'qr':
                to_sample = pivot_columns(self.error_matrix)
            elif self.selection_method == 'ED':
                # select algorithms to sample only from subset of algorithms that will run in allocated time
                valid = np.where(t_predicted <= self.n_cores * runtime_limit/2)[0]
                Y = self.Y[:rank, valid]
                # TODO: check if Y is rank-deficient, i.e. will ED problem fail
                v_opt = ED.solve(t_predicted[valid], runtime_limit/4, self.n_cores, Y, self.scalarization)
                to_sample = valid[np.where(v_opt > 0.9)[0]]
                if np.isnan(to_sample).any():
                    to_sample = np.argsort(t_predicted)[:rank]

            elif self.selection_method == 'random':
                to_sample = []
                # set of algorithms that are predicted to finish within given budget
                to_sample_candidates = np.where(t_predicted <= runtime_limit/2)[0]
                # remove algorithms that have been sampled already
                to_sample_candidates = list(set(to_sample_candidates) - self.sampled_indices)
                # if the remaining time is not sufficient for random sampling
                if len(to_sample_candidates) == 0:
                    to_sample = np.array([np.argmin(t_predicted)])
                else:
                    to_sample = np.random.choice(to_sample_candidates, min(self.n_cores, len(to_sample_candidates)), replace=False)
            else:
                to_sample = np.arange(0, self.new_row.shape[1])

            if len(to_sample) == 0 and len(self.sampled_indices) == 0:
                # if no columns are selected in first iteration (log det instability), sample n fastest columns
                n = len(np.where(np.cumsum(np.sort(t_predicted)) <= runtime_limit/4)[0])
                if n > 0:
                    to_sample = np.argsort(t_predicted)[:n]
                else:
                    self.ensemble.fitted = False
                    return

            start = time.time()
            if self.selection_method != 'random':
                # we only need to fit models on the new dataset if it has not been fitted already
                to_sample = list(set(to_sample) - self.sampled_indices)
                if self.verbose:
                    print('Sampling {} entries of new row...'.format(len(to_sample)))

                p1 = mp.Pool(self.n_cores)
                sample_models = [Model(self.p_type, self.column_headings[i]['algorithm'],
                                       self.column_headings[i]['hyperparameters'], self.verbose, i) for i in to_sample]
                sample_model_errors = [p1.apply_async(Model.kfold_fit_validate, args=[m, x_train, y_train, 5])
                                       for m in sample_models]
                p1.close()
                p1.join()

                # predict performance of models not actually fitted on the new dataset
                self.sampled_indices = self.sampled_indices.union(set(to_sample))
                for i, error in enumerate(sample_model_errors):
                    cv_error, cv_predictions = error.get()
                    sample_models[i].cv_error, sample_models[i].cv_predictions = cv_error.mean(), cv_predictions
                    sample_models[i].sampled = True
                    self.new_row[:, to_sample[i]] = cv_error.mean()
                    self.sampled_models[to_sample[i]] = sample_models[i]
                imputed = impute(self.error_matrix, self.new_row, list(self.sampled_indices), rank=rank)

                # impute ALL entries
                # unknown = sorted(list(set(range(self.new_row.shape[1])) - self.sampled_indices))
                # self.new_row[:, unknown] = imputed[:, unknown]

                # TODO: differentiate self.new_row and self.new_row_pred like TensorOboe
                self.new_row = imputed.copy()
                self.new_row_pred = imputed.copy()

                # k-fold fit candidate learners of ensemble
                remaining = (runtime_limit - (time.time()-start)) * self.n_cores
                # add models predicted to be the best to list of candidate learners to avoid empty lists
                best_sampled_idx = list(self.sampled_indices)[int(np.argmin(self.new_row[:, list(self.sampled_indices)]))]
                assert self.sampled_models[best_sampled_idx] is not None
                candidate_indices = [best_sampled_idx]
                self.ensemble.candidate_learners.append(self.sampled_models[best_sampled_idx])
                for i in np.argsort(self.new_row[0]):
                    if t_predicted[i] + t_predicted[candidate_indices].sum() <= remaining:
                        candidate_indices.append(i)
                        # last = candidate_indices.pop()
                        # assert last == best_sampled_idx
                        # candidate_indices.append(last)
                        # if model has already been k-fold fitted, immediately add to candidate learners
                        if i in self.sampled_indices:
                            assert self.sampled_models[i] is not None
                            self.ensemble.candidate_learners.append(self.sampled_models[i])
                # candidate learners that need to be k-fold fitted
                to_fit = list(set(candidate_indices) - self.sampled_indices)
            else:
                remaining = (runtime_limit - (time.time()-start)) * self.n_cores
                to_fit = to_sample.copy()

            # fit models predicted to have good performance and thus going to be added to the ensemble
            p2 = mp.Pool(self.n_cores)
            candidate_models = [Model(self.p_type, self.column_headings[i]['algorithm'],
                                      self.column_headings[i]['hyperparameters'], self.verbose, i) for i in to_fit]
            candidate_model_errors = [p2.apply_async(Model.kfold_fit_validate, args=[m, x_train, y_train, 5])
                                      for m in candidate_models]
            p2.close()
            p2.join()

            # update sampled indices
            self.sampled_indices = self.sampled_indices.union(set(to_fit))
            for i, error in enumerate(candidate_model_errors):
                cv_error, cv_predictions = error.get()
                candidate_models[i].cv_error, candidate_models[i].cv_predictions = cv_error.mean(), cv_predictions
                candidate_models[i].sampled = True
                self.new_row[:, to_fit[i]] = cv_error.mean()
                self.sampled_models[to_fit[i]] = candidate_models[i]
                self.ensemble.candidate_learners.append(candidate_models[i])
            # self.new_row = impute(self.error_matrix, self.new_row, list(self.sampled_indices), rank=rank)

            if self.verbose:
                print('\nFitting ensemble of max size {}...'.format(len(self.ensemble.candidate_learners)))
            # ensemble selection and fitting in the remaining time budget
            self.ensemble.fit(x_train, y_train, runtime_limit=remaining, fitted_base_learners=self.fitted_models)
            for model in self.ensemble.base_learners:
                assert model.index is not None
                self.fitted_indices.add(model.index)
                self.fitted_models[model.index] = model
            self.ensemble.fitted = True

            if self.verbose:
                print('\nAutoLearner fitting complete.')

        elif self.method == 'tensoroboe':
            if self.verbose:
                print("\nSingle round runtime target: {}".format(runtime_limit))

            # set to defaults if not provided
            ranks = ranks or approx_tensor_rank(self.error_tensor_imputed, threshold=0.01)
            runtime_limit = runtime_limit or self.runtime_limit

            if self.verbose:
                print('Fitting AutoLearner with maximum runtime {} seconds'.format(runtime_limit))

            # cold-start: pick the initial set of models to fit on the new dataset
            valid = np.where(t_predicted <= self.n_cores * runtime_limit/8)[0]
            Y = self.Y[:ranks[0], valid]

            if self.verbose:
                print("Selecting an initial set of models to evaluate ...")

            selected_columns_qr, t_sum, case = ED.pivot_columns_time(Y, t_predicted[valid], runtime_limit/8,
                                                            columns_to_avoid=None,
                                                            rank=Y.shape[0])
            if case == 'greedy_initialization':
                if self.verbose:
                    print(case)
                to_sample = valid[selected_columns_qr]

            elif case == 'qr_initialization':
                to_sample = valid[ED.greedy_stepwise_selection_with_time(Y=Y,
                                                            t=t_predicted[valid],
                                                            initialization=selected_columns_qr,
                                                            t_elapsed=t_sum,
                                                            t_max=runtime_limit/8,
                                                            idx_to_exclude=None)]
            # TODO: check if Y is rank-deficient, i.e. will ED problem fail

            if self.verbose:
                print(t_predicted[to_sample])

            if np.isnan(to_sample).any():
                to_sample = np.argsort(t_predicted)[:ranks[0]]

            if len(to_sample) == 0 and len(self.sampled_indices) == 0:
                # if no columns are selected in first iteration (log det instability), sample n fastest columns
                n = len(np.where(np.cumsum(np.sort(t_predicted)) <= runtime_limit/4)[0])
                if n > 0:
                    to_sample = np.argsort(t_predicted)[:n]
                else:
                    self.ensemble.fitted = False
                    return

            start = time.time()
            if self.selection_method != 'random':
                candidate_indices = []
                # we only need to fit models on the new dataset if it has not been fitted already
                to_sample = list(set(to_sample) - self.sampled_indices)
                if self.verbose:
                    print('Sampling {} entries of new row...'.format(len(to_sample)))

    #             return to_sample

                sampled_pipelines_single_round = [PipelineObject(p_type=self.p_type, config=self.pipeline_settings_on_dataset[i], index=i, verbose=self.verbose) for i in to_sample]

                p1 = mp.Pool(self.n_cores)
                sampled_pipeline_errors_single_round = [p1.apply_async(PipelineObject.kfold_fit_validate, args=[p, x_train, y_train, categorical, self.n_folds, runtime_limit/4, self.random_state]) for p in sampled_pipelines_single_round]
                p1.close()
                p1.join()

                if self.verbose:
                    print("pool fitting completed")

                # predict performance of models not actually fitted on the new dataset

                for i, error in enumerate(sampled_pipeline_errors_single_round):
                    cv_error, cv_predictions, t_elapsed = error.get()

                    if not np.isnan(cv_error):
                        sampled_pipelines_single_round[i].cv_error, sampled_pipelines_single_round[i].cv_predictions = cv_error, cv_predictions
                        sampled_pipelines_single_round[i].sampled = True
                        self.new_row[:, to_sample[i]] = cv_error
                        self.sampled_pipelines[to_sample[i]] = sampled_pipelines_single_round[i]
                        self.sampled_indices = self.sampled_indices.union(set([to_sample[i]]))
                        candidate_indices.append(to_sample[i])
                        self._t_predicted[to_sample[i]] = t_elapsed
                    else:
                        self._t_predicted[to_sample[i]] = max(t_elapsed, self._t_predicted[to_sample[i]])

                self.new_row_pred = impute_with_coefficients(self.Y[:ranks[0], :], self.new_row, list(self.sampled_indices))

                for idx in np.argsort(self.new_row[0, :])[:5]: # np.argsort automatically put nans at the end of the list
                    if not np.isnan(self.new_row[0, idx]):
                        if self.verbose:
                            print(self.sampled_pipelines[idx])
                        self.ensemble.candidate_learners.append(self.sampled_pipelines[idx])

                # currently disabled: impute ALL entries
                # unknown = sorted(list(set(range(self.new_row.shape[1])) - self.sampled_indices))
                # self.new_row[:, unknown] = imputed[:, unknown]

                # k-fold fit candidate learners of ensemble
                remaining = (runtime_limit - (time.time() - start)) * self.n_cores
                first = (case == 'qr_initialization') and not self.ever_once_selected_best

                if remaining > 0 or first:
                # add models predicted to be the best to list of candidate learners
                    if self.verbose:
                        if remaining < 0 and first:
                            print("Insufficient time in this doubling round, but we add models predicted to be the best at least once.")
                        print("length of sampled indices: {}".format(len(self.sampled_indices)))

        #             self.ensemble.candidate_learners.append(self.sampled_pipelines[best_sampled_idx])
                    for i in np.argsort(self.new_row_pred[0]):
                        if (first and len(candidate_indices) <= 3) or t_predicted[i] + t_predicted[candidate_indices].sum() <= remaining / 4:
                            candidate_indices.append(i)
                            # if model has already been k-fold fitted, immediately add to candidate learners
                            if i in self.sampled_indices:
                                assert self.sampled_pipelines[i] is not None
                                self.ensemble.candidate_learners.append(self.sampled_pipelines[i])
                    # candidate learners that need to be k-fold fitted
                    to_fit = list(set(candidate_indices) - self.sampled_indices)
                    if self.verbose:
                        print("{} candidate learners need to be k-fold fitted".format(to_fit))
                else:
                    if self.verbose:
                        print("Insufficient time in this doubling round.")
            else:
                remaining = (runtime_limit - (time.time() - start)) * self.n_cores
                to_fit = to_sample.copy()

            remaining = (runtime_limit - (time.time() - start)) * self.n_cores

            if remaining > 0 or first:
                if len(to_fit) > 0:
                    # fit models predicted to have good performance and thus going to be added to the ensemble
                    if self.verbose:
                        print("Fitting {} pipelines predicted to be the best ...".format(len(to_fit)))

                    candidate_pipelines = [PipelineObject(p_type=self.p_type, config=self.pipeline_settings_on_dataset[i], index=i, verbose=self.verbose) for i in to_fit]

    #                 print(remaining_global)
                    p2 = mp.Pool(self.n_cores)
                    candidate_pipeline_errors = [p2.apply_async(PipelineObject.kfold_fit_validate, args=[p, x_train, y_train, categorical, self.n_folds, remaining_global/2, self.random_state]) for p in candidate_pipelines] # set a not-quite-small limit for promising models
                    p2.close()
                    p2.join()

                    for i, error in enumerate(candidate_pipeline_errors):
                        cv_error, cv_predictions, t_elapsed = error.get()

                        if not np.isnan(cv_error):
                            candidate_pipelines[i].cv_error, candidate_pipelines[i].cv_predictions = cv_error, cv_predictions
                            candidate_pipelines[i].sampled = True
                            self.new_row[:, to_fit[i]] = cv_error
                            self.sampled_pipelines[to_fit[i]] = candidate_pipelines[i]
                            self.ensemble.candidate_learners.append(candidate_pipelines[i])
                            # update sampled indices
                            self.sampled_indices = self.sampled_indices.union(set([to_fit[i]]))
                            self._t_predicted[to_fit[i]] = t_elapsed
                        else:
                            self._t_predicted[to_fit[i]] = max(t_elapsed, self._t_predicted[to_fit[i]])

    #                 self.new_row = impute(self.error_matrix, self.new_row, list(self.sampled_indices), rank=rank)
                    self.ever_once_selected_best = True

            if self.verbose:
                print("Number of candidate learners in the ensemble: {}".format(len(self.ensemble.candidate_learners)))

            if len(self.ensemble.candidate_learners) > 0:
                self.ensemble.fitted = True

                if self.verbose:
                    print('\nFitting ensemble of maximum size {}...'.format(len(self.ensemble.candidate_learners)))
                # ensemble selection and fitting in the remaining time budget
                self.ensemble.fit(x_train, y_train, categorical, remaining, self.sampled_pipelines)
                for pipeline in self.ensemble.base_learners:
                    assert pipeline.index is not None
                    self.sampled_indices.add(pipeline.index)
                    self.sampled_pipelines[pipeline.index] = pipeline

                if self.verbose:
                    print('\nAutoLearner fitting complete.')
            else:
                if self.verbose:
                    print("Insufficient time in this round.")

    def fit(self, x_train, y_train, categorical=None):

        """Fit an AutoLearner object, iteratively doubling allowed runtime, and terminate when reaching the time limit."""

        if self.method == 'oboe':
            num_points, num_features = x_train.shape

            if self.verbose:
                print('\nShape of training dataset: {} data points, {} features'.format(num_points, num_features))

            if num_points > 10000 and num_points / num_features > self.dataset_ratio_threshold:
                num_points_new = int(min(5000, num_features * self.dataset_ratio_threshold))
                sampling_ratio = num_points_new / num_points
                print(sampling_ratio)
                df_x_train = pd.DataFrame(x_train)
                df_y_train = pd.DataFrame(y_train, columns=['labels'])
                df_resampled = df_x_train.join(df_y_train).groupby('labels').apply(pd.DataFrame.sample, frac=sampling_ratio).reset_index(drop=True)
                x_train = df_resampled.drop(['labels'], axis=1).values
                y_train = df_resampled['labels'].values
                if self.verbose:
                    print('\nTraining dataset resampled \nShape of resampled training dataset: {} data points, {} features'.format(x_train.shape[0], x_train.shape[1]))

            t_predicted = convex_opt.predict_runtime(x_train.shape, model_name=self.runtime_predictor_algorithm, method=self.method)

            # split data into training and validation sets
            try:
                x_tr, x_va, y_tr, y_va = train_test_split(x_train, y_train, test_size=0.15, stratify=y_train, random_state=0)
            except ValueError:
                x_tr, x_va, y_tr, y_va = train_test_split(x_train, y_train, test_size=0.15, random_state=0)

            ranks = [approx_matrix_rank(self.error_matrix, threshold=0.05)]
            if self.build_ensemble:
                t_init = 2**np.floor(np.log2(np.sort(t_predicted)[:int(1.1*ranks[0])].sum()))
                t_init = max(1, t_init)
            else:
                t_init = self.runtime_limit / 2
            times = [t_init]
            losses = [0.5]

            # predict runtime for the training set of the new dataset
            t_predicted = convex_opt.predict_runtime(x_train.shape, runtime_matrix=self.runtime_matrix, method='oboe')

            e_hat, e_hat_pred, actual_times, sampled, ensembles = [], [], [], [], []

            start = time.time()

            def doubling():
                k, t = ranks[0], times[0]
                counter, self.best = 0, 0
                while time.time() - start < self.runtime_limit - t:
                    if self.verbose:
                        print('Fitting with k={}, t={}'.format(k, t))
#                     if self.build_ensemble:
#                         self.ensemble = Ensemble(self.p_type, self.ensemble_method, self.stacking_hyperparams)
#                     else:
#                         self.ensemble = Model_collection(self.p_type)
                    self._fit(x_tr, y_tr, categorical=None, ranks=k, runtime_limit=t, t_predicted=t_predicted)
                    if self.build_ensemble and self.ensemble.fitted:
                        if self.verbose:
                            print("\nGot a new ensemble in the round with rumtime target {} seconds".format(t))
                        loss = error(y_va, self.ensemble.predict(x_va), self.p_type)

                        # TEMPORARY: Record intermediate results

                        e_hat.append(np.copy(self.new_row))
                        e_hat_pred.append(np.copy(self.new_row_pred))
                        actual_times.append(time.time() - start)
                        sampled.append(self.sampled_indices)
                        ensembles.append(self.ensemble)
                        losses.append(loss)

                        if loss == min(losses):
                            ranks.append(k+1)
                            self.best = counter
                        else:
                            ranks.append(k)

                        counter += 1

                    times.append(2*t)
                    k = ranks[-1]
                    t = times[-1]

            class TimeoutException(Exception): pass

            @contextmanager
            def time_limit(seconds):
                def signal_handler(signum, frame):
                    raise TimeoutException("Time limit reached.")
                signal.signal(signal.SIGALRM, signal_handler)
                signal.alarm(seconds)
                try:
                    yield
                finally:
                    signal.alarm(0)

            try:
                # set aside 3 seconds for initial and final processing steps
                with time_limit(self.runtime_limit - 3):
                    doubling()
            except TimeoutException as e:
                if self.verbose:
                    print("Time limit reached.")

        elif self.method == 'tensoroboe':
            num_points, num_features = x_train.shape

            self.ever_once_selected_best = False

            if self.verbose:
                print('\nShape of training dataset: {} data points, {} features'.format(num_points, num_features))

            if num_points > 10000 and num_points / num_features > self.dataset_ratio_threshold:
                num_points_new = int(min(5000, num_features * self.dataset_ratio_threshold))
                sampling_ratio = num_points_new / num_points
                if self.verbose:
                    print("dataset too skewed in having too many data points; sampling data points with sampling ratio {}".format(sampling_ratio))
                df_x_train = pd.DataFrame(x_train)
                df_y_train = pd.DataFrame(y_train, columns=['labels'])
                df_resampled = df_x_train.join(df_y_train).groupby('labels').apply(pd.DataFrame.sample, frac=sampling_ratio).reset_index(drop=True)
                x_train = df_resampled.drop(['labels'], axis=1).values
                y_train = df_resampled['labels'].values
                if self.verbose:
                    print('\nTraining dataset resampled \nShape of resampled training dataset: {} data points, {} features'.format(x_train.shape[0], x_train.shape[1]))

            if self.verbose:
                print("Splitting training set into training and validation ..")

            # split data into training and validation sets
    #         try:
    #             x_tr, x_va, y_tr, y_va = train_test_split(x_train, y_train, test_size=0.2, stratify=y_train, random_state=0)
    #         except ValueError:
    #             x_tr, x_va, y_tr, y_va = train_test_split(x_train, y_train, test_size=0.2, random_state=0)

            # for now, do not do validation, but instead do cross-validation
            x_tr = x_train
            x_va = x_train
            y_tr = y_train
            y_va = y_train

            num_points_tr, num_features_tr = x_tr.shape

            self._t_predicted = np.maximum(self._predict_runtime(x_tr), 0.5)

            def p2f(x):
                return float(x.strip('%'))/100

            # pipeline settings on dataset
            self.pipeline_settings_on_dataset = []

            for item in self.pipeline_settings:
                item_copy = copy.deepcopy(item)
                if item_copy['dim_reducer']['algorithm'] == 'PCA':
                    item_copy['dim_reducer']['hyperparameters']['n_components'] = int(min((self.n_folds - 1) * num_points_tr/self.n_folds, p2f(item_copy['dim_reducer']['hyperparameters']['n_components']) * num_features_tr))
                elif item_copy['dim_reducer']['algorithm'] == 'SelectKBest':
                    item_copy['dim_reducer']['hyperparameters']['k'] = int(p2f(item_copy['dim_reducer']['hyperparameters']['k']) * num_features_tr)

                self.pipeline_settings_on_dataset.append(item_copy)

            start = time.time()

            ranks = [approx_tensor_rank(self.error_tensor_imputed, threshold=0.05)]

            t_init = 2**np.floor(np.log2(np.sort(self._t_predicted)[:int(5*ranks[0][-1])].sum()))
    #         t_init = max(2**np.floor(np.log2(self.runtime_limit/2)), t_init) # better on large datasets
            t_init = max(2**np.floor(np.log2(self.runtime_limit/8)), t_init)

            if self.verbose:
                print("runtime limit of initial round: {} seconds".format(t_init))
            times = [t_init]
            losses = [0.5]

            # fit and validate the best-on-average pipeline
            if self.verbose:
                print("fitting and kfold_fit_validating the best-on-average pipeline")
            best_on_average_pipeline = PipelineObject(
                p_type=self.p_type, config=self.pipeline_settings_on_dataset[self.best_on_average_pipeline_index],
                index=self.best_on_average_pipeline_index, verbose=self.verbose)
            self.ensemble.base_learners.append(best_on_average_pipeline)
            self.ensemble.fit(x_tr, y_tr, categorical)
            loss = self.ensemble.kfold_fit_validate(x_va, y_va, categorical, n_folds=self.n_folds, timeout=t_init)[0]
            losses.append(loss)

            e_hat, e_hat_pred, actual_times, sampled, ensembles = [], [], [], [], []

            if self.verbose:
                print("Doubling process started ...")
            def doubling():
                k, t = ranks[0], times[0]
                counter, self.best = 0, 0
                while time.time() - start < self.runtime_limit - t:
                    if self.verbose:
                        print('Fitting with ranks={}, t={}'.format(k, t))
    #                 if self.build_ensemble:
    #                     self.ensemble = Ensemble(self.p_type, self.ensemble_method, self.stacking_hyperparams)
    #                 else:
    #                     self.ensemble = Model_collection(self.p_type)
                    self._fit(x_tr, y_tr, categorical, self._t_predicted, ranks=k, runtime_limit=t, remaining_global=self.runtime_limit-time.time()+start)
                    if self.build_ensemble and self.ensemble.fitted:
                        if self.verbose:
                            print("\nGot a new ensemble in the round with runtime target {} seconds".format(t))
    #                     loss = error(y_va, self.ensemble.predict(x_va), self.p_type)
                        loss = self.ensemble.kfold_fit_validate(x_va, y_va, categorical, n_folds=self.n_folds, timeout=(self.runtime_limit-time.time()+start)/2)[0]

                        # TEMPORARY: Record intermediate results
                        e_hat.append(np.copy(self.new_row))
                        e_hat_pred.append(np.copy(self.new_row_pred))
                        actual_times.append(time.time() - start)
                        sampled.append(self.sampled_indices)
                        ensembles.append(self.ensemble)
                        losses.append(loss)

                        if loss == min(losses):
                            rank_new = (k[0], k[1], k[2], k[3], k[4], k[5]+1)
                            if self.verbose:
                                print("new approximate rank for the error tensor: {}".format(rank_new))
                            ranks.append(rank_new)
                            self.best = counter
                        else:
                            ranks.append(k)

                        counter += 1

                    times.append(2*t)
                    k = ranks[-1]
                    t = times[-1]

            class TimeoutException(Exception): pass

            @contextmanager
            def time_limit(seconds):
                def signal_handler(signum, frame):
                    raise TimeoutException("Time limit reached.")
                signal.signal(signal.SIGALRM, signal_handler)
                signal.alarm(seconds)
                try:
                    yield
                finally:
                    signal.alarm(0)

            try:
                # set aside 3 seconds for initial and final processing steps
                with time_limit(self.runtime_limit - 3):
                    doubling()
            except TimeoutException as e:
                if self.verbose:
                    print("Time limit reached.")

        if not self.ensemble.fitted:
            if self.build_ensemble and self.fit_ensemble_despite_timeout and self.verbose:
                print('\n Ensemble not fitted in the time limit. Fitting ensemble of max size {}...'.format(len(self.ensemble.candidate_learners)))
            # ensemble selection and fitting in the remaining time budget
            if (not self.build_ensemble) or (self.build_ensemble and self.fit_ensemble_despite_timeout):
                self.ensemble.fit(x_train, y_train, runtime_limit=None, fitted_base_learners=self.fitted_models)
                for model in self.ensemble.base_learners:
                    assert model.index is not None
                    self.fitted_indices.add(model.index)
                    self.fitted_models[model.index] = model
                self.ensemble.fitted = True

        if self.build_ensemble:
            # after all iterations, restore best model

            try:
                self.new_row = e_hat[self.best]
                self.ensemble = ensembles[self.best]
                self.predict_with_most_common_class = False

                return {'ranks': ranks[:-1], 'runtime_limits': times[:-1], 'validation_loss': losses,
                    'filled_new_row': e_hat, 'predicted_new_row': e_hat_pred, 'actual_runtimes': actual_times, 'sampled_indices': sampled, 'models': ensembles}
            except IndexError:
                self.y_train_mode = mode(y_train.flatten())[0][0]
                if not self.ensemble.fitted:
                    print("No ensemble built within time limit. The auto-learner is set to predict by the mode of class labels in the training set. Please try increasing the time limit or allocate more computational resources.")
                    self.predict_with_most_common_class = True
                else:
                    self.predict_with_most_common_class = False

    def predict(self, x_test):
        """Generate predictions on test data.

        Args:
            x_test (np.ndarray): Features of the test dataset.
        Returns:
            np.ndarray: Predicted labels.
        """
        if self.build_ensemble:
            if self.predict_with_most_common_class:
                return np.full(x_test.shape[0], self.y_train_mode)
            else:
                return self.ensemble.predict(x_test)
        else: # just select a collection of promising models
            return self.ensemble.predict(x_test) # the self.ensemble object here is a Model_collection

    def _predict_runtime(self, x_train):
        # predict runtime for the training set of the new dataset.
        if self.verbose:
            print("Predicting pipeline running time ..")
        return convex_opt.predict_runtime(x_train.shape, saved_model='Class', model=self.runtime_predictor)

# The code below is deprecated
    def refit(self, x_train, y_train):
        """Refit an existing AutoLearner object on a new dataset. This will simply retrain the base-learners and
        stacked learner of an existing model, and so algorithm and hyperparameter selection may not be optimal.
        Args:
            x_train (np.ndarray): Features of the training dataset.
            y_train (np.ndarray): Labels of the training dataset.
        """
        assert self.ensemble.fitted, "Cannot refit unless model has been fit."
        self.ensemble.refit(x_train, y_train)

    def get_models(self):
        """Get details of the selected machine learning models and the ensemble.
        """
        return self.ensemble.get_models()

    def get_pipeline_accuracies(self, y_test):
        """
        Get accuracies of selected pipelines.
        """
        return self.ensemble.get_pipeline_accuracies(y_test)

    def get_model_accuracies(self, y_test):
        """
        Get accuracies of selected pipelines.
        """
        return self.get_pipeline_accuracies(y_test)

    def _greedy_initial_selection(self, x_train, y_train, t_predicted, runtime_limit):
        if self.verbose:
            print("Fitting fast pipelines that perform well on average.")

        candidate_pipeline_indices = list(set(np.where(t_predicted <= runtime_limit / 8)[0]).intersection(
            set(np.argsort(np.nanmedian(tl.unfold(self.error_tensor_imputed, mode=0), axis=0))[
                :int(len(self.pipeline_settings_on_dataset)/50)])))

        candidate_pipelines = [PipelineObject(p_type=self.p_type, config=self.pipeline_settings_on_dataset[i], index=i, verbose=self.verbose) for i in candidate_pipeline_indices]

        p1 = mp.Pool(self.n_cores)
        candidate_pipeline_errors = [p1.apply_async(PipelineObject.kfold_fit_validate, args=[p, x_train, y_train, self.n_folds, runtime_limit/4, self.random_state]) for p in candidate_pipelines]
        p1.close()
        p1.join()

        if self.verbose:
            print("Initial greedy fitting completed")

        for i, error in enumerate(candidate_pipeline_errors):
            cv_error, cv_predictions, t_elapsed = error.get()
            if not np.isnan(cv_error):
                candidate_pipelines[i].cv_error, candidate_pipelines[i].cv_predictions = cv_error, cv_predictions
                candidate_pipelines[i].sampled = True
                self.new_row[:, candidate_pipeline_indices[i]] = cv_error
                self.sampled_pipelines[candidate_pipeline_indices[i]] = candidate_pipelines[i]
                self.ensemble.candidate_learners.append(candidate_pipelines[i])
                # update sampled indices
                self.sampled_indices = self.sampled_indices.union(set([candidate_pipeline_indices[i]]))
                self._t_predicted[candidate_pipeline_indices[i]] = t_elapsed
            else:
                self._t_predicted[candidate_pipeline_indices[i]] = max(t_elapsed, self._t_predicted[candidate_pipeline_indices[i]])

        if len(self.ensemble.candidate_learners) > 0:
            self.ensemble.fitted = True
            self.ensemble.fit(x_train, y_train)
        else:
            if self.verbose:
                print("Insufficient time to fit fast and on average best-performing pipelines.")