# Consolidated Cassie environment.

from .cassiemujoco import pd_in_t, state_out_t, CassieSim, CassieVis

from .trajectory import CassieTrajectory, getAllTrajectories
from cassie.quaternion_function import *
from .rewards import *

from math import floor

import numpy as np
import os
import random
import copy

import pickle

class CassieEnv_noaccel_footdist_nojoint:
    def __init__(self, traj='walking', simrate=60, clock_based=True, state_est=True, dynamics_randomization=True, no_delta=True, reward="iros_paper", history=0):
        self.sim = CassieSim("./cassie/cassiemujoco/cassie.xml")
        self.vis = None

        self.reward_func = reward

        self.clock_based = clock_based
        self.state_est = state_est
        self.no_delta = no_delta
        self.dynamics_randomization = dynamics_randomization

        # CONFIGURE REF TRAJECTORY to use
        if traj == "aslip":
            self.speeds = np.array([x / 10 for x in range(0, 21)])
            self.trajectories = getAllTrajectories(self.speeds)
            self.num_speeds = len(self.trajectories)
            self.speed = self.speeds[0]
            self.trajectory = self.trajectories[0]
            self.aslip_traj = True
        else:
            self.aslip_traj = False
            dirname = os.path.dirname(__file__)
            if traj == "walking":
                traj_path = os.path.join(dirname, "trajectory", "stepdata.bin")
            elif traj == "stepping":
                traj_path = os.path.join(dirname, "trajectory", "more-poses-trial.bin")
            self.trajectory = CassieTrajectory(traj_path)
            self.speed = 0

        self.observation_space, self.clock_inds, self.mirrored_obs = self.set_up_state_space()

        # Adds option for state history for FF nets
        self._obs = len(self.observation_space)
        self.history = history
        self.state_history = [np.zeros(self._obs) for _ in range(self.history+1)]

        self.observation_space = np.zeros(self._obs + self._obs * self.history)
        self.action_space = np.zeros(10)

        self.P = np.array([100,  100,  88,  96,  50])
        self.D = np.array([10.0, 10.0, 8.0, 9.6, 5.0])

        self.u = pd_in_t()

        # TODO: should probably initialize this to current state
        self.cassie_state = state_out_t()

        self.simrate = simrate # simulate X mujoco steps with same pd target
                                # 60 brings simulation from 2000Hz to roughly 30Hz

        self.time    = 0 # number of time steps in current episode
        self.phase   = 0 # portion of the phase the robot is in
        self.counter = 0 # number of phase cycles completed in episode

        # NOTE: a reference trajectory represents ONE phase cycle

        # should be floor(len(traj) / simrate) - 1
        # should be VERY cautious here because wrapping around trajectory
        # badly can cause assymetrical/bad gaits
        self.phaselen = floor(len(self.trajectory) / self.simrate) - 1 if not self.aslip_traj else self.trajectory.length - 1

        # see include/cassiemujoco.h for meaning of these indices
        self.pos_idx = [7, 8, 9, 14, 20, 21, 22, 23, 28, 34]
        self.vel_idx = [6, 7, 8, 12, 18, 19, 20, 21, 25, 31]

        self.pos_index = np.array([1,2,3,4,5,6,7,8,9,14,15,16,20,21,22,23,28,29,30,34])
        self.vel_index = np.array([0,1,2,3,4,5,6,7,8,12,13,14,18,19,20,21,25,26,27,31])

        # CONFIGURE OFFSET for No Delta Policies
        if self.aslip_traj:
            ref_pos, ref_vel = self.get_ref_state(self.phase)
            self.offset = ref_pos[self.pos_idx]
        else:
            self.offset = np.array([0.0045, 0.0, 0.4973, -1.1997, -1.5968, 0.0045, 0.0, 0.4973, -1.1997, -1.5968])

        self.phase_add = 1

        # global flat foot orientation, can be useful part of reward function:
        self.neutral_foot_orient = np.array([-0.24790886454547323, -0.24679713195445646, -0.6609396704367185, 0.663921021343526])
        self.avg_lfoot_quat = np.zeros(4)
        self.avg_rfoot_quat = np.zeros(4)

        #### Dynamics Randomization ####
        # self.dynamics_randomization = False
        self.slope_rand = False
        self.joint_rand = False
        # Record default dynamics parameters
        if self.dynamics_randomization:
            self.default_damping = self.sim.get_dof_damping()
            self.default_mass = self.sim.get_body_mass()
            self.default_ipos = self.sim.get_body_ipos()
            self.default_fric = self.sim.get_geom_friction()

            weak_factor = 0.7
            strong_factor = 1.3

            pelvis_damp_range = [[self.default_damping[0], self.default_damping[0]],
                                [self.default_damping[1], self.default_damping[1]],
                                [self.default_damping[2], self.default_damping[2]],
                                [self.default_damping[3], self.default_damping[3]],
                                [self.default_damping[4], self.default_damping[4]],
                                [self.default_damping[5], self.default_damping[5]]]

            hip_damp_range = [[self.default_damping[6]*weak_factor, self.default_damping[6]*strong_factor],
                            [self.default_damping[7]*weak_factor, self.default_damping[7]*strong_factor],
                            [self.default_damping[8]*weak_factor, self.default_damping[8]*strong_factor]]  # 6->8 and 19->21

            achilles_damp_range = [[self.default_damping[9]*weak_factor,  self.default_damping[9]*strong_factor],
                                    [self.default_damping[10]*weak_factor, self.default_damping[10]*strong_factor],
                                    [self.default_damping[11]*weak_factor, self.default_damping[11]*strong_factor]] # 9->11 and 22->24

            knee_damp_range     = [[self.default_damping[12]*weak_factor, self.default_damping[12]*strong_factor]]   # 12 and 25
            shin_damp_range     = [[self.default_damping[13]*weak_factor, self.default_damping[13]*strong_factor]]   # 13 and 26
            tarsus_damp_range   = [[self.default_damping[14], self.default_damping[14]]]             # 14 and 27
            heel_damp_range     = [[self.default_damping[15], self.default_damping[15]]]                           # 15 and 28
            fcrank_damp_range   = [[self.default_damping[16]*weak_factor, self.default_damping[16]*strong_factor]]   # 16 and 29
            prod_damp_range     = [[self.default_damping[17], self.default_damping[17]]]                           # 17 and 30
            foot_damp_range     = [[self.default_damping[18]*weak_factor, self.default_damping[18]*strong_factor]]   # 18 and 31

            side_damp = hip_damp_range + achilles_damp_range + knee_damp_range + shin_damp_range + tarsus_damp_range + heel_damp_range + fcrank_damp_range + prod_damp_range + foot_damp_range
            self.damp_range = pelvis_damp_range + side_damp + side_damp

            hi = 1.1
            lo = 0.9
            m = self.default_mass
            pelvis_mass_range      = [[lo*m[1],  hi*m[1]]]  # 1
            hip_mass_range         = [[lo*m[2],  hi*m[2]],  # 2->4 and 14->16
                                    [lo*m[3],  hi*m[3]],
                                    [lo*m[4],  hi*m[4]]]

            achilles_mass_range    = [[lo*m[5],  hi*m[5]]]  # 5 and 17
            knee_mass_range        = [[lo*m[6],  hi*m[6]]]  # 6 and 18
            knee_spring_mass_range = [[lo*m[7],  hi*m[7]]]  # 7 and 19
            shin_mass_range        = [[lo*m[8],  hi*m[8]]]  # 8 and 20
            tarsus_mass_range      = [[lo*m[9],  hi*m[9]]]  # 9 and 21
            heel_spring_mass_range = [[lo*m[10], hi*m[10]]] # 10 and 22
            fcrank_mass_range      = [[lo*m[11], hi*m[11]]] # 11 and 23
            prod_mass_range        = [[lo*m[12], hi*m[12]]] # 12 and 24
            foot_mass_range        = [[lo*m[13], hi*m[13]]] # 13 and 25

            side_mass = hip_mass_range + achilles_mass_range \
                        + knee_mass_range + knee_spring_mass_range \
                        + shin_mass_range + tarsus_mass_range \
                        + heel_spring_mass_range + fcrank_mass_range \
                        + prod_mass_range + foot_mass_range

            self.mass_range = [[0, 0]] + pelvis_mass_range + side_mass + side_mass

            self.damp_noise = np.zeros(len(self.damp_range))
            self.mass_noise = np.zeros(len(self.mass_range))
            self.fric_noise = np.zeros(3)

        # self.delta_x_min, self.delta_x_max = self.default_ipos[3] - 0.05, self.default_ipos[3] + 0.05
        # self.delta_y_min, self.delta_y_max = self.default_ipos[4] - 0.05, self.default_ipos[4] + 0.05

        ### Trims ###
        self.joint_offsets = np.zeros(14)
        self.com_vel_offset = 0
        self.y_offset = 0

        ### Random commands during training ###
        self.speed_schedule = np.zeros(4)
        self.orient_add = 0

        # Keep track of actions, torques
        self.prev_action = np.zeros(10)
        self.curr_action = None
        self.prev_torque = None

        # for RNN policies
        self.critic_state = None

        # Reward terms
        self.l_foot_orient = 0
        self.r_foot_orient = 0
        self.neutral_foot_orient = np.array([-0.24790886454547323, -0.24679713195445646, -0.6609396704367185, 0.663921021343526])
        self.l_high = False
        self.r_high = False
        self.lfoot_vel = np.zeros(3)
        self.rfoot_vel = np.zeros(3)
        self.l_foot_cost = 0
        self.r_foot_cost = 0
        self.l_foot_cost_even = 0
        self.r_foot_cost_even = 0
        self.l_foot_cost_smooth = 0
        self.r_foot_cost_smooth = 0
        self.smooth_cost = 0
        self.torque_cost = 0
        self.prev_torque = None

        self.debug = False

    def set_up_state_space(self):

        mjstate_size   = 40
        state_est_size = 36

        speed_size     = 1

        clock_size    = 2

        # TODO: When setup foot position for non state est, fix mirror obs as well
        # Find the mirrored trajectory
        if self.aslip_traj:
            ref_traj_size = 18
            mirrored_traj = [6,7,8,9,10,11,0.1,1,2,3,4,5,12,13,14,15,16,17]
        else:
            ref_traj_size = 40
            mirrored_traj = np.array([0.1, 1, 2, 3, 4, 5, -13, -14, 15, 16, 17, 18, 19, -6, -7, 8, 9, 10, 11, 12,
                                        20, 21, 22, 23, 24, 25, -33, -34, 35, 36, 37, 38, 39, -26, -27, 28, 29, 30, 31, 32])

        if self.state_est:
            base_mir_obs = np.array([ 3, 4, 5, 0.1, 1, 2, 6, -7, 8, -9, -15, -16, 17, 18, 19, -10, -11, 12, 13, 14, 20, -21, 22,
                                    -23, 24, -25, -31, -32, 33, 34, 35, -26, -27, 28, 29, 30])
            # base_mir_obs = np.array([ 3, 4, 5, 0.1, 1, 2, 6, 7, 8, 9, -15, -16, 17, 18, 19, -10, -11, 12, 13, 14, 20, 21, 22,
            #                         23, 24, 25, -31, -32, 33, 34, 35, -26, -27, 28, 29, 30, 36, 37, 38,
            #                         42, 43, 44, 39, 40, 41, 48, 49, 50, 45, 46, 47])
            obs_size = state_est_size#len(base_mir_obs)
        else:
            base_mir_obs = np.array([0.1, 1, 2, -3, 4, -5, -13, -14, 15, 16, 17, 18, 19, -6, -7, 8, 9, 10, 11, 12, 20, -21, 22, -23, 24, -25, -33, -34, 35, 36, 37, 38, 39, -26, -27, 28, 29, 30, 31, 32])
            obs_size = mjstate_size
        if self.clock_based:
            append_obs = np.array([len(base_mir_obs) + i for i in range(clock_size+speed_size)])
            mirrored_obs = np.concatenate([base_mir_obs, append_obs])
            clock_inds = append_obs[0:clock_size].tolist()
            obs_size += clock_size + speed_size
        else:
            mirrored_traj_sign = np.multiply(np.sign(mirrored_traj), obs_size+np.floor(np.abs(mirrored_traj)))
            mirrored_obs = np.concatenate([base_mir_obs, mirrored_traj_sign])
            clock_inds = None
            obs_size += ref_traj_size

        # NOTE: mirror loss only set up for clock based with state estimation so far.
        observation_space = np.zeros(obs_size)

        check_arr = np.arange(obs_size, dtype=np.float64)
        check_arr[0] = 0.1
        # print("mir obs check: ", np.all(np.sort(np.abs(mirrored_obs)) == check_arr))
        # print("mir obs check: ", np.sort(np.abs(mirrored_obs)))
        # print("mir obs check: ", np.sort(np.abs(mirrored_obs)) == check_arr)

        mirrored_obs = mirrored_obs.tolist()
        # print("mirrored_obs: ", mirrored_obs)
        # print("mir obs len: ", len(mirrored_obs))
        # print("obs_size: ", obs_size)

        # exit()
        # observation_space = np.concatenate([observation_space, np.zeros(32+26+3)])

        return observation_space, clock_inds, mirrored_obs

    def step_simulation(self, action):

        if self.aslip_traj and self.phase == self.phaselen - 1:
            ref_pos, ref_vel = self.get_ref_state(0)
        else:
            # maybe make ref traj only send relevant idxs?
            ref_pos, ref_vel = self.get_ref_state(self.phase + self.phase_add)

        if not self.no_delta:
            target = action + ref_pos[self.pos_idx]
        else:
            target = action + self.offset

        if self.joint_rand:
            target -= self.joint_offsets[0:10]

        curr_fpos = np.zeros(6)
        self.sim.foot_pos(curr_fpos)
        prev_foot = copy.deepcopy(curr_fpos)
        self.u = pd_in_t()
        for i in range(5):
            # TODO: move setting gains out of the loop?
            # maybe write a wrapper for pd_in_t ?
            self.u.leftLeg.motorPd.pGain[i]  = self.P[i]
            self.u.rightLeg.motorPd.pGain[i] = self.P[i]

            self.u.leftLeg.motorPd.dGain[i]  = self.D[i]
            self.u.rightLeg.motorPd.dGain[i] = self.D[i]

            self.u.leftLeg.motorPd.torque[i]  = 0 # Feedforward torque
            self.u.rightLeg.motorPd.torque[i] = 0

            self.u.leftLeg.motorPd.pTarget[i]  = target[i]
            self.u.rightLeg.motorPd.pTarget[i] = target[i + 5]

            self.u.leftLeg.motorPd.dTarget[i]  = 0
            self.u.rightLeg.motorPd.dTarget[i] = 0

        self.cassie_state = self.sim.step_pd(self.u)
        self.sim.foot_pos(curr_fpos)
        self.lfoot_vel = (curr_fpos[0:3] - prev_foot[0:3]) / 0.0005
        self.rfoot_vel = (curr_fpos[3:6] - prev_foot[3:6]) / 0.0005
        foot_forces = self.sim.get_foot_forces()
        if self.l_high and foot_forces[0] > 0:
            self.l_high = False
        elif not self.l_high and curr_fpos[2] >= 0.19:
            self.l_high = True
        if self.r_high and foot_forces[1] > 0:
            self.r_high = False
        elif not self.r_high and curr_fpos[5] >= 0.19:
            self.r_high = True

    def step_sim_basic(self, action):

        if self.aslip_traj and self.phase == self.phaselen - 1:
            ref_pos, ref_vel = self.get_ref_state(0)
        else:
            # maybe make ref traj only send relevant idxs?
            ref_pos, ref_vel = self.get_ref_state(self.phase + self.phase_add)

        if not self.no_delta:
            target = action + ref_pos[self.pos_idx]
        else:
            target = action + self.offset

        if self.joint_rand:
            target -= self.joint_offsets[0:10]

        self.u = pd_in_t()
        for i in range(5):
            # TODO: move setting gains out of the loop?
            # maybe write a wrapper for pd_in_t ?
            self.u.leftLeg.motorPd.pGain[i]  = self.P[i]
            self.u.rightLeg.motorPd.pGain[i] = self.P[i]

            self.u.leftLeg.motorPd.dGain[i]  = self.D[i]
            self.u.rightLeg.motorPd.dGain[i] = self.D[i]

            self.u.leftLeg.motorPd.torque[i]  = 0 # Feedforward torque
            self.u.rightLeg.motorPd.torque[i] = 0

            self.u.leftLeg.motorPd.pTarget[i]  = target[i]
            self.u.rightLeg.motorPd.pTarget[i] = target[i + 5]

            self.u.leftLeg.motorPd.dTarget[i]  = 0
            self.u.rightLeg.motorPd.dTarget[i] = 0

        self.cassie_state = self.sim.step_pd(self.u)

    def step(self, action, return_omniscient_state=False):

        self.l_foot_orient = 0
        self.r_foot_orient = 0
        self.l_foot_cost = 0
        self.r_foot_cost = 0
        self.l_foot_cost_even = 0
        self.r_foot_cost_even = 0
        self.l_foot_cost_smooth = 0
        self.r_foot_cost_smooth = 0
        self.torque_cost = 0
        self.hipyaw_cost = 0
        foot_pos = np.zeros(6)
        self.smooth_cost = 0

        for _ in range(self.simrate):
            self.step_simulation(action)
            qvel = self.sim.qvel()

            # Foot Orientation Cost
            self.l_foot_orient += 20*(1 - np.inner(self.neutral_foot_orient, self.sim.xquat("left-foot")) ** 2)
            self.r_foot_orient += 20*(1 - np.inner(self.neutral_foot_orient, self.sim.xquat("right-foot")) ** 2)

            # Hip Yaw velocity cost
            self.hipyaw_cost += (np.abs(qvel[6]) + np.abs(qvel[19])) / 3

            # Foot height cost
            des_height = 0.15
            self.sim.foot_pos(foot_pos)
            foot_forces = self.sim.get_foot_forces()
            r_ground_cost = foot_pos[5]**2 +  np.linalg.norm(self.rfoot_vel)
            r_height_cost = 40*(des_height - foot_pos[5])**2
            r_td_cost = 40*(foot_pos[5])**2 * self.rfoot_vel[2]**2
            l_ground_cost = foot_pos[2]**2 +  np.linalg.norm(self.lfoot_vel)
            l_height_cost = 40*(des_height - foot_pos[2])**2
            l_td_cost = 40*(foot_pos[2])**2 * self.lfoot_vel[2]**2
            # Right foot cost
            if foot_forces[0] ==  0:
                # If left foot in air, then keep right foot in place and on ground
                self.r_foot_cost += r_ground_cost
            else:
                # Otherwise, left foot is on ground and can lift right foot up.
                if not self.r_high:
                    # If right foot not reach desired height yet, then use distance cost
                    self.r_foot_cost += r_height_cost
                else:
                    # Otherwise right foot reached height, bring back down slowly. Weight distance by z vel
                    self.r_foot_cost += r_td_cost
            # Left foot cost
            if foot_forces[1] ==  0:
                # If right foot in air, then keep left foot in place and on ground
                self.l_foot_cost += l_ground_cost
            else:
                # Otherwise, right foot is on ground and can lift left foot up.
                if not self.r_high:
                    # If left foot not reach desired height yet, then use distance cost
                    self.l_foot_cost += l_height_cost
                else:
                    # Otherwise left foot reached height, bring back down slowly. Weight distance by z vel
                    self.l_foot_cost += l_td_cost

            # Foot height cost even
            # For first half of cycle, lift up left foot and keep right foot on ground.
            if self.phase < self.phaselen / 2.0:
                self.r_foot_cost_even += r_ground_cost
                if not self.l_high:
                    # If left foot not reach desired height yet, then use distance cost
                    self.l_foot_cost_even += l_height_cost
                else:
                    # Otherwise left foot reached height, bring back down slowly. Weight distance by z vel
                    self.l_foot_cost_even += l_td_cost
            # For second half of cycle, lift up right foot and keep left foot on ground
            else:
                self.l_foot_cost_even += l_ground_cost
                if not self.r_high:
                    # If right foot not reach desired height yet, then use distance cost
                    self.r_foot_cost_even += r_height_cost
                else:
                    # Otherwise right foot reached height, bring back down slowly. Weight distance by z vel
                    self.r_foot_cost_even += r_td_cost

            # Foot height cost smooth
            # Left foot starts on ground and then lift up and then back on ground.
            # Right foot starts in air and then put on ground and then lift back up.
            one2one_clock = 0.5*(np.cos(2*np.pi/(self.phaselen+1)*self.phase) + 1)
            zero2zero_clock = 0.5*(np.cos(2*np.pi/(self.phaselen+1)*(self.phase-(self.phaselen+1)/2)) + 1)
            self.l_foot_cost_smooth += zero2zero_clock*l_height_cost + one2one_clock*l_ground_cost
            self.r_foot_cost_smooth += one2one_clock*r_height_cost + zero2zero_clock*r_ground_cost

            # Torque costs
            curr_torques = np.array(self.cassie_state.motor.torque[:])
            if self.prev_torque is not None:
                self.smooth_cost += 0.0001*np.linalg.norm(np.square(curr_torques - self.prev_torque))
            else:
                self.smooth_cost += 0
            self.prev_torque = curr_torques
            self.torque_cost += 0.00006*np.linalg.norm(np.square(curr_torques))

        self.l_foot_orient              /= self.simrate
        self.r_foot_orient              /= self.simrate
        self.l_foot_cost                /= self.simrate
        self.r_foot_cost                /= self.simrate
        self.l_foot_cost_even           /= self.simrate
        self.r_foot_cost_even           /= self.simrate
        self.l_foot_cost_smooth         /= self.simrate
        self.r_foot_cost_smooth         /= self.simrate
        self.smooth_cost                /= self.simrate
        self.torque_cost                /= self.simrate
        self.hipyaw_cost                /= self.simrate

        # print("hipyaw cost: ", self.hipyaw_cost)
        # print("torquecost: ", self.torque_cost)

        height = self.sim.qpos()[2]
        self.curr_action = action

        self.time  += 1
        self.phase += self.phase_add

        if (self.aslip_traj and self.phase >= self.phaselen) or self.phase > self.phaselen:
            self.phase = 0
            self.counter += 1

        # Early termination
        done = not(height > 0.4 and height < 3.0)

        reward = self.compute_reward(action)

        # reset avg foot quaternion
        self.avg_lfoot_quat = np.zeros(4)
        self.avg_rfoot_quat = np.zeros(4)

        # update previous action
        self.prev_action = action

        # TODO: make 0.3 a variable/more transparent
        if reward < 0.3:
            done = True

        if return_omniscient_state:
            return self.get_full_state(), self.get_omniscient_state(), reward, done, {}
        else:
            return self.get_full_state(), reward, done, {}

    def step_basic(self, action, return_omniscient_state=False):

        for _ in range(self.simrate):
            self.step_sim_basic(action)

        self.time  += 1
        self.phase += self.phase_add

        if (self.aslip_traj and self.phase >= self.phaselen) or self.phase > self.phaselen:
            self.phase = 0
            self.counter += 1

        if return_omniscient_state:
            return self.get_full_state(), self.get_omniscient_state()
        else:
            return self.get_full_state()

    def reset(self):

        self.phase = random.randint(0, self.phaselen)
        self.time = 0
        self.counter = 0

        self.state_history = [np.zeros(self._obs) for _ in range(self.history+1)]

        qpos, qvel = self.get_ref_state(self.phase)
        orientation = random.randint(-10, 10) * np.pi / 25
        quaternion = euler2quat(z=orientation, y=0, x=0)
        qpos[3:7] = quaternion
        self.y_offset = 0#random.uniform(-3.5, 3.5)
        qpos[1] = self.y_offset

        self.sim.set_qpos(qpos)
        self.sim.set_qvel(qvel)

        # Need to reset u? Or better way to reset cassie_state than taking step
        self.cassie_state = self.sim.step_pd(self.u)

        if self.aslip_traj:
            random_speed_idx = random.randint(0, self.num_speeds-1)
            self.speed = self.speeds[random_speed_idx]
            # print("current speed: {}".format(self.speed))
            self.trajectory = self.trajectories[random_speed_idx] # switch the current trajectory
            self.phaselen = self.trajectory.length - 1
        else:
            self.speed = (random.randint(0, 30)) / 10
        # Make sure that if speed is above 2, freq is at least 1.2
        # if self.speed > 1.3:# or np.any(self.speed_schedule > 1.6):
        #     self.phase_add = 1.3 + 0.7*random.random()
        # else:
        self.phase_add = 1 + random.random()

        self.orient_add = 0#random.randint(-10, 10) * np.pi / 25

        self.com_vel_offset = 0#0.1*np.random.uniform(-0.1, 0.1, 2)

        if self.dynamics_randomization:
            #### Dynamics Randomization ####
            self.damp_noise = np.clip([np.random.uniform(a, b) for a, b in self.damp_range], 0, None)
            self.mass_noise = np.clip([np.random.uniform(a, b) for a, b in self.mass_range], 0, None)
            # com_noise = [0, 0, 0] + [np.random.uniform(self.delta_x_min, self.delta_x_min)] + [np.random.uniform(self.delta_y_min, self.delta_y_max)] + [0] + list(self.default_ipos[6:])
            # fric_noise = [np.random.uniform(0.95, 1.05)] + [np.random.uniform(5e-4, 5e-3)] + [np.random.uniform(5e-5, 5e-4)]#+ list(self.default_fric[2:])
            # fric_noise = []
            # translational = np.random.uniform(0.6, 1.2)
            # torsional = np.random.uniform(1e-4, 1e-2)
            # rolling = np.random.uniform(5e-5, 5e-4)
            # for _ in range(int(len(self.default_fric)/3)):
            #     fric_noise += [translational, torsional, rolling]
            self.fric_noise = np.clip([np.random.uniform(0.6, 1.2), np.random.uniform(1e-4, 1e-2), np.random.uniform(5e-5, 5e-4)], 0, None)
            self.sim.set_dof_damping(self.damp_noise)
            self.sim.set_body_mass(self.mass_noise)
            # self.sim.set_body_ipos(com_noise)
            self.sim.set_geom_friction(self.fric_noise, "floor")
            self.sim.set_const()

        if self.slope_rand:
            rand_angle = np.pi/180*np.random.uniform(-5, 5, 2)
            floor_quat = euler2quat(z=0, y=rand_angle[0], x=rand_angle[1])
            self.sim.set_geom_quat(floor_quat, "floor")
        if self.joint_rand:
            self.joint_offsets = np.random.uniform(-0.03, 0.03, 14)

        # Reward terms
        self.l_foot_orient = 0
        self.r_foot_orient = 0
        self.l_foot_cost = 0
        self.r_foot_cost = 0
        self.l_foot_cost_even = 0
        self.r_foot_cost_even = 0
        self.l_foot_cost_smooth = 0
        self.r_foot_cost_smooth = 0
        self.smooth_cost = 0
        self.torque_cost = 0
        self.prev_torque = None

        return self.get_full_state()

    def reset_for_test(self, full_reset=False):
        self.phase = 0
        self.time = 0
        self.counter = 0
        self.orient_add = 0
        self.y_offset = 0
        self.phase_add = 1

        self.state_history = [np.zeros(self._obs) for _ in range(self.history+1)]

        if self.aslip_traj:
            self.speed = 0
            # print("current speed: {}".format(self.speed))
            self.trajectory = self.trajectories[0] # switch the current trajectory
            self.phaselen = self.trajectory.length - 1
        else:
            self.speed = 0

        if not full_reset:
            qpos, qvel = self.get_ref_state(self.phase)
            self.sim.set_qpos(qpos)
            self.sim.set_qvel(qvel)

            # Need to reset u? Or better way to reset cassie_state than taking step
            self.cassie_state = self.sim.step_pd(self.u)
        else:
            self.sim.full_reset()
            self.reset_cassie_state()

        if self.dynamics_randomization:
            self.damp_noise = self.default_damping
            self.mass_noise = self.default_mass
            self.fric_noise = self.default_fric
            self.sim.set_dof_damping(self.default_damping)
            self.sim.set_body_mass(self.default_mass)
            # self.sim.set_body_ipos(self.default_ipos)
            self.sim.set_geom_friction(self.default_fric)
            self.sim.set_const()

        if self.slope_rand:
            self.sim.set_geom_quat(np.array([1, 0, 0, 0]), "floor")

        # Reward terms
        self.l_foot_orient = 0
        self.r_foot_orient = 0
        self.l_foot_cost = 0
        self.r_foot_cost = 0
        self.l_foot_cost_even = 0
        self.r_foot_cost_even = 0
        self.l_foot_cost_smooth = 0
        self.r_foot_cost_smooth = 0
        self.smooth_cost = 0
        self.torque_cost = 0
        self.prev_torque = None

        return self.get_full_state()

    def reset_cassie_state(self):
        # Only reset parts of cassie_state that is used in get_full_state
        self.cassie_state.pelvis.position[:] = [0, 0, 1.01]
        self.cassie_state.pelvis.orientation[:] = [1, 0, 0, 0]
        self.cassie_state.pelvis.rotationalVelocity[:] = np.zeros(3)
        self.cassie_state.pelvis.translationalVelocity[:] = np.zeros(3)
        self.cassie_state.pelvis.translationalAcceleration[:] = np.zeros(3)
        self.cassie_state.terrain.height = 0
        self.cassie_state.motor.position[:] = [0.0045, 0, 0.4973, -1.1997, -1.5968, 0.0045, 0, 0.4973, -1.1997, -1.5968]
        self.cassie_state.motor.velocity[:] = np.zeros(10)
        self.cassie_state.joint.position[:] = [0, 1.4267, -1.5968, 0, 1.4267, -1.5968]
        self.cassie_state.joint.velocity[:] = np.zeros(6)

    # Helper function for updating the speed, used in visualization tests
    def update_speed(self, new_speed):
        if self.aslip_traj:
            self.speed = new_speed
            self.trajectory = self.trajectories[(np.abs(self.speeds - self.speed)).argmin()]
            old_phaselen = self.phaselen
            self.phaselen = self.trajectory.length - 1
            # update phase
            self.phase = int(self.phaselen * self.phase / old_phaselen)
            # new offset
            ref_pos, ref_vel = self.get_ref_state(self.phase)
            self.offset = ref_pos[self.pos_idx]
        else:
            self.speed = new_speed

    def compute_reward(self, action):
        qpos = np.copy(self.sim.qpos())
        qvel = np.copy(self.sim.qvel())

        ref_pos, ref_vel = self.get_ref_state(self.phase)

        if self.reward_func == "jonah_RNN":
            return jonah_RNN_reward(self)
        elif self.reward_func == "aslip":
            return aslip_reward(self, action)
        elif self.reward_func == "aslip_TaskSpace":
            return aslip_TaskSpace_reward(self, action)
        elif self.reward_func == "iros_paper":
            return iros_paper_reward(self)
        elif self.reward_func == "5k_speed_reward":
            return old_speed_reward(self)
        elif self.reward_func == "speed_footheightvelflag":
            return speedmatch_footheightvelflag_reward(self)
        elif self.reward_func == "speed_footheightvelflag_even":
            return speedmatch_footheightvelflag_even_reward(self)
        elif self.reward_func == "speed_footheightvelflag_even_capzvel":
            return speedmatch_footheightvelflag_even_capzvel_reward(self)
        elif self.reward_func == "step_even_reward":
            return step_even_reward(self)
        elif self.reward_func == "step_even_pelheight_reward":
            return step_even_pelheight_reward(self)
        elif self.reward_func == "step_smooth_pelheight_reward":
            return step_smooth_pelheight_reward(self)
        elif self.reward_func == "speed_footheightvelflag_even_footorient":
            return speedmatch_footheightvelflag_even_footorient_reward(self)
        elif self.reward_func == "speed_footheightvelflag_even_footorient_footdist":
            return speedmatch_footheightvelflag_even_footorient_footdist_reward(self)
        elif self.reward_func == "speed_footheightvelflag_even_footorient_smooth":
            return speedmatch_footheightvelflag_even_footorient_smooth_reward(self)
        elif self.reward_func == "speed_footheightvelflag_even_footorient_footdist_torquecost":
            return speedmatch_footheightvelflag_even_footorient_footdist_torquecost_reward(self)
        elif self.reward_func == "speed_footheightvelflag_even_footorient_footdist_torquecost_smooth":
            return speedmatch_footheightvelflag_even_footorient_footdist_torquecost_smooth_reward(self)
        elif self.reward_func == "speed_footheightsmooth_footorient":
            return speedmatch_footheightsmooth_footorient_reward(self)
        elif self.reward_func == "speedmatch_footheightsmooth_footorient_hipyaw_torquecost_reward":
            return speedmatch_footheightsmooth_footorient_hipyaw_torquecost_reward(self)
        else:
            raise NotImplementedError

  # get the corresponding state from the reference trajectory for the current phase
    def get_ref_state(self, phase=None):
        if phase is None:
            phase = self.phase

        if phase > self.phaselen:
            phase = 0

        desired_ind = phase * self.simrate if not self.aslip_traj else phase
        # phase_diff = desired_ind - math.floor(desired_ind)
        # if phase_diff != 0:       # desired ind is an int
        #     pos_prev = np.copy(self.trajectory.qpos[math.floor(desired_ind)])
        #     vel_prev = np.copy(self.trajectory.qvel[math.floor(desired_ind)])
        #     pos_next = np.copy(self.trajectory.qpos[math.ceil(desired_ind)])
        #     vel_next = np.copy(self.trajectory.qvel[math.ceil(desired_ind)])
        #     pos = pos_prev + phase_diff * (pos_next - pos_prev)
        #     vel = vel_prev + phase_diff * (vel_next - vel_prev)
        # else:
        # print("desired ind: ", desired_ind)
        pos = np.copy(self.trajectory.qpos[int(desired_ind)])
        vel = np.copy(self.trajectory.qvel[int(desired_ind)])

        # this is just setting the x to where it "should" be given the number
        # of cycles
        # pos[0] += (self.trajectory.qpos[-1, 0] - self.trajectory.qpos[0, 0]) * self.counter

        # ^ should only matter for COM error calculation,
        # gets dropped out of state variable for input reasons

        ###### Setting variable speed  #########
        pos[0] *= self.speed
        pos[0] += (self.trajectory.qpos[-1, 0] - self.trajectory.qpos[0, 0]) * self.counter * self.speed
        ######                          ########

        # setting lateral distance target to 0?
        # regardless of reference trajectory?
        pos[1] = 0

        if not self.aslip_traj:
            vel[0] *= self.speed

        return pos, vel

    def get_full_state(self):
        qpos = np.copy(self.sim.qpos())
        qvel = np.copy(self.sim.qvel())

        ref_pos, ref_vel = self.get_ref_state(self.phase + self.phase_add)

        # TODO: maybe convert to set subtraction for clarity
        # {i for i in range(35)} -
        # {0, 10, 11, 12, 13, 17, 18, 19, 24, 25, 26, 27, 31, 32, 33}

        # this is everything except pelvis x and qw, achilles rod quaternions,
        # and heel spring/foot crank/plantar rod angles
        # note: x is forward dist, y is lateral dist, z is height

        # makes sense to always exclude x because it is in global coordinates and
        # irrelevant to phase-based control. Z is inherently invariant to
        # trajectory despite being global coord. Y is only invariant to straight
        # line trajectories.

        # CLOCK BASED (NO TRAJECTORY)
        if self.clock_based:
            clock = [np.sin(2 * np.pi *  self.phase / (self.phaselen+1)),
                    np.cos(2 * np.pi *  self.phase / (self.phaselen+1))]

            ext_state = np.concatenate((clock, [self.speed]))

        # ASLIP TRAJECTORY
        elif self.aslip_traj and not self.clock_based:
            if(self.phase == 0):
                ext_state = np.concatenate(get_ref_aslip_ext_state(self, self.phaselen - 1))
            else:
                ext_state = np.concatenate(get_ref_aslip_ext_state(self, self.phase))

        # OTHER TRAJECTORY
        else:
            ext_state = np.concatenate([ref_pos[self.pos_index], ref_vel[self.vel_index]])

        # Update orientation
        new_orient = self.cassie_state.pelvis.orientation[:]
        new_translationalVelocity = self.cassie_state.pelvis.translationalVelocity[:]
        # new_translationalVelocity[0:2] += self.com_vel_offset
        quaternion = euler2quat(z=self.orient_add, y=0, x=0)
        iquaternion = inverse_quaternion(quaternion)
        new_orient = quaternion_product(iquaternion, self.cassie_state.pelvis.orientation[:])
        if new_orient[0] < 0:
            new_orient = -new_orient
        new_translationalVelocity = rotate_by_quaternion(self.cassie_state.pelvis.translationalVelocity[:], iquaternion)
        motor_pos = self.cassie_state.motor.position[:]
        if self.joint_rand:
            motor_pos += self.joint_offsets[0:10]

        # Use state estimator
        robot_state = np.concatenate([
            self.cassie_state.leftFoot.position[:],     # left foot position
            self.cassie_state.rightFoot.position[:],     # right foot position
            new_orient,                                 # pelvis orientation
            motor_pos,                                     # actuated joint positions

            new_translationalVelocity,                       # pelvis translational velocity
            self.cassie_state.pelvis.rotationalVelocity[:],                          # pelvis rotational velocity
            self.cassie_state.motor.velocity[:],                                     # actuated joint velocities

        ])

        #TODO: Set up foot position for non state est
        if self.state_est:
            state = np.concatenate([robot_state, ext_state])
        else:
            state = np.concatenate([qpos[self.pos_index], qvel[self.vel_index], ext_state])

        self.state_history.insert(0, state)
        self.state_history = self.state_history[:self.history+1]

        return np.concatenate(self.state_history)

    def render(self):
        if self.vis is None:
            self.vis = CassieVis(self.sim, "./cassie/cassiemujoco/cassie.xml")

        return self.vis.draw(self.sim)

# Currently unused
# def get_omniscient_state(self):
#     full_state = self.get_full_state()
#     omniscient_state = np.hstack((full_state, self.sim.get_dof_damping(), self.sim.get_body_mass(), self.sim.get_body_ipos(), self.sim.get_ground_friction))
#     return omniscient_state

#nbody layout:
# 0:  worldbody (zero)
# 1:  pelvis

# 2:  left hip roll
# 3:  left hip yaw
# 4:  left hip pitch
# 5:  left achilles rod
# 6:  left knee
# 7:  left knee spring
# 8:  left shin
# 9:  left tarsus
# 10:  left heel spring
# 12:  left foot crank
# 12: left plantar rod
# 13: left foot

# 14: right hip roll
# 15: right hip yaw
# 16: right hip pitch
# 17: right achilles rod
# 18: right knee
# 19: right knee spring
# 20: right shin
# 21: right tarsus
# 22: right heel spring
# 23: right foot crank
# 24: right plantar rod
# 25: right foot

# qpos layout
# [ 0] Pelvis x
# [ 1] Pelvis y
# [ 2] Pelvis z
# [ 3] Pelvis orientation qw
# [ 4] Pelvis orientation qx
# [ 5] Pelvis orientation qy
# [ 6] Pelvis orientation qz
# [ 7] Left hip roll         (Motor [0])
# [ 8] Left hip yaw          (Motor [1])
# [ 9] Left hip pitch        (Motor [2])
# [10] Left achilles rod qw
# [11] Left achilles rod qx
# [12] Left achilles rod qy
# [13] Left achilles rod qz
# [14] Left knee             (Motor [3])
# [15] Left shin                        (Joint [0])
# [16] Left tarsus                      (Joint [1])
# [17] Left heel spring
# [18] Left foot crank
# [19] Left plantar rod
# [20] Left foot             (Motor [4], Joint [2])
# [21] Right hip roll        (Motor [5])
# [22] Right hip yaw         (Motor [6])
# [23] Right hip pitch       (Motor [7])
# [24] Right achilles rod qw
# [25] Right achilles rod qx
# [26] Right achilles rod qy
# [27] Right achilles rod qz
# [28] Right knee            (Motor [8])
# [29] Right shin                       (Joint [3])
# [30] Right tarsus                     (Joint [4])
# [31] Right heel spring
# [32] Right foot crank
# [33] Right plantar rod
# [34] Right foot            (Motor [9], Joint [5])

# qvel layout
# [ 0] Pelvis x
# [ 1] Pelvis y
# [ 2] Pelvis z
# [ 3] Pelvis orientation wx
# [ 4] Pelvis orientation wy
# [ 5] Pelvis orientation wz
# [ 6] Left hip roll         (Motor [0])
# [ 7] Left hip yaw          (Motor [1])
# [ 8] Left hip pitch        (Motor [2])
# [ 9] Left knee             (Motor [3])
# [10] Left shin                        (Joint [0])
# [11] Left tarsus                      (Joint [1])
# [12] Left foot             (Motor [4], Joint [2])
# [13] Right hip roll        (Motor [5])
# [14] Right hip yaw         (Motor [6])
# [15] Right hip pitch       (Motor [7])
# [16] Right knee            (Motor [8])
# [17] Right shin                       (Joint [3])
# [18] Right tarsus                     (Joint [4])
# [19] Right foot            (Motor [9], Joint [5])


#!/usr/bin/python3
import pigpio
import time
import board
import busio
import adafruit_ads1x15.ads1115 as ADS
from adafruit_ads1x15.analog_in import AnalogIn
import numpy as np

class Heater(object):
    """
    Hardware PWM is used
    using GPIO18 to control output
    using ADC channel 1 as input
    using GPIO13 and ADC channel 3 as reference
    """

    # setup initial temp and frequency for pwm
    def __init__(self, initTemp=25.0, frequency=9000):
        super(Heater, self).__init__()
        self.initTemp = initTemp
        self.frequency = frequency
        self.MaxTemp = 33.0
        self.pi = pigpio.pi()

        # 9000hz give a relative nice curve
        self.pi.hardware_PWM(13, self.frequency, 0)
        self.pi.hardware_PWM(18, self.frequency, 0)

        # Create the I2C bus
        i2c = busio.I2C(board.SCL, board.SDA)
        # Create the ADC object using the I2C bus
        ads = ADS.ADS1115(i2c)
        # Create single-ended input on channel 1
        self.chanRef = AnalogIn(ads, ADS.P1)
        self.chanCali = AnalogIn(ads, ADS.P3)

        self.readObejectiveTemp()
        self.setTemp(initTemp)

    # define exit behavior
    def __del__(self):
        # seems something here is wrong
        # self.pi.hardware_PWM(13, 0, 0)
        # self.pi.hardware_PWM(18, 0, 0)
        self.pi.stop()
        print("Heater exit success")
        print("!!! Unplug the socket !!!")

    def readObejectiveTemp(self):
        self.obejectiveTemp = self.chanRef.voltage*10

    def readCalibrationTemp(self):
        self.caliTemp = self.chanCali.voltage*10

    # set the temp on Heater within tolerance
    # I think 0.1 is a good choice
    def setTemp(self, targetTemp, tolerance=0.1):
        self.temp = targetTemp
        self.pi.hardware_PWM(18, self.frequency, int(
            targetTemp/self.MaxTemp*1000000))
        self.pi.hardware_PWM(13, self.frequency, int(
            targetTemp/self.MaxTemp*1000000))
        # remaining part is not implement

    def outputBanner(self):
        return("{} {} {}\n".format('Setpoint', 'Output', 'Objective'))

    def outputTemp(self):
        self.readObejectiveTemp()
        self.readCalibrationTemp()
        return("{:>5.2f}  {:>5.2f}  {:>5.2f} \n".format(
            self.temp, self.caliTemp, self.obejectiveTemp))

class TempCycle(Heater):
    """
    initTemp: start temp for heater
    finalTemp: end temp for heater
    cycleTime: cycle how many times. 1 cycle ~ 30min
    increase_rate: max rate measured is 0.3C/min,set to 0.2C/min by default
    decrease_rate: max rate measured is 1/6 C/min,set to 0.13 by default
    tempStep: step between temps
    outputStep: time between two lines, a loop is around 60s
    """

    def __init__(self, initTemp, finalTemp,
                 cycleNumber=20,
                 increase_rate=0.2, decrease_rate=0.10, tempStep=0.2,
                 outputStep=10):
        super(TempCycle, self).__init__(initTemp, frequency=9000)
        self.initTemp = initTemp
        self.finalTemp = finalTemp
        self.increase_rate = increase_rate
        self.decrease_rate = decrease_rate
        self.tempStep = tempStep
        self.cycleNumber = cycleNumber
        self.outputStep = outputStep
        if self.initTemp > self.finalTemp:
            self.increase = True
        else:
            self.increase = False

    # if increase, rate = increase_rate
    # else rate = decrease_rate
    def _rate_(self):
        return self.decrease_rate + int(self.increase) *\
            (self.increase_rate-self.decrease_rate)

    def start(self):
        filename = time.strftime(
            "%m-%d-at-%H:%M-year%YCyclelog.txt", time.localtime())
        with open(filename, 'w') as f:
            self.setTemp(self.initTemp)

            f.write(self.outputBanner())
            f.write("timestep={}s\n".format(self.outputStep))
            while self.cycleNumber:

                # from start to end
                # finaltemp is include in the next loop
                for setpoint in np.arange(self.initTemp, self.finalTemp, self.tempStep):
                    startTime = time.time()
                    self.setTemp(setpoint)
                    rate = self._rate_()
                    while time.time()-startTime < self.tempStep/rate*60:
                        f.write(self.outputTemp())
                        time.sleep(self.outputStep)
                    self.increase = False

                # from end to start
                # inittemp is include in the next loop
                for setpoint in np.arange(self.finalTemp, self.initTemp, -self.tempStep):
                    startTime = time.time()
                    self.setTemp(setpoint)
                    rate = self._rate_()
                    while time.time()-startTime < self.tempStep/rate*60:
                        f.write(self.outputTemp())
                        time.sleep(self.outputStep)
                    self.increase = False

                self.cycleNumber -= 1

class TempSchedule(Heater):
    """
    schedule: temp list for heater
    reside_time: time list for each temp.choose the longer one
                compared to calculated one.
    """

    def __init__(self, schedule, reside_time,
                 cycleNumber=20,
                 increase_rate=0.2, decrease_rate=0.10, tempStep=0.2,
                 outputStep=10):
        super(TempSchedule, self).__init__(initTemp=25, frequency=9000)
        self.schedule = schedule
        self.reside_time = reside_time
        self.increase_rate = increase_rate
        self.decrease_rate = decrease_rate
        self.tempStep = tempStep
        self.cycleNumber = cycleNumber
        self.outputStep = outputStep

    def prepare(self):
        self.tempLength = len(self.schedule)
        self.increaseList = [self.schedule[(i+1) % self.tempLength] - self.schedule[i]
                             for i in range(self.tempLength)]
        self.increaseRateList = [self.increase_rate if x > 0 else self.decrease_rate
                                 for x in self.increaseList]
        self.calculatedTime = [self.increaseList[i]/self.increaseRateList[i]
                               for i in range(self.tempLength)]
        for i in range(self.tempLength):
            if self.reside_time[i] < abs(self.calculatedTime[i]):
                self.reside_time[i] = abs(self.calculatedTime[i])

    def start(self):
        self.prepare()
        filename = time.strftime(
            "%m-%d-at-%H:%M-year%YCyclelog.txt", time.localtime())
        with open(filename, 'w') as f:
            f.write(self.outputBanner())
            f.write("timestep={}s\n".format(self.outputStep))
            while self.cycleNumber:
                for i in range(self.tempLength):
                    setpoint = self.schedule[i]
                    rate = self.increaseRateList[i]
                    staytime = self.reside_time[i]
                    startTime = time.time()
                    self.setTemp(setpoint)
                    while time.time()-startTime < staytime*60:
                        f.write(self.outputTemp())
                        time.sleep(self.outputStep)
                self.cycleNumber -= 1

if __name__ == '__main__':
    test = Heater(25.0)
    del test
    test = TempCycle(25, 26)
    del test


import warnings

import numpy as np
import mxnet as mx
import mxnet.ndarray as nd
import mxnet.gluon as gluon
from mxnet import autograd
from Attention1 import LSTM
from dataset1 import dataset1
from utils.register_experiment import Register
from mxnet.gluon.data import DataLoader
from multiprocessing import cpu_count

from tqdm import tqdm

import json
import argparse
import numpy as np
from sklearn.metrics import recall_score, confusion_matrix

def recall(y, y_hat):
    y = y.asnumpy()
    y_hat = y_hat.asnumpy()
    return recall_score(y, y_hat), confusion_matrix(y, y_hat).ravel()

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Arguments for Attentionmodel')
    parser.add_argument('--dataset', type=str, help = "Dataset to use for the training")
    parser.add_argument('--filters', type=str, help = "Filters to apply for selecting element in database")
    parser.add_argument('--saveDir', type=str, help = "Directory for checkpoint saving")
    parser.add_argument('--dct', type=str, help = "File for the dictionary")
    parser.add_argument('--SEQ_LENGTH', type=int, help="Fixed size length to expand or srink text")
    parser.add_argument('--EMBEDDING_DIM', type=int, help="Size of the embedding dimention")
    parser.add_argument('--HIDDEN', type=int, help="Size of the hidden layer")
    parser.add_argument('--LAYERS', type=int, help="Number of hidden layers")
    parser.add_argument('--DROPOUT', type=float, help="Dropout value")
    parser.add_argument('--BATCH_SIZE', type=int, help="Batch size")
    parser.add_argument('--EPOCHS', type=int, help="Number of epochs for training")
    parser.add_argument('--utils', type=str, help = "Helper directory")
    parser.add_argument('--db', type=str, help = "DB name", required=True)
    parser.add_argument('--collection', type=str, help = "DB collection")
    parser.add_argument('--host', type=str, help = "DB host")
    parser.add_argument('--port', type=int, help = "Port number of db")

    args = parser.parse_args()

    warnings.filterwarnings("ignore")

    SEQ_LENGTH = args.SEQ_LENGTH
    EMBEDDING_DIM = args.EMBEDDING_DIM
    HIDDEN = args.HIDDEN
    LAYERS = args.LAYERS
    DROPOUT = args.DROPOUT
    BATCH_SIZE = args.BATCH_SIZE
    EPOCHS = args.EPOCHS
    CPU_COUNT = cpu_count()
    ctx = mx.gpu()

    ds = dataset1(args.dataset, json.loads(args.filters), SEQ_LENGTH, args.dct)

    r = Register(args.host, args.port, args.db, args.collection)
    r.newExperiment(r.getLastExperiment() + 1, 'Attention')

    net = LSTM(ds.dctLen(), SEQ_LENGTH, EMBEDDING_DIM, HIDDEN, LAYERS, DROPOUT)
    net.initialize(mx.init.Normal(sigma=0.01), ctx=ctx)
    trainer = gluon.Trainer(net.collect_params(), 'adam', {'learning_rate': 0.00001, 'wd': 0.0001})
    loss = gluon.loss.SigmoidBinaryCrossEntropyLoss(from_sigmoid=False)
    hidden = net.begin_state(func=mx.nd.zeros, batch_size=BATCH_SIZE, ctx=mx.cpu(0))

    for epochs in range(0, EPOCHS):
        pbar = tqdm(total=len(ds) // BATCH_SIZE)
        total_L = 0.0
        accuracy = []
        hidden = net.begin_state(func=mx.nd.zeros, batch_size=BATCH_SIZE, ctx=ctx)
        dl = DataLoader(ds, batch_size = BATCH_SIZE, last_batch = 'rollover', num_workers=CPU_COUNT // 2)
        recall_list = []
        cfMatrix = []
        acc = mx.metric.Accuracy()
        for batch, labels in dl:
            pbar.update(1)
            with autograd.record():
                output, _ = net(batch.copyto(ctx), hidden)
                L = loss(output, labels.copyto(ctx))
                L.backward()
                pred = output.argmax(axis=1)
                y = labels.argmax(axis=1)
                acc.update(y, pred)
                rec, mat = recall(y, pred)
                recall_list.append(rec)
                cfMatrix.append(mat)
            trainer.step(BATCH_SIZE)
            total_L += mx.nd.sum(L).asscalar()
        pbar.close()
        print("epoch : {}, Loss : {}, Accuracy : {}, recall : {}".format(epochs, total_L, acc.get()[1],
                                                                         np.mean(recall_list)))
        r.addResult({'epoch': epochs,
                     'train': {'accuracy': acc.get()[1], 'loss': total_L, 'recall': np.mean(recall_list),
                               'Confusion Matrix': list(map(int, sum(cfMatrix)))}}, r.getLastExperiment() + 1)
        net.save_parameters(args.saveDir + "_{:04d}.params".format(epochs))
    r.addParams({'SEQ_LENGTH': SEQ_LENGTH,
                 'EMBEDDING_DIM': EMBEDDING_DIM,
                 'HIDDEN': HIDDEN,
                 'LAYERS': LAYERS,
                 'DROPOUT': DROPOUT,
                 'BATCH_SIZE': BATCH_SIZE,
                 'EPOCHS': EPOCHS,
                 'dataset' : args.dataset}, r.getLastExperiment() + 1)


# Copyright 2021 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
'''
This file gets the jointpoints from image.
'''
from __future__ import division

import os
import json
import numpy as np
import cv2 as cv

from mindspore import Tensor, float32, context
from mindspore import load_checkpoint, load_param_into_net

from yolo.yolo import YOLOV3DarkNet53
from yolo.utils import load_image, detect, statistic_normalize_img
from src.config import config
from src.dataset import flip_pairs
from src.utils.transforms import flip_back, get_affine_transform, bbox2sc
from src.utils.inference import get_final_preds
from src.FastPose import createModel
from src.utils.nms import pose_nms
from src.utils.fn import vis_frame

def detect_bbox():
    '''
    detect_bbox
    '''
    img, ori_img_shape = load_image(config.detect_image, config.yolo_image_size)
    network = YOLOV3DarkNet53(is_training=False)
    param_dict = load_checkpoint(config.yolo_ckpt)
    print('loading model fastpose_ckpt from {}'.format(config.yolo_ckpt))
    load_param_into_net(network, param_dict)
    prediction = network(Tensor(img))
    output_big, output_me, output_small = prediction
    output_big = output_big.asnumpy()
    output_me = output_me.asnumpy()
    output_small = output_small.asnumpy()
    bboxes = detect([output_big, output_me, output_small], ori_img_shape, config.yolo_threshold)
    if config.save_bbox_image:
        image = cv.imread(config.detect_image)
        image = image.astype(np.float32)
        for i in range(bboxes.shape[0]):
            w, h, wi, hi = int(bboxes[i, 0]), int(bboxes[i, 1]),\
                           int(bboxes[i, 2]), int(bboxes[i, 3])
            cv.rectangle(image, (w, h), (w + wi, h + hi), (255, 255, 255), thickness=2)
        name = config.detect_image.split('/')[-1].split('.')[0]
        cv.imwrite(os.path.join(config.result_path, name + "_withBbox.jpg"), image)
    return bboxes

def inference(bboxes):
    '''
    inference
    '''
    image_width = config.MODEL.IMAGE_SIZE[0]
    image_height = config.MODEL.IMAGE_SIZE[1]
    aspect_ratio = image_width * 1.0 / image_height
    scales, centers = bbox2sc(bboxes, aspect_ratio)
    model = createModel()
    ckpt_name = config.fast_pose_ckpt
    print('loading model fastpose_ckpt from {}'.format(ckpt_name))
    load_param_into_net(model, load_checkpoint(ckpt_name))
    image_size = np.array(config.MODEL.IMAGE_SIZE, dtype=np.int32)
    data_numpy = cv.imread(config.detect_image, cv.IMREAD_COLOR | cv.IMREAD_IGNORE_ORIENTATION)

    inputs = []
    bbox_num = bboxes.shape[0]
    image_size = np.array(config.MODEL.IMAGE_SIZE, dtype=np.int32)
    for i in range(bbox_num):
        s, c = scales[i], centers[i]
        trans = get_affine_transform(c, s, 0, image_size, inv=0)
        image = cv.warpAffine(data_numpy, trans, (int(image_size[0]), int(image_size[1])), flags=cv.INTER_LINEAR)
        image_data = np.transpose(statistic_normalize_img(image, statistic_norm=True), (2, 0, 1))
        inputs.append(image_data)
    inputs = np.array(inputs, dtype=np.float32)
    output = model(Tensor(inputs, float32)).asnumpy()
    if config.TEST.FLIP_TEST:
        inputs_flipped = Tensor(inputs[:, :, :, ::-1], float32)
        output_flipped = model(inputs_flipped)
        output_flipped = flip_back(output_flipped.asnumpy(), flip_pairs)

        if config.TEST.SHIFT_HEATMAP:
            output_flipped[:, :, :, 1:] = \
                output_flipped.copy()[:, :, :, 0:-1]

            output = (output + output_flipped) * 0.5
    preds, maxvals = get_final_preds(output, centers, scales)

    return preds, maxvals

def DataWrite(result):
    '''
    DataWrite
    '''
    img_name = result['imgname'].split('/')[-1].split('.')[0]
    save_path = os.path.join(config.result_path, img_name)

    img = cv.imread(config.detect_image)
    img = vis_frame(img, result)
    cv.imwrite(save_path + ".jpg", img)

    for i in range(len(result['result'])):
        result['result'][i]['keypoints'] = result['result'][i]['keypoints'].tolist()
        result['result'][i]['kp_score'] = result['result'][i]['kp_score'].tolist()
        result['result'][i]['proposal_score'] = result['result'][i]['proposal_score'].tolist()

    with open(save_path + ".json", "w") as f:
        json.dump(result, f)
    print('result was written successfully in {}'.format(config.result_path))

def main():
    context.set_context(mode=context.GRAPH_MODE,
                        device_target="Ascend", save_graphs=False, device_id=5)
    bboxes = detect_bbox()
    pose_preds, pose_scores = inference(bboxes)

    result = pose_nms(bboxes[:, 0:4], bboxes[:, 4], pose_preds, pose_scores)
    result = {'imgname': config.detect_image, 'result': result}
    DataWrite(result)

if __name__ == "__main__":
    main()


import argparse
import time
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim

from models import modules, net, resnet, densenet
import net_mask
import loaddata
import util
import numpy as np
import sobel
import net_mask
import os

import pdb

parser = argparse.ArgumentParser(description='single depth estimation')
parser.add_argument('--epochs', default=10, type=int,
                    help='number of total epochs to run')
parser.add_argument('--start-epoch', default=0, type=int,
                    help='manual epoch number (useful on restarts)')
parser.add_argument('--lr', '--learning-rate', default=0.0001, type=float,
                    help='initial learning rate')
parser.add_argument('--weight-decay', '--wd', default=1e-4, type=float,
                    help='weight decay (default: 1e-4)')

def define_model(encoder='resnet'):
    if encoder is 'resnet':
        original_model = resnet.resnet50(pretrained = True)
        Encoder = modules.E_resnet(original_model)
        model = net.model(Encoder, num_features=2048, block_channel = [256, 512, 1024, 2048])
    if encoder is 'densenet':
        original_model = densenet.densenet161(pretrained=True)
        Encoder = modules.E_densenet(original_model)
        model = net.model(Encoder, num_features=2208, block_channel = [192, 384, 1056, 2208])
    if encoder is 'senet':
        original_model = senet.senet154(pretrained='imagenet')
        Encoder = modules.E_senet(original_model)
        model = net.model(Encoder, num_features=2048, block_channel = [256, 512, 1024, 2048])

    return model

def main():
    global args
    args = parser.parse_args()

    model_selection = 'resnet'
    model = define_model(encoder = model_selection)

    original_model2 = net_mask.drn_d_22(pretrained=True)
    model2 = net_mask.AutoED(original_model2)

    if torch.cuda.device_count() == 8:
        model = torch.nn.DataParallel(model, device_ids=[0, 1, 2, 3, 4, 5, 6, 7]).cuda()
        model2 = torch.nn.DataParallel(model2, device_ids=[0, 1, 2, 3, 4, 5, 6, 7]).cuda()
        batch_size = 64
    elif torch.cuda.device_count() == 4:
        model = torch.nn.DataParallel(model, device_ids=[0, 1, 2, 3]).cuda()
        model2 = torch.nn.DataParallel(model2, device_ids=[0, 1, 2, 3]).cuda()
        batch_size = 32
    else:
        model = torch.nn.DataParallel(model).cuda()
        model2 = torch.nn.DataParallel(model2).cuda()
        batch_size = 8
    model.load_state_dict(torch.load('./pretrained_model/model_' + model_selection))

    cudnn.benchmark = True
    optimizer = torch.optim.Adam(model2.parameters(), args.lr, weight_decay=args.weight_decay)

    train_loader = loaddata.getTrainingData(batch_size)
    for epoch in range(args.start_epoch, args.epochs):
        train(train_loader, model, model2, optimizer, epoch)

    torch.save(model2.state_dict(), './net_mask/mask_'+model_selection)

def train(train_loader, model, model2, optimizer, epoch):
    batch_time = AverageMeter()
    losses = AverageMeter()
    totalNumber = 0
    errorSum = {'MSE': 0, 'RMSE': 0, 'ABS_REL': 0, 'LG10': 0,
                'MAE': 0,  'DELTA1': 0, 'DELTA2': 0, 'DELTA3': 0}
    model.eval()
    model2.train()

    cos = nn.CosineSimilarity(dim=1, eps=0)
    get_gradient = sobel.Sobel().cuda()

    end = time.time()
    for i, sample_batched in enumerate(train_loader):
        image, depth_ = sample_batched['image'], sample_batched['depth']

        image = torch.autograd.Variable(image).cuda()
        depth_ = torch.autograd.Variable(depth_).cuda()

        ones = torch.ones(depth_.size(0), 1, depth_.size(2),depth_.size(3)).float().cuda()
        ones = torch.autograd.Variable(ones)

        depth = model(image.clone()).detach()

        optimizer.zero_grad()
        mask = model2(image)
        output = model(image*mask)

        depth_grad = get_gradient(depth)
        output_grad = get_gradient(output)
        depth_grad_dx = depth_grad[:, 0, :, :].contiguous().view_as(depth)
        depth_grad_dy = depth_grad[:, 1, :, :].contiguous().view_as(depth)
        output_grad_dx = output_grad[:, 0, :, :].contiguous().view_as(depth)
        output_grad_dy = output_grad[:, 1, :, :].contiguous().view_as(depth)

        depth_normal = torch.cat((-depth_grad_dx, -depth_grad_dy, ones), 1)
        output_normal = torch.cat((-output_grad_dx, -output_grad_dy, ones), 1)

        loss_depth = torch.log(torch.abs(output - depth) + 0.5).mean()
        loss_dx = torch.log(torch.abs(output_grad_dx - depth_grad_dx) + 0.5).mean()
        loss_dy = torch.log(torch.abs(output_grad_dy - depth_grad_dy) + 0.5).mean()
        loss_normal = torch.abs(1 - cos(output_normal, depth_normal)).mean()

        loss_rec = loss_depth + loss_normal + (loss_dx + loss_dy)
        loss_sparse = mask.mean()

        loss = loss_rec + loss_sparse*5

        losses.update(loss_sparse.data[0], image.size(0))
        loss.backward()
        optimizer.step()

        batch_time.update(time.time() - end)
        end = time.time()

        batchSize = depth.size(0)

        errors = util.evaluateError(output,depth)
        errorSum = util.addErrors(errorSum, errors, batchSize)
        totalNumber = totalNumber + batchSize
        averageError = util.averageErrors(errorSum, totalNumber)

        print('Epoch: [{0}][{1}/{2}]\t'
          'Time {batch_time.val:.3f} ({batch_time.sum:.3f})\t'
          'Loss {loss.val:.4f} ({loss.avg:.4f})'
          .format(epoch, i, len(train_loader), batch_time=batch_time, loss=losses))
    print('errors: ', averageError)

class AverageMeter(object):
    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

if __name__ == '__main__':
    main()


#!/usr/bin/env python
from scipy import *
import struct, sys, re
from scipy import linalg
import optparse
import rdU
from utils import W2kEnvironment, Ry_in_eV
import indmffile
import copy

class Struct:
    def __init__(self, case):
        self.case = case
        self.extn = 'struct'
        self.parse()

    def parse(self):
        '''The file is structured for reading by fortran code,
        so data is positioned by line and by column.'''
        f = open(self.case + '.' + self.extn, 'r')
        self.title = f.next().strip()

        line = f.next()
        self.lattice = line[0:4].strip()
        self.nat = int(line[27:30])
        self.latticename = line[30:].strip()

        self.mode = f.next()[13:17]

        line = f.next()
        self.a, self.b, self.c, self.alpha, self.beta, self.gamma = [float(line[i:i+10]) for i in range(0,60,10)]

        self.iatom  = []
        self.mult   = []
        self.isplit = []
        self.aname  = []
        self.npt    = []
        self.r0     = []
        self.rmt    = []
        self.Znuc   = []
        self.pos    = []
        self.rotloc = []

        for iat in range(self.nat):
            line = f.next()
            self.iatom.append( int(line[4:8]) )
            pos = [[float(line[col:col+10]) for col in 12+13*arange(3)]]

            line = f.next()
            mult = int(line[15:17])
            self.mult.append(mult)

            self.isplit.append( int(line[34:37]) )

            for mu in range(mult-1):
                line = f.next()
                pos.append( [float(line[col:col+10]) for col in 12+13*arange(3)] )

            self.pos.append(pos)

            line = f.next()
            self.aname.append( line[0:10].strip() )
            self.npt.append( int(line[15:20]) )
            self.r0.append( float(line[25:35]) )
            self.rmt.append( float(line[40:50]) )
            self.Znuc.append( int(float(line[55:65])) )

            rt=[]
            for i in range(3):
                line = f.next()
                rt.append( [float(line[col:col+10]) for col in 20+10*arange(3)] )
            self.rotloc.append(rt)

        self.Ng = int(f.next().split()[0])
        self.iz=[]
        self.tau=[]
        for i in range(self.Ng):
            tiz=zeros((3,3), dtype=float)
            ttau=zeros(3, dtype=float)
            for j in range(3):
                line = f.next()
                tiz[j,:] = array([int(line[0:2]), int(line[2:4]), int(line[4:6])])
                ttau[j] = float(line[6:16])
            self.iz.append(tiz)
            self.tau.append(ttau)
            f.next()
        #print self.iz

        f.close()

    def debugprint(self, f=sys.stdout):
        print >> f, '***** Structure File Contents *****'
        print >> f, 'title =', self.title
        print >> f, 'Number of sorts, nat =', self.nat
        print >> f, 'unit cell parameters (a,b,c) =', self.a, self.b, self.c
        print >> f, 'angles (alpha, beta, gamma) =', self.alpha, self.beta, self.gamma

        printlist = [
            (self.aname,  'Atom name, aname ='),
            (self.iatom,  'Atom index, iatom ='),
            (self.mult,   'Number of atoms of this type, mult ='),
            (self.isplit, 'Symmetry of the atom, isplit ='),
            (self.npt,    'Number of radial points, npt ='),
            (self.r0,     'First point in radial mesh, r0 ='),
            (self.rmt,    'Muffin tin radius, rmt ='),
            (self.Znuc,   'Atomic number, Znuc ='),
            ]

        for i in range(self.nat):
            print >> f, '---------- atom type', i, '------------'

            for var,docstr in printlist:
                print >> f, docstr, var[i]

            print >> f, 'Position(s) inside the unit cell:'
            for m in range(self.mult[i]):
                print >> f, '    pos =', self.pos[i][m]

            print >> f, 'Local rotation matrix:'
            for row in self.rotloc[i]:
                print >> f, row

        print >> f, '***********************************'

    def flat(self, notflat):
        '''Return a flat view of given data as a list.
        Example: if w.mult = [2,4] and w.aname = ['V', 'O']
        w.flatten(w.aname) -> ['V', 'V', 'O', 'O', 'O', 'O']'''
        if notflat is self.pos:
            listoflists = self.pos
        else:
            listoflists = [[elem]*mult for elem,mult in zip(notflat, self.mult)]
        return reduce(operator.add, listoflists)

def DirectVectors(BR,RN):
    #v0 = cross(BR[1],BR[2])
    #v1 = cross(BR[2],BR[0])
    #v2 = cross(BR[0],BR[1])
    #vol = dot(v0,BR[0])
    #array([v0/vol,v1/vol,v2/vol])
    return dot(RN, linalg.inv(BR).transpose() )

def Print(fh, R, tHr):
    R0 = R # dot(DVinv,R)
    Rs=[]
    for i in range(3):
        if abs(round(R0[i])-R0[i])<1e-3:
            Rs.append( "%d" % int(round(R0[i])) )
        else:
            Rs.append("%.1f" % R0[i])

    print >> fh, 'Hopping[(%s,%s,%s)]=[' % tuple(Rs)
    for l1 in range(shape(tHr)[0]):
        print >> fh, '[',
        for l2 in range(shape(tHr)[1]):
            print >> fh, "%8.4f+%8.4fj," % (tHr[l1,l2].real, tHr[l1,l2].imag),
        print >> fh, '],'
    print >> fh, ']'

def ReadBRs(case):
    "reads BR1 and BR2 from case.rotlm"
    fr = open(case+'.rotlm')
    BR = zeros((2,3,3), dtype=float)
    for ii in range(2):
        fr.next()
        for i in range(3):
            BR[ii,:,i] = map(float,fr.next().split())

    BRX = array(linalg.inv(BR[1]) * matrix(BR[0]))
    return (BR,BRX)

def FindMaxMinBand(Ek, Emin, Emax):
    "Finds minimum and maximum band to be used in downfolding"

    istart=None
    iend=None
    for i in range(len(Ek)):
        if Ek[i]>Emin:
            istart = i
            break

    for i in range(istart,len(Ek)):
        if Ek[i]>Emax:
            iend = i
            break
    return array([istart, iend])

def TryAddBands(Nse_global, Ek, Emin, Emax, SmearN, SmearX):

    Nse = Nse_global[:]

    Em = Ek[Nse[0]-1] # closest energy below the cutoff
    Ep = Ek[Nse[1]]   # closest energy above the cutoff

    print Em, Ep

    EmCanAdd, EpCanAdd = False, False

    if Emin-Em < SmearN : EmCanAdd = True
    if Ep-Emax < SmearX : EpCanAdd = True

    Succ = True

    if EmCanAdd and EpCanAdd:
        if Emin-Em < Ep-Emax:
            Nse[0]-=1
            print 'Both could be added. Adding below.'
        else:
            Nse[1]+=1
            print 'Both could be added. Adding above.'
    elif EmCanAdd :
        Nse[0]-=1
        print 'Can add below'
    elif EpCanAdd :
        Nse[1]+=1
        print 'Can add above'
    else:
        print 'Can not find enough bands in your interval. Increase Emin, Emax or smearings'
        Succ = False

    print Nse

    return (Nse, Succ)

# RS = (i0,i1,i2)*DV[:,:]
# RS*DV^-1 = (i0,i1,i2)

def GiveRs(div, DV):
    "Gives real space vectors"
    Rs=[]
    for i0 in range(div[0]/2,-div[0]/2,-1):
        for i1 in range(div[1]/2,-div[1]/2,-1):
            for i2 in range(div[2]/2,-div[2]/2,-1):
                Rs.append(DV[0]*i0 + DV[1]*i1 + DV[2]*i2)
    return Rs

if __name__ == '__main__':
    """ Takes the DMFT projector Udmft.0 and LDA eigenvalues (case.energy) to produces downfolding hopping parameters.
    It also requires two parameters specifying the bands used in downfolding. The lower cutoff is determined by energy
    '--Emin' (counted from EF in eV), while the upped cutoff is determined by the number of bands '--nbnd'.
    """

    usage = """usage: %prog [ options ]

    The script takes the DMFT projector Udmft.0 and LDA eigenvalues (case.energy)
    and produces downfolded hopping parameters. It requires parameters
    specifying bands used in downfolding. The lower/upper cutoff is determined
    by energy '--Emin' or '-N' /'--Emax' or '-X' (counted from EF in eV).
    Additional bands are added if needed in the interval
    '--SmearN' or '-n' / '--SmerX' or '-x'

    --------------------   Emax+SmearX
      added if needed
    ---------------------  Emax

      bands always used in downfolding

    ---------------------  Emin
      added if needed
    ---------------------  Emin-SmearN
    """

    # Finds what is case
    w2k = W2kEnvironment()

    parser = optparse.OptionParser(usage)
    parser.add_option("-N", "--Emin",  dest="Emin", type="float", default=None, help="Bands above Emin (counted from EF in eV) will be consider in downfolding.")
    parser.add_option("-X", "--Emax",  dest="Emax", type="float", default=None, help="Bands below Emax (counted from EF in eV) will be consider in downfolding.")
    parser.add_option("-n", "--SmearN",  dest="SmearN", type="float", default=1., help="When determining Emin cutoff, we use smearing in eV if needed")
    parser.add_option("-x", "--SmearX",  dest="SmearX", type="float", default=1., help="When determining Emax cutoff, we use smearing in eV if needed")
    parser.add_option("-s", "--smalls",  dest="smalls", type="float", default=0.1, help="The smallest acceptable singular value")
    # Next, parse the arguments
    (options, args) = parser.parse_args()

    if options.Emin is None:
        print 'You have to specify minimum band energy! Option -E or --Emin'
        sys.exit(1)
    if options.Emax is None:
        print 'You have to specify maximum band energy! Option -M or --Emax'
        sys.exit(1)

    print 'case=%s Emin=%s, Emax=%s SmearN=%s SmearX=%s' %  (w2k.case, options.Emin, options.Emax, options.SmearN, options.SmearX)

    # File handlers to exchange files with Fortran
    # These numbers are not relevant, but need to be set
    fhp = 233
    fhe = 60
    Ntry=10

    # Reads indmf-file
    inl = indmffile.Indmfl(w2k.case)
    inl.read()
    EF = inl.EF
    print 'Fermi energy is ', EF
    # atoms in the projection for downfolding

    atms = inl.atoms.keys()
    print 'atoms in the projection for downfolding=', atms
    # Finds actual positions of these atoms
    pos = [] # zeros((max(atms),3),dtype=float)
    fi = open(w2k.case+'.outputdmfu','r')
    for line in fi:
        m = re.search('Actual position of atom *(\d+)',line)
        if m is not None:
            if int(m.group(1)) in atms:
                atom = int(m.group(1))

                print 'ATOM', atom, ' in atms=', atms

                m = re.search('Actual position of atom *(\d+) *is: (.*)',line)
                #pos[atom-1,:] = map(float,m.group(2).split())
                pos.append( map(float,m.group(2).split()) )
    pos=array(pos)
    print 'Actual position of atoms for downfolding=',  pos.tolist()

    strc = Struct(w2k.case)
    strc.parse()

    # reads BR1 and BR2 from case.rotlm
    (BR,BRX) = ReadBRs(w2k.case)
    # Computes direct vectors, knowing the reciprocal basis
    DV = DirectVectors(BR[1],BR[0])

    print 'Direct vectors:'
    for i in range(3):
        print ("%7.3f "*3) % tuple(DV[i])

    print 'Please enter the direct vectors for large Brillouin zone'
    DVn = zeros(1)
    while shape(DVn) != (3,3):
        DVn = array(eval(sys.stdin.readline()))
    #DVn = array([[0.5,0.5,0],[0.5,-0.5,0],[0,0,1]])
    print 'But switching to direct vectors of the one atom unit cell:'
    for i in range(3):
        print ("%7.3f "*3) % tuple(DVn[i])

    # Reads div from case.klist
    firstline = open(w2k.case+'.klist').next()
    m = re.search('div: *\((\s*\d+\s*\d+\s*\d+)\)', firstline)
    if m is not None:
        div = map(int,m.groups()[0].split())

    Rs = GiveRs(div, DVn)
    print 'Real space vectors used in the Fourier transform to tigh-binding:'
    for i,R in enumerate(Rs):
        print "%3d"%i, "%7.3f %7.3f %7.3f" % tuple(R)

    # Read Udmft
    filename = 'Udmft.0'
    (nkp, nsymop, norbitals) = rdU.fopen(fhp, filename)

    nindo = rdU.read1(fhp, norbitals)

    print 'nkp=', nkp, 'nsymop=', nsymop, 'norbitals=', norbitals, 'nindo=', nindo

    delta_A_B = pos[1]-pos[0]
    print 'delta_A_B=', delta_A_B

    rdU.feopen(fhe, w2k.case+'.energy', strc.nat)

    Hk=[]
    kps=[]
    wgh=[]
    for ikp in range(nkp):

        (iikp, nbands, tmaxdim2, tnorbitals, nemin) = rdU.read2(fhp)
        (k, kname, wg, ios, nen) = rdU.feread1(fhe)

        print 'delta*k=', delta_A_B*k

        if ios!=0:
            print 'Number of k-points in case.energy is smaller then in Udfmt.0'
            break

        # Read LDA energies from case.energy
        Ek = rdU.feread2(fhe, nen) * Ry_in_eV - EF

        # finds minimum and maximum band for downfolding
        #(bnd_s, bnd_e) = FindMaxMinBand(Ek, options.Emin, nemin, options.nbnd, options.dnbnd)
        #e = Ek[nemin+bnd_s-1:nemin+bnd_e-1]
        #print 'bnd_s=', bnd_s, 'bnd_e=', bnd_e, 'Emm=', e[0], e[-1]
        #print 'k=', k

        for isym in range(nsymop):

            kn = dot(strc.iz[isym],k)
            k_nat = [sum(BRX[i,:]*kn) for i in range(3)]

            kps.append(kn)
            wgh.append(wg)

            iisym = rdU.read3(fhp)
            if iisym!=isym+1:
                print 'ERROR: isym!=iisym', isym, iisym

            for iorb in range(norbitals):

                U = []
                for ind in range(nindo[iorb]):
                    Udat = rdU.read4(fhp, nbands)
                    U.append( Udat )
                U = array(U)

                if iorb==0:
                    Utot = U
                else:
                    Utot = vstack((Utot,U))

            # finds minimum and maximum band for downfolding
            Nse = FindMaxMinBand(Ek, options.Emin, options.Emax)

            print 'len(Utot)=', len(Utot)

            while (Nse[1]-Nse[0] < len(Utot) ):
                print 'Warning: Number of bands in the interval is ', Nse[1]-Nse[0], 'while number of orbitals is ', len(Utot)
                (Nse, Succ) = TryAddBands(Nse, Ek, options.Emin, options.Emax, options.SmearN, options.SmearX)
                if not Succ: sys.exit(1)

            NsingularValues=0
            while (NsingularValues<len(Utot)):
                (bnd_s,bnd_e) = Nse-(nemin-1)
                e = Ek[nemin-1+bnd_s:nemin-1+bnd_e]
                print 'bnd_s=', bnd_s, 'bnd_e=', bnd_e, 'n=', bnd_e-bnd_s, 'Emm=', e[0], e[-1]
                # Performing SVD on projector
                (u, s, v) = linalg.svd(Utot[:,bnd_s:bnd_e])
                print 'singular-values:', ikp, s.tolist()
                NsingularValues = len(filter(lambda x: x>options.smalls, s))
                print 'NsingularValues=',  NsingularValues

                if (NsingularValues<len(Utot)):
                    (Nse, Succ) = TryAddBands(Nse, Ek, options.Emin, options.Emax, options.SmearN, options.SmearX)
                    if not Succ: sys.exit(1)

            # downfolding transformation
            Rt = dot(u[:,:len(s)],v[:len(s),:])
            # downfolded hamiltonian
            tHk = dot(conjugate(Rt) * e, transpose(Rt))

            N0 = len(Utot)/2
            print 'N0=', N0

            zHk = copy.deepcopy(tHk[:N0,:N0])
            zHk += tHk[:N0,:N0]
            zHk += tHk[:N0,N0:]*exp(2*pi*1j*sum(delta_A_B*kn) ) # 1,2
            zHk += tHk[N0:,:N0]*exp(-2*pi*1j*sum(delta_A_B*kn) )# 2,1
            zHk *= 0.5
            Hk.append(zHk)

    rdU.fclose(fhp)
    rdU.feclose(fhe)

    Hk = array(Hk)

    nrm = sum(wgh)
    wgh = wgh/nrm

    # Fourier transform to produce hopping parameters
    Hr=[]
    for R in Rs:
        tHr=zeros(shape(Hk[0]),dtype=complex)
        dsum=0j
        for ik in range(len(kps)):
            expk = exp(1j*2*pi*dot(kps[ik],R)) * wgh[ik]
            tHr += expk * Hk[ik]
            dsum += expk
        print 'R=[', "%4.1f, "*3 % tuple(R), '] exp=', "%10.5f" % abs(dsum)
        Hr.append(tHr)

    # Printing of the Hamiltonian
    fh = open('hamiltonian.dat', 'w')

    print >> fh, '# Tighbinding Hamiltonian with entries: ', inl.legends
    print >> fh
    #print >> fh, "lattice_unit=[", tuple(DVn[0]), ',', tuple(DVn[1]), ',', tuple(DVn[2]), "]"
    print >> fh, "lattice_unit=", identity(3).tolist()

    print >> fh, "orbital_positions=",zeros((3,3)).tolist()
    print >> fh
    #print '['
    #for atm in inl.siginds.keys():
    #    for i in range(len(inl.siginds[atm])):
    #        print >> fh, pos[atm-1].tolist(),','
    #print
    #print >> fh, ']\n'

    # sorting such that first come larger hoppings
    for irm,R in enumerate(Rs):
        if sum(abs(Rs[irm]))<1e-10: break

    ir_index = range(len(Rs))
    ir_index.sort( lambda a,b: cmp(sum(abs(Hr[b])), sum(abs(Hr[a]))) )
    ir0 = ir_index.index(irm)
    ir_index = [irm] + ir_index[:ir0] + ir_index[ir0+1:]

    #print ir_index

    DVn_inv = linalg.inv(DVn)
    print >> fh, 'Hopping={}'
    for ir in ir_index :
        # RS*DV^-1 = (i0,i1,i2)
        Rn = dot(Rs[ir] ,DVn_inv)
        Print(fh, Rn, Hr[ir])
        print >> fh

    fh.close()


# -*- coding:utf8 -*-
# @TIME     : 2021/4/12
# @Author   : SuHao

# https://github.com/BobLiu20/YOLOv3_PyTorch/blob/master/nets/yolo_loss.py

import math
import torch
import torch.nn as nn
from utils.parse_config import parse_model_config
import numpy as np

def bbox_iou(box1, box2, x1y1x2y2=True):
    """
    Returns the IoU of two bounding boxes
    """
    if not x1y1x2y2:
        # Transform from center and width to exact coordinates
        b1_x1, b1_x2 = box1[:, 0] - box1[:, 2] / 2, box1[:, 0] + box1[:, 2] / 2
        b1_y1, b1_y2 = box1[:, 1] - box1[:, 3] / 2, box1[:, 1] + box1[:, 3] / 2
        b2_x1, b2_x2 = box2[:, 0] - box2[:, 2] / 2, box2[:, 0] + box2[:, 2] / 2
        b2_y1, b2_y2 = box2[:, 1] - box2[:, 3] / 2, box2[:, 1] + box2[:, 3] / 2
    else:
        # Get the coordinates of bounding boxes
        b1_x1, b1_y1, b1_x2, b1_y2 = box1[:,0], box1[:,1], box1[:,2], box1[:,3]
        b2_x1, b2_y1, b2_x2, b2_y2 = box2[:,0], box2[:,1], box2[:,2], box2[:,3]

    # get the corrdinates of the intersection rectangle
    inter_rect_x1 =  torch.max(b1_x1, b2_x1)
    inter_rect_y1 =  torch.max(b1_y1, b2_y1)
    inter_rect_x2 =  torch.min(b1_x2, b2_x2)
    inter_rect_y2 =  torch.min(b1_y2, b2_y2)
    # Intersection area
    inter_area =    torch.clamp(inter_rect_x2 - inter_rect_x1 + 1, min=0) * \
                    torch.clamp(inter_rect_y2 - inter_rect_y1 + 1, min=0)
    # Union Area
    b1_area = (b1_x2 - b1_x1 + 1) * (b1_y2 - b1_y1 + 1)
    b2_area = (b2_x2 - b2_x1 + 1) * (b2_y2 - b2_y1 + 1)

    iou = inter_area / (b1_area + b2_area - inter_area + 1e-16)

    return iou

class YOLOLoss(nn.Module):
    def __init__(self, config_path, yolo=0):
        '''
        :param config_path: 从cfg配置文件中解析各种参数
        :param yolo 为0,1,2 指定YOLO层
        '''
        super(YOLOLoss, self).__init__()
        self.module_defs = parse_model_config(config_path)                              # 解析网络配置文件
        self.hyperparams = self.module_defs.pop(0)              # 超参数
        self.anchors = []                                       # 存放anchor的长宽
        for module_def in self.module_defs:
            if module_def["type"] == "yolo":                          # 寻找YOLO块
                if module_def["mask"][0] == str(3*yolo):           # 寻找是否是对应序号的yolo块
                    anchor_idxs = [int(x) for x in module_def["mask"].split(",")]       # Anchor的序号，yolov3中每个特征图有3个Anchor
                    anchors = [int(x) for x in module_def["anchors"].split(",")]
                    anchors = [(anchors[i], anchors[i + 1]) for i in range(0, len(anchors), 2)]
                    self.anchors = [anchors[i] for i in anchor_idxs]         # 提取3个Anchor
                    break
        self.num_anchors = len(self.anchors)                    # anchor的数量
        self.num_classes = int(module_def["classes"])           # 类别的数量
        self.img_dim = int(self.hyperparams["width"])
        self.ignore_thres = float(module_def["ignore_thresh"])
        self.obj_scale = 1
        self.noobj_scale = 1
        self.metrics = {}
        self.grid_size = torch.tensor(0)  # grid size

        self.lambda_xy = 2.5
        self.lambda_wh = 2.5
        self.lambda_conf = 1.0
        self.lambda_cls = 1.0

        self.mse_loss = nn.MSELoss()            # 最小均方误差
        self.bce_loss = nn.BCELoss()          # 二分类交叉熵

    def forward(self, pred, targets=None):
        # yolo层的输入是特征图，即pred是特征图，维度是样本数量*(（5+类别数量）*3)*13*13
        # Tensors for cuda support
        if targets is not None:
            FloatTensor = torch.cuda.FloatTensor if pred.is_cuda else torch.FloatTensor
            LongTensor = torch.cuda.LongTensor if pred.is_cuda else torch.LongTensor
            num_samples = pred.size(0)      # 样本数量
            import math
            grid_size = int(math.sqrt(pred.size(2) / self.num_anchors))       # 此时网格尺寸应该是13或26或52
            self.stride = self.img_dim / grid_size
            prediction = (
                pred.view(num_samples, self.num_anchors, grid_size, grid_size, self.num_classes + 5)        # pred的维度是n*(self.num_anchors*grid_size*grid_size)*85
                    .contiguous()
                #  当调用contiguous()时，会强制拷贝一份tensor，让它的布局和从头创建的一模一样，但是两个tensor完全没有联系。
            )
            pred_boxes = prediction[..., 0:4] / self.stride
            x = pred_boxes[..., 0] - pred_boxes[..., 0].floor()
            y = pred_boxes[..., 1] - pred_boxes[..., 1].floor()
            w = pred_boxes[..., 2]
            h = pred_boxes[..., 3]
            scaled_anchors = FloatTensor([(a_w/self.stride, a_h/self.stride) for a_w, a_h in self.anchors])    # 对Anchor的坐标位置进行缩放
            anchor_w = scaled_anchors[:, 0:1].view((1, self.num_anchors, 1, 1))
            anchor_h = scaled_anchors[:, 1:2].view((1, self.num_anchors, 1, 1))
            w = torch.log(w / anchor_w + 1e-16)
            h = torch.log(h / anchor_h + 1e-16)
            conf = prediction[..., 4]
            pred_cls = prediction[..., 5:]
            #  build target
            scaled_anchors = [(a_w/self.stride, a_h/self.stride) for a_w, a_h in self.anchors]
            mask, noobj_mask, tx, ty, tw, th, tconf, tcls = self.get_target(targets, scaled_anchors,
                                                                           grid_size, grid_size,
                                                                           self.ignore_thres)
            mask, noobj_mask = mask.cuda(), noobj_mask.cuda()
            tx, ty, tw, th = tx.cuda(), ty.cuda(), tw.cuda(), th.cuda()
            tconf, tcls = tconf.cuda(), tcls.cuda()
            #  losses.
            loss_x = self.bce_loss(x * mask, tx * mask)
            loss_y = self.bce_loss(y * mask, ty * mask)
            loss_w = self.mse_loss(w * mask, tw * mask)
            loss_h = self.mse_loss(h * mask, th * mask)
            loss_conf = self.bce_loss(conf * mask, mask) + \
                0.5 * self.bce_loss(conf * noobj_mask, noobj_mask * 0.0)
            loss_cls = self.bce_loss(pred_cls[mask == 1], tcls[mask == 1])
            #  total loss = losses * weight
            loss = loss_x * self.lambda_xy + loss_y * self.lambda_xy + \
                loss_w * self.lambda_wh + loss_h * self.lambda_wh + \
                loss_conf * self.lambda_conf + loss_cls * self.lambda_cls
            return loss
        else:
            return pred

    def get_target(self, target, anchors, in_w, in_h, ignore_threshold):
        bs = target.size(0)

        mask = torch.zeros(bs, self.num_anchors, in_h, in_w, requires_grad=False)
        noobj_mask = torch.ones(bs, self.num_anchors, in_h, in_w, requires_grad=False)
        tx = torch.zeros(bs, self.num_anchors, in_h, in_w, requires_grad=False)
        ty = torch.zeros(bs, self.num_anchors, in_h, in_w, requires_grad=False)
        tw = torch.zeros(bs, self.num_anchors, in_h, in_w, requires_grad=False)
        th = torch.zeros(bs, self.num_anchors, in_h, in_w, requires_grad=False)
        tconf = torch.zeros(bs, self.num_anchors, in_h, in_w, requires_grad=False)
        tcls = torch.zeros(bs, self.num_anchors, in_h, in_w, self.num_classes, requires_grad=False)
        for b in range(bs):
            for t in range(target.shape[1]):
                if target[b, t].sum() == 0:
                    continue
                # Convert to position relative to box
                gx = target[b, t, 1] * in_w
                gy = target[b, t, 2] * in_h
                gw = target[b, t, 3] * in_w
                gh = target[b, t, 4] * in_h
                # Get grid box indices
                gi = int(gx)
                gj = int(gy)
                # Get shape of gt box
                gt_box = torch.FloatTensor(np.array([0, 0, gw, gh])).unsqueeze(0)
                # Get shape of anchor box
                anchor_shapes = torch.FloatTensor(np.concatenate((np.zeros((self.num_anchors, 2)),
                                                                  np.array(anchors)), 1))
                # Calculate iou between gt and anchor shapes
                anch_ious = bbox_iou(gt_box, anchor_shapes)
                # Where the overlap is larger than threshold set mask to zero (ignore)
                noobj_mask[b, anch_ious > ignore_threshold, gj, gi] = 0
                # Find the best matching anchor box
                best_n = np.argmax(anch_ious)

                # Masks
                mask[b, best_n, gj, gi] = 1
                # Coordinates
                tx[b, best_n, gj, gi] = gx - gi
                ty[b, best_n, gj, gi] = gy - gj
                # Width and height
                tw[b, best_n, gj, gi] = math.log(gw / anchors[best_n][0] + 1e-16)
                th[b, best_n, gj, gi] = math.log(gh / anchors[best_n][1] + 1e-16)
                # object
                tconf[b, best_n, gj, gi] = 1
                # One-hot encoding of label
                tcls[b, best_n, gj, gi, int(target[b, t, 0])] = 1

        return mask, noobj_mask, tx, ty, tw, th, tconf, tcls

class YOLOLoss_total(nn.Module):
    def __init__(self, config_path):
        super(YOLOLoss_total, self).__init__()
        self.config_path = config_path

    def forward(self, x, targets=None):
        self.L = len(x)
        self.loss = 0
        self.Loss_list = []
        self.outputs = []
        self.outs = []
        self.params = {"anchors": [], "strides": []}
        if targets is not None:
            for i in range(self.L):
                self.Loss_list.append(YOLOLoss(self.config_path, yolo=self.L - 1 - i))
                loss = self.Loss_list[-1](x[i], targets)
                self.loss += loss
            return self.loss
        else:
            for i in range(self.L):
                self.Loss_list.append(YOLOLoss(self.config_path, yolo=self.L - 1 - i))
                output = self.Loss_list[-1](x[i], targets)
                self.outputs.append(output)
            self.outputs = torch.cat(self.outputs, 2)
            self.outputs = torch.squeeze(self.outputs, dim=1)
            return self.outputs


import torch
import numpy as np
from PIL import Image
import os
import cv2

model = torch.jit.load('zak_test.torchscript')

path = '/home/zak/Documents/WayveExperimentalAi/temp/ningaloo--2021-01-28-08-40-21--4f1a733c-3c36-4bd5-a687-ce1dd562fcd5/cameras/front-forward--rgb'
# path = '/home/zak/WayveData/remote/naas_image_storage/ningaloo-images/ningaloo--1.3.0044--default/2020-10-26--00-53-08--naas--follow-traffic-light-v3--fe869bbf/cameras/front-forward/rgb'
# path = '/home/zak/Documents/contrastive-unpaired-translation/datasets/grumpifycat/trainA'

fnames = os.listdir(path)

fnames = sorted(fnames)
# fnames = ['%dunixus.jpeg'%ts for ts in sorted([int(fname[:-11]) for fname in fnames])][12:][::3]

for i, fname in enumerate(fnames):
    print(fname)
    img = Image.open(os.path.join(path, fname))
    # img = img.resize((img.size[0]//2, img.size[1]//2))

    # img = img.resize((256,256))
    img = np.asarray(img)
    imgX = torch.as_tensor(img).permute(2,0,1).unsqueeze(0).cuda()
    with torch.no_grad():
        out = model(imgX)
    out = out[0].permute(1,2,0)
    out = out * .5 + .5
    out = (out*255).clamp(0,255).byte().cpu().numpy()

    out = np.hstack((img, out))
    out = out[:,:,::-1]
    cv2.imshow('', out)
    # cv2.imwrite('temp2x/%08d.jpg'%i, out)
    cv2.waitKey(1)


""" In this example, we're going to use NEvoPy's implementation of a classic
fixed-topology neuroevolution algorithm to learn how to play Flappy Bird!

We're going to use the "FlappyBird-v0" gym environment (available here:
https://github.com/Talendar/flappy-bird-gym) to simulate the game.

Note:
    When I ran this example, I wasn't able to get a neural network capable of
    playing the game. This example is listed here just to demonstrate how to use
    a fixed-topology genome.

Authors:
    Gabriel Nogueira (Talendar)
"""

import os

import flappy_bird_gym
import matplotlib.pyplot as plt
import nevopy as ne
import numpy as np
import skimage.color
import skimage.transform

# In some scenarios, using a GPU with TensorFlow during neuroevolution might
# slow things down considerably. Since in this example we're dealing with small
# images in batches of 1 item, it's not worth using the GPU. The line below
# makes TensorFlow to use only CPUs.
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

# Settings for the ray processing scheduler.
RAY_CONFIG = dict(
    num_gpus=0,
    worker_gpu_frac=0,
)

# Defining a base genome for our population. It will serve as a model for the
# other genomes of the population, which will have the same topology as it.
BASE_GENOME = ne.fixed_topology.FixedTopologyGenome(

    # List with the layers of the base genome. Other genomes in the population
    # will have similar layers (with the same topology).
    layers=[
        ne.fixed_topology.layers.TFConv2DLayer(
            filters=8,
            kernel_size=(10, 10),
            strides=(5, 5),
            mating_func=ne.fixed_topology.layers.mating.exchange_units_mating,
        ),
        ne.fixed_topology.layers.TFFlattenLayer(),
        ne.fixed_topology.layers.TFDenseLayer(32, activation="relu"),
        ne.fixed_topology.layers.TFDenseLayer(1, activation="sigmoid")
    ],
)

class GymPreProcessImgCallback(ne.utils.GymCallback):
    """ We'll use this callback to pre-process the images before feeding them
    to the genomes.
    """

    def on_obs_processing(self, wrapped_obs):
        img = skimage.color.rgb2gray(wrapped_obs.value)
        img = np.moveaxis(img, source=1, destination=0)  # width <-> height
        img = img[10:-90]

        w, h = img.shape[0] // 6, img.shape[1] // 6
        img = skimage.transform.resize(img, (w, h))

        img = img.astype(np.float32)
        img = np.expand_dims(img, axis=0)
        img = np.expand_dims(img, axis=-1)

        wrapped_obs.value = img

def make_env():
    """ Makes a new gym 'FlappyBird-rgb-v0' environment. """
    return flappy_bird_gym.make("FlappyBird-rgb-v0")

if __name__ == "__main__":
    # Instantiating our GymPreProcessImgCallback:
    pre_proc_img_gym_cb = GymPreProcessImgCallback()

    # Let's start by visualizing some stuff.
    with make_env() as test_env:
        # Obtaining a raw input sample:
        sample_input = ne.utils.MutableWrapper(test_env.reset())
        print(f"Raw input shape: {sample_input.value.shape}")

        # Pre-processing the input with our callback:
        pre_proc_img_gym_cb.on_obs_processing(sample_input)
        print(f"Pre-processed input shape: {sample_input.value.shape}")
        plt.imshow(sample_input.value[0], cmap="gray")
        plt.show()

        # Feeding the sample image to our base genome, so it can register
        # the shape of the inputs:
        print(f"Sample output: {BASE_GENOME(sample_input.value)}")

        # Visualizing our base genome's topology:
        print("\nVisualizing base genome's topology...")
        BASE_GENOME.visualize()

    # NEvoPy provides a general fitness function that works with most gym
    # environments, so we don't have to worry about implementing it.
    fitness_function = ne.utils.GymFitnessFunction(
        make_env=make_env,
        default_num_episodes=5,
        callbacks=[pre_proc_img_gym_cb],
    )

    # Creating a new random population of fixed-topology genomes:
    # (we're going to use the default settings)
    population = ne.genetic_algorithm.GeneticPopulation(
        size=64,
        base_genome=BASE_GENOME,
        processing_scheduler=ne.processing.RayProcessingScheduler(**RAY_CONFIG),
    )

    # We can tell NEvoPy to halt the evolutionary process when a certain fitness
    # value has been achieved by the population. For that, we use a specific
    # type of callback.
    early_stopping_cb = ne.callbacks.FitnessEarlyStopping(
        fitness_threshold=5000,
        min_consecutive_generations=5,
    )

    # Evolving the population (this might take some minutes):
    history = population.evolve(generations=25,
                                fitness_function=fitness_function,
                                callbacks=[early_stopping_cb])

    # Retrieving the best genome in the population:
    best_genome = population.fittest()

    # Visualizing the topology of our best genome:
    best_genome.visualize()

    # Visualizing the progression of the population's fitness:
    history.visualize(log_scale=False)

    # Now, let's compare the performance of our best genome with the performance
    # of a random agent.
    print("\nEvaluation (25 episodes):")
    print(f". Random genome score: {fitness_function(None, num_eps=25)}")
    print(f". Evolved genome: {fitness_function(best_genome, num_eps=25)}")

    # Visualizing our best genome and a random agent playing:
    while True:
        key = int(input("\nEvolution complete!\n"
                        "   [0] Visualize random genome\n"
                        "   [1] Visualize evolved genome\n"
                        "   [2] Exit\n"
                        "Choose an option: "))
        if key == 0:
            print(f"\nRandom genome: "
                  f"{fitness_function(None, num_eps=1, visualize=True)}")
        elif key == 1:
            print(f"\nEvolved genome: "
                  f"{fitness_function(best_genome, num_eps=1, visualize=True)}")
        elif key == 2:
            break
        else:
            print("\nInvalid key!")