DESC = '''
Helper script to record/examine a rollout (record/ render/ playback/ recover) on an environment\n
  > Examine options:\n
    - Record:   Record an execution. (Useful for kinesthetic demonstrations on hardware)\n
    - Render:   Render back the execution. (sim.forward)\n
    - Playback: Playback the rollout action sequence in openloop (sim.step(a))\n
    - Recover:  Playback actions recovered from the observations \n
  > Render options\n
    - either onscreen, or offscreen, or just rollout without rendering.\n
  > Save options:\n
    - save resulting rollouts as RoboHive/Roboset format, and as 2D plots\n
USAGE:\n
    $ python examine_rollout.py --env_name door-v0 \n
    $ python examine_rollout.py --env_name door-v0 --rollout_path my_rollouts.h5 --repeat 10 \n
    $ python logger/examine_logs.py --env_name rpFrankaRobotiqData-v0 --rollout_path teleOp_trace.h5 --rollout_format RoboSet --render offscreen --compress_paths False -c left_cam -c right_cam -c top_cam -c Franka_wrist_cam --plot_paths True
'''

from myosuite.utils.paths_utils import plot as plotnsave_paths
from myosuite.utils import tensor_utils
import gym
import click
import numpy as np
import time
import os

@click.command(help=DESC)
@click.option('-e', '--env_name', type=str, help='environment to load', required=True)
@click.option('-p', '--rollout_path', type=str, help='absolute path of the rollout', default=None)
@click.option('-f', '--rollout_format', type=click.Choice(['RoboHive', 'RoboSet']), help='Data format', default='RoboHive')
@click.option('-m', '--mode', type=click.Choice(['record', 'render', 'playback', 'recover']), help='How to examine rollout', default='playback')
@click.option('-h', '--horizon', type=int, help='Rollout horizon, when mode is record', default=-1)
@click.option('-s', '--seed', type=int, help='seed for generating environment instances', default=123)
@click.option('-n', '--num_repeat', type=int, help='number of repeats for the rollouts', default=1)
@click.option('-r', '--render', type=click.Choice(['onscreen', 'offscreen', 'none']), help='visualize onscreen or offscreen', default='onscreen')
@click.option('-c', '--camera_name', multiple=True, type=str, default=[None,], help=('list of camera names for rendering'))
@click.option('-fs', '--frame_size', type=tuple, default=(424, 240), help=('Camera frame size for rendering'))
@click.option('-o', '--output_dir', type=str, default='./', help=('Directory to save the outputs'))
@click.option('-on', '--output_name', type=str, default=None, help=('The name to save the outputs as'))
@click.option('-sp', '--save_paths', type=bool, default=False, help=('Save the rollout paths'))
@click.option('-cp', '--compress_paths', type=bool, default=True, help=('compress paths. Remove obs and env_info/state keys'))
@click.option('-pp', '--plot_paths', type=bool, default=False, help=('2D-plot of individual paths'))
@click.option('-ea', '--env_args', type=str, default=None, help=('env args. E.g. --env_args "{\'is_hardware\':True}"'))
@click.option('-ns', '--noise_scale', type=float, default=0.0, help=('Noise amplitude in randians}'))
@click.option('-ie', '--include_exteroception', type=bool, default=False, help=('Include exteroception'))
def examine_logs(env_name, rollout_path, rollout_format, mode, horizon, seed, num_repeat, render, camera_name, frame_size, output_dir, output_name, save_paths, compress_paths, plot_paths, env_args, noise_scale, include_exteroception):

    # seed and load environments
    np.random.seed(seed)
    env = gym.make(env_name) if env_args==None else gym.make(env_name, **(eval(env_args)))
    env = env.env
    env.seed(seed)

    # Start a "trace" for recording rollouts
    if rollout_format=='RoboHive':
        from myosuite.logger.grouped_datasets import Trace
    elif rollout_format=='RoboSet':
        from myosuite.logger.roboset_logger import RoboSet_Trace as Trace
    else:
        raise TypeError("unknown rollout_format format")
    trace = Trace("Rollouts")

    # Load old traces as "paths"; none if record
    if mode == 'record':
        assert horizon>0, "Rollout horizon must be specified when recording rollout"
        assert output_name is not None, "Specify the name of the recording"
        if save_paths is False:
            print("Warning: Recording is not being saved. Enable save_paths=True to log the recorded path")
        paths = [None,]*num_repeat # Mark old traces as None
    else:
        assert rollout_path is not None, "Rollout path is required for mode:{} ".format(mode)
        if output_dir == './': # overide the default
            output_dir = os.path.dirname(rollout_path)
        if output_name is None: # default to the rollout name
            rollout_name = os.path.split(rollout_path)[-1]
            output_name, output_type = os.path.splitext(rollout_name)
        paths = Trace.load(trace_path=rollout_path, trace_type=rollout_format)

    # Resolve rendering
    if render == 'onscreen':
        env.mujoco_render_frames = True
    elif render =='offscreen':
        env.mujoco_render_frames = False
    elif render == None:
        env.mujoco_render_frames = False

    # Rollout paths
    for i_loop in range(num_repeat):

        # Rollout path
        print("Starting rollout loop:{}".format(i_loop))
        for path_name, path_data in paths.items():

            # initialize path -----------------------------
            ep_t0 = time.time()
            path_name+='-'+str(i_loop)
            print("Starting {} rollout".format(path_name))
            trace.create_group(path_name)

            # init: reset to starting state
            if path_data and rollout_format=='RoboHive':
                # reset to init state
                if "state" in path_data['env_infos'].keys():
                    path_state = tensor_utils.split_tensor_dict_list(path_data['env_infos']['state'])
                    env.reset(reset_qpos=path_state[0]['qpos'], reset_qvel=path_state[0]['qvel'])
                    env.set_env_state(path_state[0])
                else:
                    env.reset()
            elif path_data and rollout_format=='RoboSet':
                path_data = path_data['data']
                # reset to init state
                reset_qpos = env.init_qpos.copy()
                nq_arm = len(path_data['qp_arm'][0])
                reset_qpos[:nq_arm] = path_data['qp_arm'][0]
                nq_ee = len(path_data['qp_ee'][0])
                reset_qpos[nq_arm:nq_arm+nq_ee] = path_data['qp_ee'][0]
                env.reset(reset_qpos=reset_qpos)
            elif mode=="record":
                env.reset()
            else:
                raise TypeError(f"Unknown rollout_format{rollout_format} / mode:{mode}")
            trace_horizon = horizon if mode=='record' else path_data['time'].shape[0]-1

            # Rollout path --------------------------------
            obs, rwd, done, env_info = env.forward(update_exteroception=include_exteroception)
            ep_rwd = rwd
            for i_step in range(trace_horizon+1):

                # Get step's actions ----------------------

                # Record Execution. Useful for kinesthetic demonstrations on hardware
                if mode=='record':
                    act = env.action_space.sample() # dummy random sample

                # Directly create the scene
                elif mode=='render':
                    # populate actions merely for logging
                    if rollout_format=='RoboSet':
                        act = np.concatenate([path_data['ctrl_arm'][i_step], path_data['ctrl_ee'][i_step]])
                    elif rollout_format=='RoboHive' and "state" in path_data['env_infos'].keys():
                        act = path_data['actions'][i_step]
                    else:
                        raise NotImplementedError("Settings not found")

                # Apply actions in open loop
                elif mode=='playback':
                    if rollout_format=='RoboSet':
                        act = np.concatenate([path_data['ctrl_arm'][i_step], path_data['ctrl_ee'][i_step]])
                    elif rollout_format=='RoboHive':
                        act = path_data['actions'][i_step]

                # Recover actions from states
                elif mode=='recover':
                    # assumes position controls
                    if rollout_format=='RoboSet':
                        act = np.concatenate([path_data['qp_arm'][i_step], path_data['qp_ee'][i_step]])
                    elif rollout_format=='RoboHive':
                        act = path_data['env_infos']['obs_dict']['qp'][i_step]
                    if noise_scale:
                        act = act + env.np_random.uniform(high=noise_scale, low=-noise_scale, size=len(act)).astype(act.dtype)
                    if env.normalize_act:
                        act = env.robot.normalize_actions(controls=act)

                # nan actions for last log entry
                if i_step == trace_horizon:
                    act = np.nan*np.ones(env.action_space.shape)

                # log values at time=t ----------------------------------
                if compress_paths:
                    obs = [] # don't save obs, env_infos has obs_dict
                    if 'state' in env_info.keys(): del env_info['state']  # don't save state, obs_dict has env necessities

                # log: time, obs, act, rwd, info, done
                datum_dict = dict(
                        time=env.time,
                        observations=obs,
                        actions=act.copy(),
                        rewards=rwd,
                        env_infos=env_info,
                        done=done,
                    )
                trace.append_datums(group_key=path_name,dataset_key_val=datum_dict)

                # log: offscreen frames
                if render =='offscreen':
                    for cam in camera_name:
                        curr_frame = env.sim.renderer.render_offscreen(
                            camera_id=cam,
                            width=frame_size[0],
                            height=frame_size[1],
                            device_id=0
                        )
                        # cam_key = cam+'_rgb'
                        trace.append_datum(group_key=path_name,dataset_key=cam, dataset_val=curr_frame)

                # step/forward env using actions from t=>t+1 ----------------------
                if mode=='render' and i_step < trace_horizon:
                    if rollout_format=='RoboSet':
                        env.sim.data.qpos[:nq_arm]= path_data['qp_arm'][i_step+1]
                        env.sim.data.qpos[nq_arm:nq_arm+nq_ee]= path_data['qp_ee'][i_step+1]
                        env.sim.data.qvel[:nq_arm]= path_data['qv_arm'][i_step+1]
                        env.sim.data.qvel[nq_arm:nq_arm+nq_ee]= path_data['qv_ee'][i_step+1]
                        env.sim.data.time = path_data['time'][i_step+1]
                    elif rollout_format=='RoboHive' and "state" in path_data['env_infos'].keys():
                        env.set_env_state(path_state[i_step+1])
                    else:
                        raise NotImplementedError("Settings not found")
                    obs, rwd, done, env_info = env.forward(update_exteroception=include_exteroception)
                    ep_rwd += rwd
                elif i_step < trace_horizon: # incase last step actions (nans) can cause issues in step
                    obs, rwd, done, env_info = env.step(act, update_exteroception=include_exteroception)
                    ep_rwd += rwd

            # save offscreen buffers as video and clear the dataset
            if render == 'offscreen':
                trace.render(output_dir=output_dir, output_format="mp4", groups=[path_name], datasets=camera_name, input_fps=1/env.dt)
                trace.remove_dataset(group_keys=[path_name], dataset_key=camera_name)

            # Finish rollout
            old_stat = f"(Old rollout rewards: {sum(path_data['rewards'][:])})" if path_data else ""
            print(f"Finishing {path_name[:-2]} rollout in {(time.time()-ep_t0):0.4} sec. Total rewards {ep_rwd} "+old_stat)

        # Finish loop
        print("Finished rollout loop:{}".format(i_loop))

    # plot paths ???: Needs upgrade to the new logger
    trace.stack()
    time_stamp = time.strftime("%Y%m%d-%H%M%S")
    # plot paths
    if plot_paths:
        file_name = os.path.join(output_dir, output_name + '{}'.format(time_stamp))
        plotnsave_paths(trace.trace, env=env, fileName_prefix=file_name)

    # Close and save paths
    trace.close()
    if save_paths:
        file_name = os.path.join(output_dir, output_name + '{}_paths.h5'.format(time_stamp))
        trace.save(trace_name=file_name)

if __name__ == '__main__':
    examine_logs()


import tensorflow as tf
from tensorflow.keras import layers

class ResidualBlock(layers.Layer):
    def __init__(self, **kwargs):
        super(ResidualBlock, self).__init__(**kwargs)

    def build(self, input_shape):
        """
        input_shape = (batch_size, height, width, channels)
        """
        channels = input_shape.as_list()[-1]
        self.conv1 = layers.Conv2D(channels, (3, 3), strides=(1, 1), padding='same', use_bias=False, name='conv1')
        self.conv2 = layers.Conv2D(channels, (3, 3), strides=(1, 1), padding='same', use_bias=False, name='conv2')
        self.instance_norm1 = layers.LayerNormalization(axis=(1, 2), epsilon=1e-5, name='i_norm1')
        self.instance_norm2 = layers.LayerNormalization(axis=(1, 2), epsilon=1e-5, name='i_norm2')
        super(ResidualBlock, self).build(input_shape)  # Be sure to call this at the end

    def call(self, inputs, **kwargs):
        """
        inputs - shape = (batch_size, height, width, channels)
        """
        x = self.conv1(inputs)  # (batch_size, height, width, channels)
        x = self.instance_norm1(x)
        x = layers.ReLU()(x)
        x = self.conv2(x)  # (batch_size, height, width, channels)
        x = self.instance_norm2(x)

        x0 = inputs
        return x + x0  # (batch_size, height, width, channels)

def make_generator_model(label_dim, g_conv_dim):
    # Concatenate inputs
    input_img = layers.Input(shape=(128, 128, 3), name='gen_input_img')  # (None, 128, 128, 3)

    input_lbl = layers.Input(shape=(label_dim,), name='gen_input_lbl')  # (None, label_dim)
    lbl_reshape = layers.Reshape((1, 1, label_dim))(input_lbl)  # (None, 1, 1, label_dim)
    lbl_stack = layers.Lambda(lambda x: tf.tile(x, (1, 128, 128, 1)))(lbl_reshape)  # (None, 128, 128, label_dim)

    input_concat = layers.Concatenate(name='gen_input_concat')([input_img, lbl_stack])  # (None, 128, 128, 3+label_dim)

    # Downsampling part

    x = layers.Conv2D(g_conv_dim, (7, 7), strides=(1, 1), padding='same',
                      use_bias=False, name='gen_conv1')(input_concat)  # (None, 128, 128, g_conv_dim)
    x = layers.LayerNormalization(axis=(1, 2), epsilon=1e-5, name='gen_i_norm1')(x)
    x = layers.ReLU()(x)

    x = layers.Conv2D(g_conv_dim * 2, (4, 4), strides=(2, 2), padding='same',
                      use_bias=False, name='gen_conv2')(x)  # (None, 64, 64, g_conv_dim*2)
    x = layers.LayerNormalization(axis=(1, 2), epsilon=1e-5, name='gen_i_norm2')(x)
    x = layers.ReLU()(x)

    x = layers.Conv2D(g_conv_dim * 4, (4, 4), strides=(2, 2), padding='same',
                      use_bias=False, name='gen_conv3')(x)  # (None, 32, 32, g_conv_dim*4)
    x = layers.LayerNormalization(axis=(1, 2), epsilon=1e-5, name='gen_i_norm3')(x)
    x = layers.ReLU()(x)

    # Bottleneck part
    x = ResidualBlock(name='gen_res_block1')(x)
    x = ResidualBlock(name='gen_res_block2')(x)
    x = ResidualBlock(name='gen_res_block3')(x)
    x = ResidualBlock(name='gen_res_block4')(x)
    x = ResidualBlock(name='gen_res_block5')(x)
    x = ResidualBlock(name='gen_res_block6')(x)  # (None, 32, 32, g_conv_dim*4)

    # Upsampling part
    x = layers.Conv2DTranspose(g_conv_dim * 2, (4, 4), strides=(2, 2), padding='same',
                               use_bias=False, name='gen_deconv1')(x)  # (None, 64, 64, g_conv_dim*2)
    x = layers.LayerNormalization(axis=(1, 2), epsilon=1e-5, name='gen_i_norm4')(x)
    x = layers.ReLU()(x)

    x = layers.Conv2DTranspose(g_conv_dim, (4, 4), strides=(2, 2), padding='same',
                               use_bias=False, name='gen_deconv2')(x)  # (None, 128, 128, g_conv_dim)
    x = layers.LayerNormalization(axis=(1, 2), epsilon=1e-5, name='gen_i_norm5')(x)
    x = layers.ReLU()(x)

    x = layers.Conv2D(3, (7, 7), strides=(1, 1), padding='same',
                      use_bias=False, name='gen_last_conv')(x)  # (None, 128, 128, 3)
    output_layer = layers.Activation('tanh')(x)

    return tf.keras.Model([input_img, input_lbl], output_layer)


import cv2
import glob
import matplotlib.pyplot as plt
from matplotlib.patches import Polygon
import numpy as np
import os

from moviepy.editor import VideoFileClip

# My humble tribute to Michael Jordan,
# the best basketball player ever.
np.random.seed(23)

def calibrate_camera(images, nx=9, ny=6):
    """"
    images = paths to the images used to calibrate the camera
    nx = Number of inside corners in x
    ny = Number of inside corners in y
    """

    # Prepare object points, like (0,0,0), (1,0,0), (2,0,0) ....,(6,5,0)
    objp = np.zeros((nx * ny, 3), np.float32)
    objp[:,:2] = np.mgrid[0:nx, 0:ny].T.reshape(-1, 2)

    # Arrays to store object points and image points from all the images.
    objpoints = [] # 3d points in real world space
    imgpoints = [] # 2d points in image plane.

    # Step through the list and search for chessboard corners
    img = None
    for fname in images:
        img = cv2.imread(fname)
        gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)

        # Find the chessboard corners
        ret, corners = cv2.findChessboardCorners(gray, (nx, ny), None)

        # If found, add object points, image points
        if ret == True:
            objpoints.append(objp)
            imgpoints.append(corners)

    # Camera calibration, given object points, image points, and the shape of the grayscale image
    if (len(objpoints) > 0):
        # Camera successfully calibrated.
        print("Camera successfully calibrated.")
        ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(objpoints, imgpoints, gray.shape[::-1], None, None)
    else:
        # Unable to calibrate the camera.
        print("Unable to calibrate the camera.")
        ret, mtx, dist, rvecs, tvecs = (None, None, None, None, None)

    return ret, mtx, dist, rvecs, tvecs

def draw_lane(img, warped_img, left_points, right_points, Minv):
    # Create an image to draw the lines on
    warp_zero = np.zeros_like(warped_img).astype(np.uint8)
    color_warp = np.dstack((warp_zero, warp_zero, warp_zero))

    # Recast the x and y points into usable format for cv2.fillPoly()
    left_fitx = left_points[0]
    right_fitx = right_points[0]
    ploty = left_points[1]

    pts_left = np.array([np.transpose(np.vstack([left_fitx, ploty]))])
    pts_right = np.array([np.flipud(np.transpose(np.vstack([right_fitx, ploty])))])
    pts = np.hstack((pts_left, pts_right))

    # Draw the lane onto the warped blank image
    cv2.fillPoly(color_warp, np.int_([pts]), (0,255, 0))

    # Warp the blank back to original image space using inverse perspective matrix (Minv)
    newwarp = cv2.warpPerspective(color_warp, Minv, (img.shape[1], img.shape[0]))

    # Combine the result with the original image
    return cv2.addWeighted(img, 1, newwarp, 0.3, 0)

def get_histogram(img):
    return np.sum(img[img.shape[0]//2:, :], axis=0)

def detect_lines(img, return_img=False):
    # Take a histogram of the bottom half of the image
    histogram = get_histogram(img)

    if return_img:
        # Create an output image to draw on and  visualize the result
        out_img = np.dstack((img, img, img))*255

    # Find the peak of the left and right halves of the histogram
    # These will be the starting point for the left and right lines
    midpoint = np.int(histogram.shape[0]//2)
    leftx_base = np.argmax(histogram[:midpoint])
    rightx_base = np.argmax(histogram[midpoint:]) + midpoint

    # Choose the number of sliding windows
    nwindows = 9

    # Set height of windows
    window_height = np.int(img.shape[0]//nwindows)

    # Identify the x and y positions of all nonzero pixels in the image
    nonzero = img.nonzero()
    nonzerox = np.array(nonzero[1])
    nonzeroy = np.array(nonzero[0])

    # Current positions to be updated for each window
    leftx_current = leftx_base
    rightx_current = rightx_base

    # Set the width of the windows +/- margin
    margin = 100

    # Set minimum number of pixels found to recenter window
    minpix = 50

    # Create empty lists to receive left and right lane pixel indices
    left_lane_inds = []
    right_lane_inds = []

    # Step through the windows one by one
    for window in range(nwindows):
        # Identify window boundaries in x and y (and right and left)
        win_y_low = img.shape[0] - (window + 1)*window_height
        win_y_high = img.shape[0] - window*window_height
        win_xleft_low = leftx_current - margin
        win_xleft_high = leftx_current + margin
        win_xright_low = rightx_current - margin
        win_xright_high = rightx_current + margin

        if return_img:
            # Draw the windows on the visualization image
            cv2.rectangle(out_img,(win_xleft_low,win_y_low),(win_xleft_high,win_y_high), (0,255,0), 3)
            cv2.rectangle(out_img,(win_xright_low,win_y_low),(win_xright_high,win_y_high), (0,255,0), 3)

        # Identify the nonzero pixels in x and y within the window
        good_left_inds = ((nonzeroy >= win_y_low) & (nonzeroy < win_y_high) &
                          (nonzerox >= win_xleft_low) &  (nonzerox < win_xleft_high)).nonzero()[0]
        good_right_inds = ((nonzeroy >= win_y_low) & (nonzeroy < win_y_high) &
                           (nonzerox >= win_xright_low) &  (nonzerox < win_xright_high)).nonzero()[0]
        # Append these indices to the lists
        left_lane_inds.append(good_left_inds)
        right_lane_inds.append(good_right_inds)

        # If you found > minpix pixels, recenter next window on their mean position
        if len(good_left_inds) > minpix:
            leftx_current = np.int(np.mean(nonzerox[good_left_inds]))
        if len(good_right_inds) > minpix:
            rightx_current = np.int(np.mean(nonzerox[good_right_inds]))

    # Concatenate the arrays of indices
    left_lane_inds = np.concatenate(left_lane_inds)
    right_lane_inds = np.concatenate(right_lane_inds)

    # Extract left and right line pixel positions
    leftx = nonzerox[left_lane_inds]
    lefty = nonzeroy[left_lane_inds]
    rightx = nonzerox[right_lane_inds]
    righty = nonzeroy[right_lane_inds]

    # Fit a second order polynomial to each
    left_fit = np.polyfit(lefty, leftx, 2)
    right_fit = np.polyfit(righty, rightx, 2)

    ploty = np.linspace(0, img.shape[0] - 1, img.shape[0])
    left_fitx = left_fit[0]*ploty**2 + left_fit[1]*ploty + left_fit[2]
    right_fitx = right_fit[0]*ploty**2 + right_fit[1]*ploty + right_fit[2]

    if return_img:
        out_img[nonzeroy[left_lane_inds], nonzerox[left_lane_inds]] = [255, 0, 0]
        out_img[nonzeroy[right_lane_inds], nonzerox[right_lane_inds]] = [0, 0, 255]

        # Draw left and right lines
        for index in range(img.shape[0]):
            cv2.circle(out_img, (int(left_fitx[index]), int(ploty[index])), 3, (255,255,0))
            cv2.circle(out_img, (int(right_fitx[index]), int(ploty[index])), 3, (255,255,0))

        return (left_fit, right_fit), (left_fitx, ploty), (right_fitx, ploty), out_img.astype(int)

    return (left_fit, right_fit), (left_fitx, ploty), (right_fitx, ploty)

def detect_similar_lines(img, line_fits=None, return_img=False):
    if line_fits is None:
        return detect_lines(img, return_img)

    left_fit = line_fits[0]
    right_fit = line_fits[1]

    nonzero = img.nonzero()
    nonzeroy = np.array(nonzero[0])
    nonzerox = np.array(nonzero[1])
    margin = 100
    left_lane_inds = ((nonzerox > (left_fit[0]*(nonzeroy**2) + left_fit[1]*nonzeroy +
    left_fit[2] - margin)) & (nonzerox < (left_fit[0]*(nonzeroy**2) +
    left_fit[1]*nonzeroy + left_fit[2] + margin)))

    right_lane_inds = ((nonzerox > (right_fit[0]*(nonzeroy**2) + right_fit[1]*nonzeroy +
    right_fit[2] - margin)) & (nonzerox < (right_fit[0]*(nonzeroy**2) +
    right_fit[1]*nonzeroy + right_fit[2] + margin)))

    # Again, extract left and right line pixel positions
    leftx = nonzerox[left_lane_inds]
    lefty = nonzeroy[left_lane_inds]
    rightx = nonzerox[right_lane_inds]
    righty = nonzeroy[right_lane_inds]

    # If any of the lines could not be found,
    # perform a more exhaustive search
    if (leftx.size == 0 or rightx.size == 0):
        return detect_lines(img, return_img)

    # Fit a second order polynomial to each
    left_fit = np.polyfit(lefty, leftx, 2)
    right_fit = np.polyfit(righty, rightx, 2)

    # Generate x and y values for plotting
    ploty = np.linspace(0, img.shape[0] - 1, img.shape[0] )
    left_fitx = left_fit[0]*ploty**2 + left_fit[1]*ploty + left_fit[2]
    right_fitx = right_fit[0]*ploty**2 + right_fit[1]*ploty + right_fit[2]

    if return_img:
        # Create an image to draw on and an image to show the selection window
        out_img = np.dstack((img, img, img))*255
        window_img = np.zeros_like(out_img)

        # Generate a polygon to illustrate the search window area
        # And recast the x and y points into usable format for cv2.fillPoly()
        left_line_window1 = np.array([np.transpose(np.vstack([left_fitx - margin, ploty]))])
        left_line_window2 = np.array([np.flipud(np.transpose(np.vstack([left_fitx + margin, ploty])))])
        left_line_pts = np.hstack((left_line_window1, left_line_window2))
        right_line_window1 = np.array([np.transpose(np.vstack([right_fitx - margin, ploty]))])
        right_line_window2 = np.array([np.flipud(np.transpose(np.vstack([right_fitx + margin, ploty])))])
        right_line_pts = np.hstack((right_line_window1, right_line_window2))

        # Draw the lane onto the warped blank image
        cv2.fillPoly(window_img, np.int_([left_line_pts]), (0,255, 0))
        cv2.fillPoly(window_img, np.int_([right_line_pts]), (0,255, 0))
        out_img = cv2.addWeighted(out_img, 1, window_img, 0.3, 0)

        # Color in left and right line pixels
        out_img[nonzeroy[left_lane_inds], nonzerox[left_lane_inds]] = [255, 0, 0]
        out_img[nonzeroy[right_lane_inds], nonzerox[right_lane_inds]] = [0, 0, 255]

        for index in range(img.shape[0]):
            cv2.circle(out_img, (int(left_fitx[index]), int(ploty[index])), 3, (255,255,0))
            cv2.circle(out_img, (int(right_fitx[index]), int(ploty[index])), 3, (255,255,0))

        return (left_fit, right_fit), (left_fitx, ploty), (right_fitx, ploty), out_img.astype(int)

    return (left_fit, right_fit), (left_fitx, ploty), (right_fitx, ploty)

# Define perspective transform function
def warp(img, src_coordinates=None, dst_coordinates=None):
    # Define calibration box in source (original) and destination (desired or warped) coordinates
    img_size = (img.shape[1], img.shape[0])

    if src_coordinates is None:
        src_coordinates = np.float32(
            [[280,  700],  # Bottom left
             [595,  460],  # Top left
             [725,  460],  # Top right
             [1125, 700]]) # Bottom right

    if dst_coordinates is None:
        dst_coordinates = np.float32(
            [[250,  720],  # Bottom left
             [250,    0],  # Top left
             [1065,   0],  # Top right
             [1065, 720]]) # Bottom right

    # Compute the perspective transfor, M
    M = cv2.getPerspectiveTransform(src_coordinates, dst_coordinates)

    # Compute the inverse perspective transfor also by swapping the input parameters
    Minv = cv2.getPerspectiveTransform(dst_coordinates, src_coordinates)

    # Create warped image - uses linear interpolation
    warped = cv2.warpPerspective(img, M, img_size, flags=cv2.INTER_LINEAR)

    return warped, M, Minv

def combine_threshs(grad_x, grad_y, mag_binary, dir_binary, col_binary, ksize=15):
    # Combine the previous thresholds
    combined = np.zeros_like(dir_binary)
    combined[((grad_x == 1) & (grad_y == 1)) | ((mag_binary == 1) & (dir_binary == 1)) | (col_binary == 1)] = 1

    return combined

# Define a function that applies Sobel x and y,
# then computes the direction of the gradient
# and applies a threshold.
def dir_thresh(img, sobel_kernel=9, thresh=(0, np.pi/2)):
    # Grayscale
    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)

    # Calculate the x and y gradients
    sobelx = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=sobel_kernel)
    sobely = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=sobel_kernel)

    # Take the absolute value of the gradient direction,
    # apply a threshold, and create a binary image result
    absgraddir = np.arctan2(np.absolute(sobely), np.absolute(sobelx))
    binary_output =  np.zeros_like(absgraddir)
    binary_output[(absgraddir >= thresh[0]) & (absgraddir <= thresh[1])] = 1

    # Return the binary image
    return binary_output

# Define a function that applies Sobel x and y,
# then computes the magnitude of the gradient
# and applies a threshold
def mag_thres(img, sobel_kernel=3, thresh=(0, 255)):
    # Convert to grayscale

    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    # Take both Sobel x and y gradients
    sobelx = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=sobel_kernel)
    sobely = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=sobel_kernel)
    # Calculate the gradient magnitude
    gradmag = np.sqrt(sobelx**2 + sobely**2)
    # Rescale to 8 bit
    scale_factor = np.max(gradmag)/255
    gradmag = (gradmag/scale_factor).astype(np.uint8)
    # Create a binary image of ones where threshold is met, zeros otherwise
    binary_output = np.zeros_like(gradmag)
    binary_output[(gradmag >= thresh[0]) & (gradmag <= thresh[1])] = 1

    # Return the binary image
    return binary_output

def col_thresh(img, thresh=(170, 255)):
    # HLS
    hls = cv2.cvtColor(img, cv2.COLOR_RGB2HLS)

    # Select S channel because it is usually the best performant
    # for this task. R channel also performs similarly.
    s_channel = hls[:,:,2]

    binary_output = np.zeros_like(s_channel)
    binary_output[(s_channel >= thresh[0]) & (s_channel <= thresh[1])] = 1

    # Return the binary image
    return binary_output
    # img_hls = cv2.cvtColor(img, cv2.COLOR_RGB2HLS)
    # img_hls = cv2.medianBlur(img_hls, 5)
    # img_lab = cv2.cvtColor(img, cv2.COLOR_RGB2LAB)
    # img_lab = cv2.medianBlur(img_lab, 5)

    # big_kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (31, 31))
    # small_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (7, 7))

    # greenery = (img_lab[:, :, 2].astype(np.uint8) > 130) & cv2.inRange(img_hls, (0, 0, 50), (35, 190, 255))

    # road_mask = np.logical_not(greenery).astype(np.uint8) & (img_hls[:, :, 1] < 250)
    # road_mask = cv2.morphologyEx(road_mask, cv2.MORPH_OPEN, small_kernel)
    # road_mask = cv2.dilate(road_mask, big_kernel)

    # img2, contours, hierarchy = cv2.findContours(road_mask, cv2.RETR_LIST, cv2.CHAIN_APPROX_NONE)

    # biggest_area = 0
    # for contour in contours:
    #     area = cv2.contourArea(contour)
    #     if area>biggest_area:
    #         biggest_area = area
    #         biggest_contour = contour
    # road_mask = np.zeros_like(road_mask)
    # cv2.fillPoly(road_mask, [biggest_contour],  1)
    # return road_mask

# Define a function that applies Sobel x or y,
# then takes an absolute value and applies a threshold.
# Note: calling your function with orient='x', thresh_min=5, thresh_max=100
# should produce output like the example image shown above this quiz.
def abs_sobel_thresh(img, orient='x', sobel_kernel=3, thresh=(0, 255)):
    # Convert to grayscale
    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)

    # Apply x or y gradient with the OpenCV Sobel() function
    # and take the absolute value
    if orient == 'x':
        abs_sobel = np.absolute(cv2.Sobel(gray, cv2.CV_64F, 1, 0, sobel_kernel))
    if orient == 'y':
        abs_sobel = np.absolute(cv2.Sobel(gray, cv2.CV_64F, 0, 1, sobel_kernel))

    # Rescale back to 8 bit integer
    scaled_sobel = np.uint8(255*abs_sobel/np.max(abs_sobel))

    # Create a copy and apply the threshold
    binary_output = np.zeros_like(scaled_sobel)

    # Here I'm using inclusive (>=, <=) thresholds, but exclusive is ok too
    binary_output[(scaled_sobel >= thresh[0]) & (scaled_sobel <= thresh[1])] = 1

    # Return the result
    return binary_output

class ProcessImage:

    def __init__(self, images):
        # Make a list of calibration images
        images = glob.glob(images)

        # Calibrate camera
        self.ret, self.mtx, self.dist, self.rvecs, self.tvecs = calibrate_camera(images)
        self.lines_fit = None

    def __call__(self, img):
        # Undistord image
        img = cv2.undistort(img, self.mtx, self.dist, None, self.mtx)

        # Calculate directional gradient
        grad_binary = abs_sobel_thresh(img, orient='x', sobel_kernel=15, thresh=(30, 100))

        # Calculate gradient magnitude
        mag_binary = mag_thres(img, sobel_kernel=15, thresh=(50, 100))

        # Calculate gradient direction
        dir_binary = dir_thresh(img, sobel_kernel=15, thresh=(0.7, 1.3))

        # Calculate color threshold
        col_binary = col_thresh(img, thresh=(170, 255))

        grad_x = abs_sobel_thresh(img, orient='x', sobel_kernel=15, thresh=(30, 100))
        grad_y = abs_sobel_thresh(img, orient='y', sobel_kernel=15, thresh=(30, 100))

        # Combine all the thresholds to identify the lane lines
        combined = combine_threshs(grad_x, grad_y, mag_binary, dir_binary, col_binary, ksize=15)

        # Apply a perspective transform to rectify binary image ("birds-eye view")
        src_coordinates = np.float32(
            [[585, 460],  # Bottom left
             [203, 720],  # Top left
             [1127, 720],  # Top right
             [695, 460]]) # Bottom right

        dst_coordinates = np.float32(
            [[320, 0],  # Bottom left
             [320, 720],  # Top left
             [960, 720],  # Top right
             [960, 0]]) # Bottom right

        combined_warped, _, Minv = warp(combined, src_coordinates, dst_coordinates)

        self.lines_fit, left_points, right_points, out_img = detect_similar_lines(combined_warped, self.lines_fit, return_img=True)

        # Warp the detected lane boundaries back onto the original image.
        img_lane = draw_lane(img, combined_warped, left_points, right_points, Minv)

        # # Add metrics to the output img
        # out_img = add_metrics(img_lane, leftx=left_points[0], rightx=right_points[0])

        return img_lane

input_video = './challenge_video.mp4'
output_video = './challenge_video_solution.mp4'

## You may uncomment the following line for a subclip of the first 5 seconds
#clip1 = VideoFileClip(input_video).subclip(0,5)
clip1 = VideoFileClip(input_video)

# Process video frames with our 'process_image' function
process_image = ProcessImage('./camera_cal/calibration*.jpg')

white_clip = clip1.fl_image(process_image)
white_clip.write_videofile(os.path.join('output_videos', output_video), audio=False)


#!/usr/bin/env python
# coding: utf-8

# In[68]:

import numpy as np
import scipy
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
import pickle

from pathlib import Path
from pprint import pprint
from scipy.io import loadmat
from tqdm import trange, tqdm
from scipy.stats import ttest_ind
from collections import defaultdict as ddict
from easydict import EasyDict as edict
from functools import partial

from sklearn.linear_model import LogisticRegression
from sklearn.metrics import roc_auc_score, roc_curve
from sklearn.ensemble import RandomForestClassifier
from sklearn.tree import DecisionTreeClassifier
from sklearn.cluster import KMeans, SpectralClustering

# import holoviews as hv
# from holoviews import opts, dim

get_ipython().run_line_magic('matplotlib', 'inline')

# In[2]:

# import my own functions
import sys
sys.path.append("../../mypkg")
import importlib

# paras
import paras
importlib.reload(paras);
from paras import paras

# some useful constants
import constants
importlib.reload(constants)
from constants import REGION_NAMES, REGION_NAMES_WLOBE, RES_ROOT, DATA_ROOT

# In[18]:

# utils
import utils.matrix
importlib.reload(utils.matrix)
from utils.matrix import eig_sorted

import utils.misc
importlib.reload(utils.misc)
from utils.misc import paras2name, cumsum_cutoff, save_pkl_dict2folder, load_pkl_folder2dict

import utils.standardize
importlib.reload(utils.standardize)
from utils.standardize import minmax, minmax_mat, minmax_pn, minmax_plotting

import utils.brain_plot
importlib.reload(utils.brain_plot)
from utils.brain_plot import reorder_U, U_2brain_vec

# In[ ]:

# ## Some fns

# In[4]:

def _get_sig_val_trans_curve(ecpts, sig_val):
    ecpts_full = np.concatenate([[0], ecpts, [7200]])
    sig_val_trans_vec = np.zeros(7200)
    for ix in range(len(ecpts_full)-1):
        low_idx, up_idx = ecpts_full[ix:(ix+2)].astype(int)
        sig_val_trans_vec[low_idx:up_idx] = sig_val[ix]
    return sig_val_trans_vec

# In[5]:

def _cpts_2dwells(cpts, bds=[0, 7200], fs=120):
    cpts_full = np.sort(np.concatenate([bds, cpts]))
    return np.diff(cpts_full)/fs

# In[6]:

# update the count transited from one state to another.
def _update_trans_ct(cur_trans_ct, vec):
    for ix in range(len(vec)-1):
        from_ = vec[ix]
        to_ = vec[ix+1]
        cur_trans_ct[from_, to_] += 1
    return cur_trans_ct

# In[7]:

# plot the corrmat with 7 canonical nets
def _corr_plot(vecs, cutoff=0.05, trun_fn=lambda x:x, trans_fn=np.abs):
    vecs = np.array(vecs)
    if vecs.shape[0] != 68:
        vecs = vecs.T
    assert vecs.shape[0] == 68
    corrMat = np.zeros((7, vecs.shape[-1]))
    for ix in range(vecs.shape[-1]):
        curU = vecs[:, ix]
        curU = trun_fn(curU)
        for iy, kz in enumerate(_paras.canon_net_names):
            curV = _paras.canon_nets[kz]
            curV = trun_fn(curV)
            corr_v, pv = scipy.stats.pearsonr(curU, curV)
            if pv <= cutoff:
                corrMat[iy, ix] = corr_v
            else:
                corrMat[iy, ix] = 0

    plt.figure(figsize=[15, 5])
    trans_corrMat = trans_fn(corrMat)
    sns.heatmap(trans_corrMat,  yticklabels=_paras.canon_net_names,
                cmap="coolwarm", center=0,
                vmin=-1, vmax=1, annot=np.round(trans_corrMat, 2))
    return corrMat

# In[8]:

# truncate small value in vec
def _cumsum_trunc(vec, cutoff=0.9):
    vec = vec.copy()
    idxs = cumsum_cutoff(vec, cutoff)
    counter_idxs = np.delete(np.arange(len(vec)), idxs)
    vec[counter_idxs] = 0
    return vec

# In[85]:

def _minmax_wo0(vec):
    """minmax vec without consider the 0s in the vec
    """
    vec_trans = np.zeros_like(vec)
    vec_trans[vec!=0] = minmax_plotting(vec[vec!=0])
    return vec_trans

# save the files for brainplotting
def _save_brainplotting_vec(vecs, filname_prefix, is_csv=False, trans_fn=_minmax_wo0):
    vecs = np.array(vecs)
    if vecs.shape[0] != 68:
        vecs = vecs.T
    for ix in range(vecs.shape[1]):
        cur_sing_vec = vecs[:, ix]
        out_vec_raw = trans_fn(cur_sing_vec)
        out_vec = U_2brain_vec(reorder_U(out_vec_raw))
        np.savetxt(_paras.save_dir/f"./{filname_prefix}_{ix+1}.txt", out_vec)
        if is_csv:
            pd.DataFrame({"Name": REGION_NAMES,
                      "Vec": cur_sing_vec}).to_csv(_paras.save_dir/f"./{filname_prefix}_{ix+1}.csv",
                                                          index=False)

# In[ ]:

# ## Parameters

# In[10]:

pprint(paras)

# In[11]:

# this parameters only for this file
_paras = edict()
_paras.folder_name = "method3"
_paras.save_dir = RES_ROOT/_paras.folder_name
print(f"Save to {_paras.save_dir}")

# In[12]:

# Lets save some results
post_paras = edict()

# In[ ]:

# ##  Load data

# In[13]:

cur_res = load_pkl_folder2dict(_paras.save_dir)

# In[72]:

defNetsFil = list(DATA_ROOT.glob("DK_dic68.csv"))[0]
defNets = pd.read_csv(defNetsFil).T

DefNets_dict = {}
for ix in range(defNets.shape[-1]):
    curCol = defNets[ix]
    DefNets_dict[curCol[0]] = np.array(curCol[1:])
net_names = sorted(DefNets_dict.keys())

_paras.canon_nets = DefNets_dict
_paras.canon_net_names = net_names

# In[ ]:

# ## singular value (weights) analysis (Not good)

# In[13]:

# note that I take abs of B3 for both grps
abs_B3_AD = np.abs(CPDresult_AD[-1])
abs_B3_ctrl = np.abs(CPDresult_ctrl[-1])
abs_B3_both = np.concatenate([abs_B3_AD, abs_B3_ctrl], axis=0)

# In[14]:

# fit PCAs for AD and ctrl grps
pca_fit_AD = PCA(n_components=2)
pca_fit_AD.fit(abs_B3_AD)
pca_fit_ctrl = PCA(n_components=2)
pca_fit_ctrl.fit(abs_B3_ctrl)

pca_fit_both = PCA(n_components=2)
pca_fit_both.fit(abs_B3_both)

# In[15]:

# reduce the dim
singular_vals_ADs_red = [pca_fit_both.transform(np.abs(singular_vals_ADs[ix].T)) for ix in range(88)]
singular_vals_ctrls_red = [pca_fit_both.transform(np.abs(singular_vals_ctrls[ix].T)) for ix in range(88)]

# In[16]:

sig_val_curves_AD = []
for ix in range(88):
    cur_curve = get_sig_val_trans_curve(cpts_ADs_selected[ix],
                                        np.median(np.abs(singular_vals_ADs[ix]), axis=0))
    #cur_curve = get_sig_val_trans_curve(cpts_ADs_selected[ix], singular_vals_ADs_red[ix][:, 0])
    sig_val_curves_AD.append(cur_curve)

sig_val_curves_ctrl = []
for ix in range(88):
    cur_curve = get_sig_val_trans_curve(cpts_ctrls_selected[ix],
                                        np.median(np.abs(singular_vals_ctrls[ix]), axis=0))
    #cur_curve = get_sig_val_trans_curve(cpts_ctrls_selected[ix], singular_vals_ctrls_red[ix][:, 0])
    sig_val_curves_ctrl.append(cur_curve)

sig_val_curves_AD = np.array(sig_val_curves_AD)
sig_val_curves_ctrl = np.array(sig_val_curves_ctrl)

# In[17]:

center_vec = lambda x: x-x.mean()
plt_AD_vec = center_vec(np.median(sig_val_curves_AD, axis=0))
plt_ctrl_vec = center_vec(np.median(sig_val_curves_ctrl, axis=0))

# In[18]:

plt.figure(figsize=[10, 5], dpi=100)
times = np.linspace(0, 60, 7200)
plt.plot(times, plt_AD_vec, "-", label="AD")
plt.plot(times, plt_ctrl_vec, "-", label="Ctrl")
plt.xlabel("Time")
plt.legend()

# ## U analysis

# 20 is too many to have good results

# ### Save U for brain plot

# In[39]:

# 20 is too many to have good results
_save_brainplotting_vec(cur_res.CPDresult_ctrl[0], "B1", False)
_save_brainplotting_vec(np.abs(cur_res.CPDresult_ctrl[0]), "absB1")

# In[ ]:

# ## wU analysis

# In[58]:

# put all the singular values together
# always abs singluar vals
singular_vals_ADs_pool = np.abs(np.concatenate(cur_res.singular_vals_ADs, axis=1))
singular_vals_ctrls_pool = np.abs(np.concatenate(cur_res.singular_vals_ctrls, axis=1))
singular_vals_both_pool = np.abs(np.concatenate([singular_vals_ADs_pool, singular_vals_ctrls_pool], axis=1))

# In[59]:

# if not minmax, the clustering is poor
singular_vals_both_pool_mm = minmax_mat(singular_vals_both_pool, is_row=False)
singular_vals_ctrls_pool_mm = minmax_mat(singular_vals_ctrls_pool, is_row=False)
singular_vals_ADs_pool_mm = minmax_mat(singular_vals_ADs_pool, is_row=False)

# In[200]:

singular_vals_both_pool_mm = singular_vals_both_pool/np.linalg.norm(singular_vals_both_pool,axis=0)
singular_vals_ctrls_pool_mm = singular_vals_ctrls_pool/np.linalg.norm(singular_vals_ctrls_pool,axis=0)
singular_vals_ADs_pool_mm = singular_vals_ADs_pool/np.linalg.norm(singular_vals_ADs_pool,axis=0)

# ###  Clustering singular values (Only Ctrl)

# In[107]:

# find the best num of clusters
def kmeans_BIC(kmCluObj, Cn=1):
    nclu = kmCluObj.n_clusters
    nsps = len(kmCluObj.labels_)
    nDim = kmCluObj.cluster_centers_.shape[1]
    log_lik = np.log(kmCluObj.inertia_/nsps/nDim)
    BIC = log_lik + Cn*np.log(np.log(nsps*nDim))*np.log(nsps)*(nclu-1)/nsps
    #BIC = nsps*log_lik + np.log(nsps)*nclu*nDim/2
    #BIC = nsps*log_lik/(nsps-nDim*nclu)
    #BIC = nsps*log_lik + nclu*nDim
    return BIC

# In[108]:

# tuning num of clus (not used)
n_cluss = np.array(list(range(2,21)))
BICs =[]
stds = []
for ix in n_cluss:
    cur_clus = KMeans(n_clusters=ix, random_state=0).fit(singular_vals_ctrls_pool_mm.T)
    BICs.append(kmeans_BIC(cur_clus, 2))
    stds.append(np.std(np.unique(cur_clus.labels_, return_counts=True)[1]))
plt.plot(n_cluss, BICs)

# In[319]:

n_clus = 7
#clus = SpectralClustering(n_clusters=n_clus, random_state=0).fit(singular_vals_ctrls_pool_mm.T)
clus = KMeans(n_clusters=n_clus, random_state=0).fit(singular_vals_ctrls_pool_mm.T)
ctrl_labs = clus.labels_
np.unique(clus.labels_, return_counts=True)

# In[320]:

# clustering AD given centers from ctrl
d_mat = []
for ix in range(n_clus):
    cur_cens = clus.cluster_centers_[ix, :]
    d_mat.append(np.sum((singular_vals_ADs_pool_mm.T- cur_cens)**2,axis=1))
d_mat = np.array(d_mat)
AD_labs = d_mat.argmin(axis=0)
np.unique(AD_labs, return_counts=True)

# In[321]:

# centers
clu_ws = clus.cluster_centers_
#clu_ws = np.array([_cumsum_trunc(cen, 0.7) for cen in clus.cluster_centers_])
# wegithed U, U x w
cens_wU = np.matmul(np.abs(cur_res.CPDresult_ctrl[0]), clu_ws.T)

# In[338]:

plt.figure(figsize=[6, 5])
mat = np.corrcoef(np.array(cur_res.CPDresult_ctrl[0].T))
mat[np.abs(mat)<0.3] = 0
sns.heatmap(mat, cmap="bwr", center=0)

# In[323]:

plt.figure(figsize=[15, 10])
plt.subplot(131)
Us = cur_res.CPDresult_ctrl[0]
sns.heatmap(np.abs(Us))
plt.subplot(132)
sns.heatmap(clu_ws.T)
plt.subplot(133)
sns.heatmap(np.abs(cens_wU))

# In[117]:

tfn1 = partial(_cumsum_trunc, cutoff=0.8)
_corr_plot(cens_wU, trun_fn=tfn1);

# In[118]:

tfn1 = partial(_cumsum_trunc, cutoff=0.8)
trans_fn = lambda x: _minmax_wo0(tfn1(x))
_save_brainplotting_vec(cens_wU, f"clu_ws_{n_clus}", False, trans_fn)
_save_brainplotting_vec(np.abs(cens_wU), f"clu_abs_wU_{n_clus}", False, trans_fn)

# ### clustering both grps together

# In[202]:

# tuning num of clus (not used)
n_cluss = np.array(list(range(2,21)))
BICs =[]
stds = []
for ix in n_cluss:
    cur_clus = KMeans(n_clusters=ix, random_state=0).fit(singular_vals_both_pool_mm.T)
    BICs.append(kmeans_BIC(cur_clus, 2))
    stds.append(np.std(np.unique(cur_clus.labels_, return_counts=True)[1]))
plt.plot(n_cluss, BICs)

# In[315]:

n_clus = 10
clus = SpectralClustering(n_clusters=n_clus, random_state=0).fit(singular_vals_both_pool_mm.T)
#clus = KMeans(n_clusters=n_clus, random_state=0).fit(singular_vals_both_pool_mm.T)
both_labs = clus.labels_
np.unique(clus.labels_, return_counts=True)

# In[316]:

clu_Ys = np.array(["AD"]*len(singular_vals_ADs_pool.T) + ["ctrl"]*len(singular_vals_ctrls_pool.T))
ct_table = pd.crosstab(both_labs, clu_Ys)
ct_table

# In[317]:

AD_ratios = np.array(ct_table["AD"])/singular_vals_ADs_pool.shape[-1]
ctrl_ratios = np.array(ct_table["ctrl"])/singular_vals_ctrls_pool.shape[-1]
plt.bar(np.arange(0, n_clus), AD_ratios, alpha=0.5, label="AD")
plt.bar(np.arange(0, n_clus), ctrl_ratios, alpha=0.5, label="Ctrl")
plt.legend()

# In[282]:

AD_clu_idxs = [5, 8, 9]
ctrl_clu_idxs = [0, 6, 12]

AD_ws = []
for clu_idx in AD_clu_idxs:
    kpidx = np.bitwise_and(clu_Ys=="AD", both_labs==clu_idx)
    #kpidx = both_labs==clu_idx
    cur_w = singular_vals_both_pool_mm[:, kpidx].mean(axis=1)
    #cur_w = _cumsum_trunc(cur_w, 0.8)
    AD_ws.append(cur_w)

AD_ws = np.array(AD_ws)
AD_wUs = np.matmul(cur_res.CPDresult_ctrl[0], AD_ws.T)

ctrl_ws = []
for clu_idx in ctrl_clu_idxs:
    kpidx = np.bitwise_and(clu_Ys=="ctrl", both_labs==clu_idx)
    #kpidx = both_labs==clu_idx
    cur_w = singular_vals_both_pool_mm[:, kpidx].mean(axis=1)
    #cur_w = _cumsum_trunc(cur_w, 0.8)
    ctrl_ws.append(cur_w)

ctrl_ws = np.array(ctrl_ws)
ctrl_wUs = np.matmul(cur_res.CPDresult_ctrl[0], ctrl_ws.T)

# In[284]:

plt.figure(figsize=[20, 10])

plt.subplot(141)
sns.heatmap(ctrl_ws.T)
plt.subplot(142)
sns.heatmap(AD_ws.T)

plt.subplot(143)
sns.heatmap(np.abs(ctrl_wUs))
plt.subplot(144)
sns.heatmap(np.abs(AD_wUs))

# In[286]:

AD_wUs[56, :] = 0
AD_wUs[58, :] = 0
tfn1 = partial(_cumsum_trunc, cutoff=0.8)
trans_fn = lambda x: _minmax_wo0(tfn1(x))
_save_brainplotting_vec(AD_wUs, f"AD_wU", False, trans_fn)
_save_brainplotting_vec(np.abs(AD_wUs), f"AD_abs_wU", False, trans_fn)

# In[287]:

ctrl_wUs[56, :] = 0
ctrl_wUs[58, :] = 0
tfn1 = partial(_cumsum_trunc, cutoff=0.8)
trans_fn = lambda x: _minmax_wo0(tfn1(x))
_save_brainplotting_vec(ctrl_wUs, f"ctrl_wU", False, trans_fn)
_save_brainplotting_vec(np.abs(ctrl_wUs), f"ctrl_abs_wU", False, trans_fn)

# In[ ]:

# ### Transition

# In[50]:

flag = 0
AD_clu_labs = []
for cur_v in singular_vals_ADs:
    n_ele = cur_v.shape[1]
    AD_clu_labs.append(AD_labs[flag:(flag+n_ele)])
    flag += n_ele

flag = 0
ctrl_clu_labs = []
for cur_v in singular_vals_ctrls:
    n_ele = cur_v.shape[1]
    ctrl_clu_labs.append(ctrl_labs[flag:(flag+n_ele)])
    flag += n_ele

# In[51]:

trans_ct = edict()

cur_trans_ct = np.zeros((n_clus, n_clus))
for ix in range(len(AD_clu_labs)):
    cur_vec = AD_clu_labs[ix]
    cur_trans_ct = update_trans_ct(cur_trans_ct, cur_vec)
trans_ct.AD = cur_trans_ct

cur_trans_ct = np.zeros((n_clus, n_clus))
for ix in range(len(ctrl_clu_labs)):
    cur_vec = ctrl_clu_labs[ix]
    cur_trans_ct = update_trans_ct(cur_trans_ct, cur_vec)
trans_ct.ctrl = cur_trans_ct

# In[52]:

non_diag_mask = ~np.eye(n_clus,dtype=bool)

# In[53]:

def mat_2df(mat, include_diag=False, no_zero=True, cut_off=None):
    source = np.repeat(list(range(1, 1+n_clus)), n_clus)
    target = np.repeat(np.arange(1,1+n_clus).reshape(1,-1), n_clus, axis=0).reshape(-1)
    value = mat.flatten()
    idxs = source !=-100
    if not include_diag:
        idxs = np.bitwise_and(idxs, source != target)
    if no_zero:
        idxs = np.bitwise_and(idxs, value!=0)
    if cut_off is not None:
        idxs = np.bitwise_and(idxs, value>cut_off)
    source = source[idxs]
    target = target[idxs]
    value = value[idxs]
    res_df = pd.DataFrame({
        "source": source,
        "target": target,
        "value": value.astype(int)
    })
    return res_df

# In[57]:

nodes = pd.DataFrame({
    "index": list(range(1, n_clus+1)),
    "name": [f"State {ix}" for ix in range(1, n_clus+1)]
})
nodes = hv.Dataset(nodes, "index")

cut_v = np.median(np.concatenate([trans_ct.AD[non_diag_mask], trans_ct.ctrl[non_diag_mask]]))

# In[58]:

plt_data_AD = mat_2df(trans_ct.AD, False, cut_off=cut_v)

chord = hv.Chord((plt_data_AD, nodes))
chord.opts(cmap='Category20', labels='name',
           title="AD", fontsize=20,
           zrotation=45,
           edge_color=dim('source').str(),
           node_color=dim("index").str())

# In[59]:

plt_data_ctrl = mat_2df(trans_ct.ctrl, False, cut_off=cut_v)

chord = hv.Chord((plt_data_ctrl, nodes))
chord.opts(cmap='Category20', labels='name',
           title="Ctrl", fontsize=20,
           edge_color=dim('source').str(),
           node_color=dim("index").str())

# ### Dwell x wU

# #### Each segment

# In[60]:

dwells_AD = [cpts_2dwells(cpts) for cpts in cpts_ADs_selected]
dwells_ctrl = [cpts_2dwells(cpts) for cpts in cpts_ctrls_selected]

# In[61]:

dwells_sigvec_AD = ddict(list)
for ix in range(len(singular_vals_ADs)):
    cur_dwells = dwells_AD[ix]
    cur_idxs = AD_clu_labs[ix] +1 # adjust such that starting from 1 (not 0)
    [dwells_sigvec_AD[cur_idxs[ix]].append(cur_dwells[ix]) for ix in range(len(cur_idxs))]

dwells_sigvec_ctrl = ddict(list)
for ix in range(len(singular_vals_ctrls)):
    cur_dwells = dwells_ctrl[ix]
    cur_idxs = ctrl_clu_labs[ix] +1 # adjust such that starting from 1 (not 0)
    [dwells_sigvec_ctrl[cur_idxs[ix]].append(cur_dwells[ix]) for ix in range(len(cur_idxs))]

# In[62]:

sum_dwells = edict()
sum_dwells.AD = list()
sum_dwells.ctrl = list()

for ix in range(1, n_clus+1):
    sum_dwells.AD.append(np.sum(dwells_sigvec_AD[ix]))
    sum_dwells.ctrl.append(np.sum(dwells_sigvec_ctrl[ix]))

# In[63]:

for ix in range(1, 9):
    print(scipy.stats.ttest_ind(dwells_sigvec_AD[ix], dwells_sigvec_ctrl[ix]))

# In[64]:

ratio_dwells = edict()
for type_ in ["AD", "ctrl"]:
    ratio_dwells[type_] = np.array(sum_dwells[type_])/np.sum(sum_dwells[type_])

# In[65]:

plt.figure(figsize=[10, 5])
plt.bar(list(range(1, n_clus+1)), sum_dwells["AD"], alpha=0.5, label="AD")
plt.bar(list(range(1, n_clus+1)), sum_dwells["ctrl"], alpha=0.5, label="Ctrl")
plt.title("Dwell time ratio for 7 brain states")
plt.ylabel("Mean dwell time (t)")
#plt.xticks(np.array(list(range(0, n_clus)))+0.5, states, rotation=45)
plt.legend()

# ####  each subject

# In[66]:

# dwell time of each states for each subject
dwells_sigvec_AD_sub = []
for ix in range(len(singular_vals_ADs)):
    cur_dwells = dwells_AD[ix]
    cur_idxs = AD_clu_labs[ix] # adjust such that starting from 0 (not 1)
    dwells_sigvec_AD_sub.append([np.sum(cur_dwells[cur_idxs==iy]) for iy in range(n_clus)])

dwells_sigvec_ctrl_sub = []
for ix in range(len(singular_vals_ctrls)):
    cur_dwells = dwells_ctrl[ix]
    cur_idxs = ctrl_clu_labs[ix] # adjust such that starting from 0 (not 1)
    dwells_sigvec_ctrl_sub.append([np.sum(cur_dwells[cur_idxs==iy]) for iy in range(n_clus)])

dwells_sigvec_AD_sub = np.array(dwells_sigvec_AD_sub)
dwells_sigvec_ctrl_sub = np.array(dwells_sigvec_ctrl_sub)

# In[67]:

dwells_subs = np.concatenate([dwells_sigvec_AD_sub, dwells_sigvec_ctrl_sub], axis=0)
dwells_subs_df = {f"State{ix+1}":dwells_subs[:, ix] for ix in range(n_clus)}
dwells_subs_df["group"] = ["AD"]*88 + ["ctrl"] * 88
dwells_subs_df = pd.DataFrame(dwells_subs_df)

# In[68]:

dwells_subs_df.groupby("group").mean()

# In[69]:

for ix in range(n_clus):
    test_res = scipy.stats.ranksums(dwells_sigvec_AD_sub[:, ix],
                                     dwells_sigvec_ctrl_sub[:, ix])
    print(test_res)

# In[70]:

ix = 2
plt.hist(dwells_sigvec_AD_sub[:, ix], label="AD", alpha=0.3)
plt.hist(dwells_sigvec_ctrl_sub[:, ix], label="ctrl", alpha=0.3)
plt.legend()

# In[ ]:


import pwd

import os
import threading
import time

from packageship.application.cli.base import BaseCommand
from packageship.application.common.exc import InitializeError, ResourceCompetitionError

class InitServiceThread(threading.Thread):
    """
    Description: Execute the initialization thread
    """

    def __init__(self, func, param, *args, **kwargs):
        super(InitServiceThread, self).__init__(*args, **kwargs)
        self._func = func
        self._args = param
        self.error = False

    def run(self):
        try:
            self._func(*self._args)
        except (InitializeError, ResourceCompetitionError) as error:
            self.error = True
            print("\r", error)

class InitDatabaseCommand(BaseCommand):
    """
    Description: Initialize database command
    Attributes:
        parse: Command line parsing example
        params: Command line parameters
    """

    def __init__(self):
        """
        Description: Class instance initialization
        """
        super(InitDatabaseCommand, self).__init__()
        self.parse = BaseCommand.subparsers.add_parser(
            "init", help="initialization of the database"
        )
        self.params = [
            ("-filepath", "str", "specify the path of conf.yaml", "", "store")
        ]
        self._char = ["/", "-", "\\"]

    def register(self):
        """
        Description: Command line parameter injection
        Args:

        Returns:

        Raises:

        """
        super(InitDatabaseCommand, self).register()
        self.parse.set_defaults(func=self.do_command)

    @property
    def login_user(self):
        """
        Description: The user logged in to the system
        """
        return pwd.getpwuid(os.getuid())[0]

    def do_command(self, params):
        """
        Description: Action to execute command
        Args:
            params: Command line parameters
        Returns:

        Raises:

        """

        if self.login_user not in ["root", "pkgshipuser"]:
            print("The current user does not have initial execution permission")
            return
        # Determine whether the service is started
        if not self.is_service_start():
            print("The pkgship service is not started,please start the service first")
            return

        from packageship.application.initialize.integration import InitializeService

        init = InitializeService()
        file_path = params.filepath
        if file_path:
            file_path = os.path.abspath(file_path)

        _init_service_thread = InitServiceThread(
            func=init.import_depend, param=(file_path,)
        )
        _init_service_thread.setDaemon(True)
        _init_service_thread.start()

        while _init_service_thread.is_alive():
            for number in range(3):
                print(
                    "\r",
                    "initializing{}".format("." * 10),
                    self._char[number],
                    end="",
                    flush=True,
                )
                time.sleep(0.5)
        print(os.linesep)
        if not _init_service_thread.error:
            self._display_error(init)

    def _display_error(self, initialize_service):
        if initialize_service.success:
            print("Database initialize success", end=os.linesep)
        else:
            print(
                "Initialize failed database: %s" % " ".join(initialize_service.fail),
                end=os.linesep,
            )
        if initialize_service.remove_error:
            print("Remove database failed:", end=os.linesep)
            for database, message in initialize_service.remove_error.items():
                print("%s: %s" % (database, message), end=os.linesep)
        if initialize_service.reindex_error:
            print(
                os.linesep,
                "The database name changes but the content is the same during",
                " initialization,an error occurred while renaming the index:",
                end=os.linesep,
            )
            _reindex, del_index = [], []
            for reindex_error in initialize_service.reindex_error:
                _reindex.append(reindex_error["reindex"])
                del_index.extend(reindex_error["del_index"])
            print("Reindex failed:", end=os.linesep)
            print(os.linesep.join(_reindex))
            print(
                os.linesep,
                "Please manually delete the following databases:",
                end=os.linesep,
            )
            print(os.linesep.join(del_index))


import cv2
cv2.namedWindow('frame', flags=cv2.WINDOW_AUTOSIZE)
cam = cv2.VideoCapture(1)
while True:
    ret, frame = cam.read()
    frame = cv2.resize(frame, (1280, 720), interpolation=cv2.INTERSECT_PARTIAL)
    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    for i in range(0,100):
        frame = cv2.putText(frame, '17.09.2023', (900+i, 700+i), cv2.FONT_ITALIC, 2, 0+i*2, 3)
    frame = cv2.putText(frame, '17.09.2023', (900, 700), cv2.FONT_ITALIC, 2, 255, 1)
    cv2.imshow('frame', frame)
    if cv2.waitKey(1) == ord('q'):
        break


import sqlite3
#https://www.fuzzwork.co.uk/dump/sqlite-latest.sqlite.bz2
conn = sqlite3.connect('sqlite-latest.sqlite')
c = conn.cursor()

npc_null_regions = ['Curse', 'Delve', 'Geminate', 'Great Wildlands', 'Outer Ring', 'Stain', 'Syndicate', 'Venal']

npc_null_region_ids = []

for region in npc_null_regions:
    c.execute('SELECT regionID FROM mapRegions WHERE regionName=?', (region,))
    npc_null_region_ids.append(str(c.fetchone()[0]))

def dfs_paths(graph, start, jumps):
    stack = [(start, [start])]
    while stack:
        (vertex, path) = stack.pop()
        for next in graph[vertex] - set(path):
            #if (goal and next == goal) or len(path + [next]) > (jumps - 1):
            if len(path + [next]) == jumps:
            #if next == goal or len(path+[next]) == 20:
                yield path + [next]
            else:
                stack.append((next, path + [next]))

def map_solarsystem_jumps(location, connected_systems, jumps):
    solar_system = str(solar_system_by_name(location)['solarSystemID'])
    hole_connections = []
    for system in connected_systems:
        hole_connections.append(str(solar_system_by_name(system)['solarSystemID']))

    c.execute('SELECT fromSolarSystemID, toSolarSystemID FROM mapSolarSystemJumps')
    jump_map = c.fetchall()
    graph = {}
    for jump in jump_map:
        try:
            graph[str(jump[0])].add(str(jump[1]))
        except:
            graph[str(jump[0])] = set([str(jump[1])])
    graph['31000735'] = set(hole_connections)
    for connection in hole_connections:
        try:
            graph[str(connection)].add('31000735')
        except:
            graph[str(connection)] = set (['31000735'])
    paths = list(dfs_paths(graph, solar_system, jumps))
    #print paths
    #get all involved solar systems from paths
    return paths

def get_involved_regions_from_paths(paths):
    involved_solar_systems = []
    for path in paths:
        for ss in path:
            if ss not in involved_solar_systems:
                involved_solar_systems.append(ss)
    #get all involved regions from solar systems
    regions = []
    for ss in involved_solar_systems:
        regionID = solar_system_by_id(ss)['regionID']
        if regionID not in regions:
            regions.append(regionID)
    return regions

def all_region_ids():
    c.execute('SELECT * FROM mapRegions')
    regions = c.fetchall()
    return regions

def item_properties(item):
    it = {}
    it['typeID'] = item[0]
    it['typeName'] = item[2]
    it['volume'] = item[5]
    return it

def item_by_id(itemid):
    id = (itemid,)
    c.execute('SELECT * FROM invTypes WHERE typeID=?', id)
    item = c.fetchone()
    return item_properties(item)

def solar_system_properties(solar_system):
    ss = {}
    ss['regionID'] = solar_system[0]
    ss['solarSystemID'] = solar_system[2]
    ss['solarSystemName'] = solar_system[3]
    ss['security'] = solar_system[21]
    return ss

def solar_system_by_name(name):
    name = (name,)
    c.execute('SELECT * FROM mapSolarSystems WHERE solarSystemName=?', name)
    solar_system = c.fetchone()
    return solar_system_properties(solar_system)

def solar_system_by_id(id):
    id = (id,)
    c.execute('SELECT * FROM mapSolarSystems WHERE solarSystemID=?', id)
    solar_system = c.fetchone()
    return solar_system_properties(solar_system)

def station_properties(station):
    if station == None:
        return None
    s = {}
    s['stationID'] = station[0]
    s['solarSystemID'] = station[8]
    s['regionID'] = station[10]
    s['stationName'] = station[11]
    return s

def station_by_id(id):
    if id > 70000000: #unknown location id with high number ie, 1023968078820
        return None
    id = (id,)
    c.execute('SELECT * FROM staStations WHERE stationID=?', id)
    station = c.fetchone()
    return station_properties(station)


#!/usr/bin/python3
import requests
import argparse

NASAAPI = "https://api.nasa.gov/planetary/apod?"

def main():
    print(f"Welcome, {args.name}! Let's get crackin' with your API!")
    nasacreds = "api_key=" + args.key # API KEY GOES HERE
    url= NASAAPI + nasacreds
    print(url)
    apodresp = requests.get(url)
    apod = apodresp.json()

    print("YOUR IMAGE URL:", apod["url"])

# BEFORE MAIN RUNS
parser= argparse.ArgumentParser(description="Collects an API key to access a NASA API.")
            # ^ custom class object
            # creating an object that scoops up
            # our little --args at the command line

parser.add_argument("--name", help="The user's name", default="Slappy")
# teach our lil' parser object the different arguments that our code will take!
parser.add_argument("--key", required=True, help="NASA API key string")
                             # ^ you MUST use this arg!

# have the parser combine all those arguments into one object
args= parser.parse_args()

main()


from django.urls import path
from rest_framework_simplejwt.views import TokenObtainPairView

from base.views.user_views import (
    RegisterUserView,
    GetUserView,
    UserProfileView,
    UpdateUserProfileView,
    UpdateUserView,
    DeleteUserView,
    GetUserByIdView,
    CustomTokenObtainPairView
)

urlpatterns = [
    path('register/', RegisterUserView.as_view(), name='register'),
    path('login/', CustomTokenObtainPairView.as_view(), name='token_obtain_pair'),
    path('users/', GetUserView.as_view(), name='get_users'),
    path('profile/<str:username>/', UserProfileView.as_view(), name='get_user_profile'),
    path('profile/update/', UpdateUserProfileView.as_view(), name='update_user_profile'),
    path('user/update/', UpdateUserView.as_view(), name='update_user'),
    path('user/delete/<int:pk>/', DeleteUserView.as_view(), name='delete_user'),
    path('user/<int:pk>/', GetUserByIdView.as_view(), name='get_user_by_id'),
]


'''
Calc camera instrinsics given n>2 captures including https://github.com/opencv/opencv/blob/4.x/doc/pattern.png
'''

import sys

import cv2
import numpy as np

CLAHE = cv2.createCLAHE(clipLimit = 5)

BOARD_SIZE = (9, 6)
TERMINATION_CRITERIA = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)

MAX_CAPTURES=200

def main():
    print('Instructions:')
    print('  1.) Print out calibration pattern and place it somewhere it won\'t move.')
    print('  2.) Point camera at the calibration pattern, move it around until we detect the pattern.')
    print('      If detected, a window will pop up with annotations.')
    print(f'  3.) Continue moving the camera around until you capture {MAX_CAPTURES} instances of the pattern.')
    print('  4.) ???')
    print('  5.) Profit!')

    if len(sys.argv) < 2:
        print(f'Usage: {sys.argv[0]} <webcam num>')
        return

    capture_device = cv2.VideoCapture(int(sys.argv[1]))
    if not capture_device:
        print(f'Failed to get capture device {sys.argv[1]}')
        return

    objp = np.zeros((BOARD_SIZE[1]*BOARD_SIZE[0], 3), np.float32)
    objp[:,:2] = np.mgrid[0:BOARD_SIZE[0], 0:BOARD_SIZE[1]].T.reshape(-1,2)

    imgpoints = []
    objpoints = []
    num_captures = 0

    shape = None

    while True:
        key = cv2.waitKey(1)
        if key == ord('q') or key == 27:  # 27 == esc
            print('Got exit key, breaking early...')
            break

        ret, frame = capture_device.read()

        shape = frame.shape
        height, width, _ = shape

        # Convert to black and white for easier processing (TODO: can we extract details from color)
        frame_bw = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

        # The declaration of CLAHE
        # clipLimit -> Threshold for contrast limiting
        frame_enhanced = CLAHE.apply(frame_bw)  # + 20

        ret, corners = cv2.findChessboardCorners(frame_enhanced, BOARD_SIZE,
                cv2.CALIB_CB_ADAPTIVE_THRESH + cv2.CALIB_CB_FAST_CHECK)
        if not ret:
            #print('FATAL: failed to find chessboard corners in image.')
            continue

        corners = cv2.cornerSubPix(frame_enhanced, corners, (11, 11), (-1, -1),
                TERMINATION_CRITERIA)

        imgpoints.append(corners)
        objpoints.append(objp)

        # Display the resulting frame
        image_debug = cv2.drawChessboardCorners(frame, BOARD_SIZE,
                corners, ret)

        num_captures += 1
        cv2.putText(image_debug, f'{num_captures}', (0, height - 5),
                cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 100, 255), 1, cv2.LINE_AA)

        cv2.imshow('frame', image_debug)

        if num_captures >= MAX_CAPTURES:
            print('Reached max captures, breaking...')
            break

    capture_device.release()

    if shape is None:
        print('No captures, can\'t calibrate camera...')
        return

    print('Calculating calibration stuff.')
    ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(objpoints, imgpoints,
            shape[:2], None, None)

    print(f'f_x = {float(mtx[0][0])}')
    print(f'f_y = {float(mtx[1][1])}')
    print(f'c_x = {float(mtx[0][2])}')
    print(f'c_y = {float(mtx[1][2])}')
    print(f'k_1 = {dist[0][0]}')
    print(f'k_2 = {dist[0][1]}')
    print(f'p_1 = {dist[0][2]}')
    print(f'p_2 = {dist[0][3]}')
    print(f'k_3 = {dist[0][4]}')

if __name__ == '__main__':
    main()
    cv2.destroyAllWindows()


# Basic modules
import numpy as np
import cv2
import time
from datetime import datetime
from threading import Thread

# Camera and detection modules
import pyrealsense2 as rs
from cvzone.FPS import FPS
from cvzone_fork import HandDetector

# Adjusted mathematical and realsense function
import helper_functions.functions_cv2 as fcv2
import helper_functions.functions_math as fcm

# Kalman Filter modules
from filterpy.kalman import KalmanFilter
from filterpy.common import Q_discrete_white_noise
from filterpy.kalman import (
    UnscentedKalmanFilter,
    MerweScaledSigmaPoints,
    JulierSigmaPoints,
)

# # # # # # # # # # # # # # # # # # # #
# Initialize basic counter constants  #
# # # # # # # # # # # # # # # # # # # #
# Time managment
pTime = 0
cTime = 0
num = 0

# Comparison of coordinates - Top view/Side view
i = 0
j = 0

# Dead pixels
c_dead_top = 0
c_dead_side = 0

# material distance
isInside = False
distance_material = 0

# # # # # # # # # #
# Hand detection  #
lm_range = range(21)
depth_distance_lst = []

# Landmark constants
dist_init_lst = []
dist_current_lst = []
lm_rangeT = range(20)
lm_rangeComparison = range(14)
lm_coord_lst = []
boxW_fix = None
boxH_fix = None
mean_depth_lms = 0
lms_shortened = True
baseline_shortened = True
occlusion_range = range(13)
lm_counter = 0

# # # # # # # # # # # # # #
# Unscented Kalman Filter #
# # # # # # # # # # # # # #

# For extensive documentation check the filter-py
# https://filterpy.readthedocs.io/en/latest/

# time-step
dt = 0.1
eps_max = 4.0
epss = []
counter = 0
Q_scale_factor = 100

# f-, h-functions

# # Initial state transition function
# def f_state_transition(x):
#     """state transition function for a constant velocity
#     with state vector [x, x', x'', y, ..]'"""

#     F = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]], dtype=float)
#     return np.dot(F, x)
#     # return F @ x

# Adjusted transitional function using expected movement
def f_state_transition(x, x_1):
    """state transition function for a constant velocity
    with state vector [x, x', x'', y, ..]'"""

    dt = 0.1
    P = np.array([[100.0, 0.0, 0.0], [0.0, 100.0, 0.0], [0.0, 0.0, 100.0]])

    w = np.array([np.random.normal(0.0, 0.04, 3)])
    A = np.exp(-2 * 0.5 * dt) * x_1
    B = 2 * np.exp(-0.5 * dt) * np.cos(2 * np.pi * 1 * dt) * x
    C = np.sqrt(P) * w

    return A + B + C

def h_observation(x):
    """state transition function for a constant velocity
    with state vector [x, x', x'', y, ..]'"""

    return np.array([[x[0]], [x[1]], [x[2]]])
    # return H @ x

# create sigma points
points_0 = MerweScaledSigmaPoints(n=3, alpha=0.1, beta=2.0, kappa=-1)
# points = JulierSigmaPoints(n=9, kappa=3.-3, sqrt_method=np.linalg.cholesky)
# points_0 = JulierSigmaPoints(n=3, kappa=3.-3)

# Create Unscented_KF class
Unscented_KF = UnscentedKalmanFilter(
    dim_x=3, dim_z=3, dt=dt, fx=f_state_transition, hx=h_observation, points=points_0
)

# x: system state state - initial position
Unscented_KF.x = np.array([[0.0], [0.0], [0.0]])  # x position  # y  # z
Unscented_KF.x_1 = np.array([[0.0], [0.0], [0.0]])

# P: estimate uncertainty matrix
Unscented_KF.P = np.array([[100.0, 0.0, 0.0], [0.0, 100.0, 0.0], [0.0, 0.0, 100.0]])

# R: measurement covariance matrix
Unscented_KF.R = 5

# Q: process noise matrix
Unscented_KF.Q = Q_discrete_white_noise(2, dt, 0.1)

try:
    # Create pipeline
    pipeline = rs.pipeline()
    pipeline2 = rs.pipeline()

    # Create a config object
    config = rs.config()
    config2 = rs.config()

    # Configure the pipeline to stream the depth stream
    # Change this parameters according to the recorded bag file resolution
    config2.enable_stream(rs.stream.depth, 1280, 720, rs.format.z16, 30)
    config2.enable_stream(rs.stream.color, 1280, 720, rs.format.rgb8, 30)

    # start pipeline
    pipeline2.start(config2)

    # Configure the pipeline to stream the depth stream
    config.enable_stream(rs.stream.depth, 1280, 720, rs.format.z16, 30)
    config.enable_stream(rs.stream.color, 1280, 720, rs.format.rgb8, 30)

    # start pipeline
    cv2.waitKey(300)
    pipeline.start(config)

    # Create colorizer object
    colorizer = rs.colorizer()

    fpsReader = FPS()
    detector = HandDetector(detectionCon=0.8, maxHands=2)

    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # # # # # # # # # # # #
    # Camera calibration  #
    # # # # # # # # # # # #
    # rotation from side to top
    rot_rodrigues_vec_21 = np.array([[0.6568], [-0.6162], [1.4306]])
    rotM_21, _ = cv2.Rodrigues(rot_rodrigues_vec_21)
    transV_21 = np.array([[0.0979416], [0.3873274], [0.2284719]])

    # Streaming loop
    while True:
        # Get frame set
        frames2 = pipeline2.wait_for_frames()
        # Get depth frame
        depth_frame2 = frames2.get_depth_frame()
        color_frame2 = frames2.get_color_frame()

        # get intrinsics of d415 - top-view
        depth_intrinsic2 = depth_frame2.profile.as_video_stream_profile().intrinsics
        color_intrinsic2 = color_frame2.profile.as_video_stream_profile().intrinsics
        depth_to_color_extrinsic2 = depth_frame2.profile.get_extrinsics_to(
            color_frame2.profile
        )

        # Convert depth_frame to numpy array to render image in opencv
        color_image2 = np.asanyarray(color_frame2.get_data())
        align2 = rs.align(rs.stream.color)
        frames2 = align2.process(frames2)
        fps2, color_image2 = fpsReader.update(color_image2)
        color_image2 = cv2.cvtColor(color_image2, cv2.COLOR_BGR2RGB)
        hands2, color_image2 = detector.findHands(color_image2)

        # aligned depth frame - side view
        aligned_depth_frame2 = frames2.get_depth_frame()

        # Get frames
        frames = pipeline.wait_for_frames()
        # Get depth frame
        depth_frame = frames.get_depth_frame()
        color_frame = frames.get_color_frame()

        # get intrinsics of d415 - top-view
        depth_intrinsic = depth_frame.profile.as_video_stream_profile().intrinsics
        color_intrinsic = color_frame.profile.as_video_stream_profile().intrinsics
        depth_to_color_extrinsic = depth_frame.profile.get_extrinsics_to(
            color_frame.profile
        )

        # Convert depth_frame to numpy array to render image in opencv
        color_image = np.asanyarray(color_frame.get_data())
        align = rs.align(rs.stream.color)
        frames = align.process(frames)
        fps, color_image = fpsReader.update(color_image)
        color_image = cv2.cvtColor(color_image, cv2.COLOR_BGR2RGB)
        hands, color_image = detector.findHands(color_image)

        # aligned depth frame - top view
        aligned_depth_frame = frames.get_depth_frame()

        # # # # # # # # # # # # # # # # # # # #
        # Drawn rectangle for the covered saw #
        # # # # # # # # # # # # # # # # # # # #
        rect_cover = cv2.rectangle(
            color_image, (5, 260), (1280 - 5, 430), (0, 255, 255), 1
        )
        cv2.putText(
            color_image,
            "cover",
            (5, 260 - 10),
            cv2.FONT_HERSHEY_DUPLEX,
            0.9,
            (0, 255, 255),
            1,
        )
        cover_pts = np.array(
            [[5, 260], [1280 - 5, 260], [1280 - 5, 430], [5, 430]], np.int32
        )
        cover_pts = cover_pts.reshape((-1, 1, 2))

        # rectangle saw
        rect_saw = cv2.rectangle(color_image, (220, 330), (1080, 360), (0, 0, 255), 1)
        cv2.putText(
            color_image,
            "saw",
            (220, 330 - 10),
            cv2.FONT_HERSHEY_DUPLEX,
            0.9,
            (0, 0, 255),
            1,
        )
        saw_pts = np.array([[220, 330], [1080, 330], [1080, 360], [220, 360]], np.int32)
        saw_pts = saw_pts.reshape((-1, 1, 2))

        # Top view
        if i == 0:
            if len(hands) != 0:
                image_points = [hands[0]["lmList"][4][0], hands[0]["lmList"][4][1]]
                coordinates_3D = fcm.convert_depth_pixel_to_metric_coordinate(
                    aligned_depth_frame,
                    image_points[0],
                    image_points[1],
                    color_intrinsic,
                )
                if coordinates_3D[2] == 0 and c_dead_top <= 5:
                    image_points_xr = image_points[0] + 1
                    coordinates_3D = fcm.convert_depth_pixel_to_metric_coordinate(
                        aligned_depth_frame,
                        image_points_xr,
                        image_points[1],
                        color_intrinsic,
                    )
                    c_dead_top += 1
                    if coordinates_3D[2] != 0:
                        c_dead_top = 6
                        print("Top view: Dead pixel adjusted!")

                # Set the expected distance to the material
                distance_material = 0.679

                # get average lm depth
                for lm in lm_range:
                    try:
                        lm_coord_lst.append(hands[0]["lmList"][lm])
                        depth_lm = aligned_depth_frame.get_distance(
                            hands[0]["lmList"][lm][0], hands[0]["lmList"][lm][1]
                        )
                        if depth_lm != 0:
                            depth_distance_lst.append(depth_lm)

                        if lm == (lm_range[-1]):
                            mean_depth_lms = sum(depth_distance_lst) / len(
                                depth_distance_lst
                            )
                            # print('depth_distance_list:\n', depth_distance_lst)
                            print("average depth of all lms:", mean_depth_lms)
                            depth_distance_lst.clear()
                    except RuntimeError:
                        pass

                # # # # # # # # # # # # # # # # # #
                # check if hand on material/desk  #
                # # # # # # # # # # # # # # # # # #

                checkInside_lst = []
                if (distance_material - 0.035) <= mean_depth_lms <= distance_material:
                    for lm in lm_range:
                        check_rectangle = cv2.pointPolygonTest(
                            cover_pts, lm_coord_lst[lm], False
                        )
                        if check_rectangle == 1:
                            print(lm, "-lm INSIDE of cover rectangle")
                        elif check_rectangle == -1:
                            print(lm, "-lm NOT in cover rectangle")
                            checkInside_lst.append([int(check_rectangle)])
                            if (
                                all(flag == [-1] for (flag) in checkInside_lst)
                                and len(checkInside_lst) == 21
                            ):
                                # Initial expected distance while the hand sits flat on the material
                                if not dist_init_lst:
                                    for lm_index in lm_rangeT:
                                        ignored_LM_lst = [0, 4, 8, 12, 16]
                                        # print('LM_index:\n', lm_index)
                                        # print('LM_index+1:\n', lm_index+1)
                                        if lm_index in ignored_LM_lst:
                                            continue
                                        dist_init, _ = detector.findDistance(
                                            hands[0]["lmList"][lm_index],
                                            hands[0]["lmList"][lm_index + 1],
                                        )
                                        dist_init_lst.append(dist_init)

                    # find current distance between relevant landmarks
                    # [1-2, 2-3, 3-4, 5-6, 6-7, 7-8,
                    #  9-10, 10-11, 11-12, 13-14, 14-15,
                    #  15-16, 17-18, 18-19, 19-20]
                    if dist_init_lst:
                        for lm_index in lm_rangeT:
                            ignored_LM_lst = [0, 4, 8, 12, 16]
                            if lm_index in ignored_LM_lst:
                                continue
                            dist_current, _ = detector.findDistance(
                                hands[0]["lmList"][lm_index],
                                hands[0]["lmList"][lm_index + 1],
                            )
                            # cv2.circle(color_image, (int(hands[0]['lmList'][lm_index][0]), int(hands[0]['lmList'][lm_index][1])), radius=12, color=(255, 0, 255),
                            #            thickness=-1)
                            dist_current_lst.append(dist_current)

                        # check if 7-8 shortened
                        for finger in occlusion_range:
                            ignored_counter = [1, 2, 4, 5, 7, 8, 10, 11]
                            if finger in ignored_counter:
                                continue
                            # check if 7-8 shortened
                            if dist_current_lst[finger + 2] <= (
                                dist_init_lst[finger + 2] * 0.95
                            ):
                                # check if 6-7 shortened
                                if dist_current_lst[finger + 1] <= (
                                    dist_init_lst[finger + 1] * 0.95
                                ):
                                    # check if 5-6 shortened: no prediction possible
                                    if dist_current_lst[finger] <= (
                                        dist_init_lst[finger] * 0.95
                                    ):
                                        pass

                                    # correct 7
                                    detector.predictLandmark(
                                        dist_init_lst[finger + 1],
                                        dist_current_lst[finger],
                                        (hands[0]["lmList"][finger + 2 + lm_counter]),
                                        (hands[0]["lmList"][finger + 3 + lm_counter]),
                                        color_image,
                                    )
                                    # correct 8
                                    detector.predictLandmark(
                                        (
                                            dist_init_lst[finger + 1]
                                            + dist_init_lst[finger + 2]
                                        ),
                                        dist_current_lst[3],
                                        (hands[0]["lmList"][finger + 2 + lm_counter]),
                                        (hands[0]["lmList"][finger + 3 + lm_counter]),
                                        color_image,
                                    )
                                    baseline_shortened = False
                                    lm_counter += 1
                                if baseline_shortened:
                                    # correct 8
                                    detector.predictLandmark(
                                        dist_init_lst[finger + 2],
                                        dist_current_lst[finger + 1],
                                        (hands[0]["lmList"][finger + 3 + lm_counter]),
                                        (hands[0]["lmList"][finger + 4 + lm_counter]),
                                        color_image,
                                    )
                                    lm_counter += 1
                                baseline_shortened = True

                    lm_coord_lst.clear()
                    dist_current_lst.clear()
                    lm_counter = 0

                # coordinates as np.array
                coordinates_vector = np.array(
                    [[coordinates_3D[0]], [coordinates_3D[1]], [coordinates_3D[2]]]
                )

                # check if current coordinates equal origin
                if coordinates_3D[2] == 0:
                    break
                i = 1

        # Side view
        if j == 0:
            if len(hands2) != 0:
                try:
                    image_points2 = [
                        hands2[0]["lmList"][4][0],
                        hands2[0]["lmList"][4][1],
                    ]
                    coordinates_3D_2 = fcm.convert_depth_pixel_to_metric_coordinate(
                        aligned_depth_frame2,
                        image_points2[0],
                        image_points2[1],
                        color_intrinsic2,
                    )
                    # check for missing depth
                    if coordinates_3D_2[2] == 0 and c_dead_side <= 5:
                        image_points2_xr = image_points2[0] + 1
                        coordinates_3D_2 = fcm.convert_depth_pixel_to_metric_coordinate(
                            aligned_depth_frame,
                            image_points2_xr,
                            image_points2[1],
                            color_intrinsic,
                        )
                        c_dead_side += 1
                        if coordinates_3D_2[2] != 0:
                            c_dead_side = 6
                            print("side view: dead pixel adjusted")
                            # break

                    # coordinates as np.array
                    coordinates_vector_2 = np.array(
                        [
                            [coordinates_3D_2[0]],
                            [coordinates_3D_2[1]],
                            [coordinates_3D_2[2]],
                        ]
                    )
                    coordinates_21 = np.add(
                        np.matmul(rotM_21, coordinates_vector_2), transV_21
                    )

                    # check if current coordinates equal origin
                    if coordinates_3D_2[2] == 0:
                        break
                    j = 1
                except:
                    pass

        # Apply the sequential Unscented_KF (predict, update, ...)
        # for the current landmark

        Unscented_KF.predict()
        Unscented_KF.update(coordinates_3D)
        Unscented_KF.predict()
        Unscented_KF.update(coordinates_3D_2)

        # Adjust position at t-1 to the last expected spot
        Unscented_KF.x_1 = Unscented_KF.x

        # calculate residual
        y, S = Unscented_KF.y, Unscented_KF.S
        eps = y.T @ np.linalg.inv(S) @ y
        epss.append(eps)
        if eps > eps_max:
            Unscented_KF.Q *= Q_scale_factor
            counter += 1
        elif counter > 0:
            Unscented_KF.Q /= Q_scale_factor
            counter -= 1

        i = j = 0

        # Resize image for visualization
        color_image_re = fcv2.image_resize(color_image, height=480)
        # Resize image for visualization
        color_image2_re = fcv2.image_resize(color_image2, height=480)

        # Render image in opencv window
        cTime = time.time()
        fps = 1 / (cTime - pTime)
        pTime = cTime

        # Actual windows - remove comments as needed
        cv2.imshow("Color Stream Side", color_image2_re)
        # cv2.putText(color_image2, str(int(fps)), (10, 70), cv2.FONT_HERSHEY_PLAIN, 3, (22, 0, 255), 2)  # fps-Anzeige
        # cv2.imshow("Depth Stream Side", aligned_depth_frame2)
        # color_image2_re = cv2.resize(color_image2, (720, 1280))

        # Actual windows - remove comments as needed
        cv2.imshow("Color Stream Top", color_image_re)
        # cv2.putText(color_image, str(int(fps)), (10, 70), cv2.FONT_HERSHEY_PLAIN, 3, (22, 0, 255), 2)  # fps-Anzeige
        # cv2.imshow("Depth Stream", aligned_depth_frame)

        k = cv2.waitKey(1)

        if k == 27:
            cv2.destroyAllWindows()
            break

finally:
    pass


'''
Problem 20
n! means n × (n − 1) × ... × 3 × 2 × 1

For example, 10! = 10 × 9 × ... × 3 × 2 × 1 = 3628800,
and the sum of the digits in the number 10! is 3 + 6 + 2 + 8 + 8 + 0 + 0 = 27.

Find the sum of the digits in the number 100!
'''

# Factorial through recursion. Alternatively, you can just you factorial method from python's math library.
def factorial(num):
    if num == 0:
    	return 1
    else:
    	return (num * factorial(num - 1))

x = factorial(100)
sumOfDigits = 0

# Iterate through the digits of the result and plus on each digit to sumOfDigits
for digit in str(x):
	sumOfDigits += int(digit)

print(sumOfDigits)


#!/usr/bin/env python2

from __future__ import division, print_function, absolute_import

import roslib
import rospy
import numpy as np

import rosbag

from tf.transformations import euler_from_quaternion
from geometry_msgs.msg import TransformStamped, Twist
from position_controller import PositionController

from std_msgs.msg import String

class ROSControllerNode(object):
    """ROS interface for controlling ARDrone, using position controller"""

    def __init__(self):
        """Initialize the ROSControllerNode class."""
        self.model_name = 'ARDroneCarre'

        # Publishers
        self.pub_cmd_vel = rospy.Publisher('cmd_vel_RHC',
                                            Twist,
                                            queue_size = 300)

        # Subscribers
        self.sub_vicon_data = rospy.Subscriber('/vicon/{0}/{0}'.format(
                                              self.model_name),
                                              TransformStamped, self.process_vicon_data)

        self.sub_des_pos = rospy.Subscriber('des_pos', String, self.process_desired_position)

        # Initialize messages for publishing
        self.cmd_vel = Twist()

        # Run the onboard controller at 200 Hz
        self.onboard_loop_frequency = 200.

        # Calling the position controller to pass the data
        self.pos_class = PositionController()

        # Run this ROS node at the onboard loop frequency
        self.run_pub_cmd_vel = rospy.Timer(rospy.Duration(1. /
            self.onboard_loop_frequency), self.process_commands)

        # Initialize time
        self.nonunix_time = 0
        self.old_time = rospy.get_time()
        self.start_time = rospy.get_time()
        self.dt = 0

        # Initialize points
        self.current_point = 1
        self.num_points_old = 0

    def process_commands(self, event):
        """Determine the motor speeds and and publishes these."""

        # Calculate time step
        current_time = rospy.get_time()
        self.dt = current_time - self.old_time

        # Get commands from the position controller
        commands = self.pos_class.update_position_controller(self.dt)
        [roll_command, pitch_command, yaw_rate_command, climb_rate_command] = commands

        # Set cmd_vel message values to commands
        self.cmd_vel.linear.x = roll_command
        self.cmd_vel.linear.y = pitch_command
        self.cmd_vel.linear.z = climb_rate_command
        self.cmd_vel.angular.z = yaw_rate_command

        # Publish commands
        self.pub_cmd_vel.publish(self.cmd_vel)

        # Update timestamp
        self.old_time = current_time

    def process_vicon_data(self, vicon_data_msg):

        self.pos_class.current_trans_x = vicon_data_msg.transform.translation.x
        self.pos_class.current_trans_y = vicon_data_msg.transform.translation.y
        self.pos_class.current_trans_z = vicon_data_msg.transform.translation.z

        self.pos_class.current_rot_x = vicon_data_msg.transform.rotation.x
        self.pos_class.current_rot_y = vicon_data_msg.transform.rotation.y
        self.pos_class.current_rot_z = vicon_data_msg.transform.rotation.z
        self.pos_class.current_rot_w = vicon_data_msg.transform.rotation.w

    def process_desired_position(self, pos_msg):

        # Get message containing desired trajectory coordinates from desired_positions
        self.des_pos_msg = np.fromstring(pos_msg.data, dtype = float, sep = ' ')

        # Calculate total number of waypoints
        num_points = self.des_pos_msg.size/4

        # Update number of waypoints
        if num_points != self.num_points_old:
            self.num_points_old = num_points
            self.current_point = 0

        # Convert 1x12 array into 3 arrays of 4 elements (x,y,z)
        trajectory = np.reshape(self.des_pos_msg, (-1, 4))

        self.pos_class.desired_x = trajectory[self.current_point - 1, 0]    # x coordinates
        self.pos_class.desired_y = trajectory[self.current_point - 1, 1]    # y coordinates
        self.pos_class.desired_z = trajectory[self.current_point - 1, 2]    # z coordinates

        self.nonunix_time += self.dt
        # setting deviation tolerance, the higher tolerance, the faster the drone navigates
        if (((self.pos_class.desired_x - 0.5) < self.pos_class.current_trans_x < (self.pos_class.desired_x + 0.5)) and
            ((self.pos_class.desired_y - 0.5) < self.pos_class.current_trans_y < (self.pos_class.desired_y + 0.5)) and
            ((self.pos_class.desired_z - 0.5) < self.pos_class.current_trans_z < (self.pos_class.desired_z + 0.5))):
                if self.current_point < num_points and self.nonunix_time >= 0.01:
                    self.current_point += 1
                    self.nonunix_time = 0

if __name__ == '__main__':
    rospy.init_node('ros_controller')
    ROSControllerNode()
    rospy.spin()


from . import MOTMODEL_REGISTRY

from .base import MOTDINOBASE
import torch
import copy
from collections import deque
from torch import nn
from motlib.mot_models.structures import Instances
from motlib.mot_models.network.dino_mot.tracker.util import RuntimeTrackerBase, TrackerPostProcess
from .default_model import MOTDINO

__all__ = ['TimeTrackDINO']

@MOTMODEL_REGISTRY.register()
class TimeTrackDINO(MOTDINO):

    def _post_process_single_image(self, frame_res, track_instances, is_last):
        track_instances = Instances.cat([self._generate_empty_tracks(), track_instances])

        decoder_out = frame_res['track_dec_out']
        pad_size = decoder_out.get('pad_size', 0)
        det_query_num = decoder_out.get('det_query_num', 0)
        track_query_num = decoder_out.get('track_query_num', 0)

        with torch.no_grad():
            if self.training:
                track_scores = frame_res['pred_logits'][0, :].sigmoid().max(dim=-1).values
            else:
                prob = frame_res['pred_logits'][0, :].sigmoid()
                track_scores, labels = prob.max(-1)

        if self.training:
            pad_size = frame_res['dn_meta']['pad_size']
            track_instances.output_embedding = frame_res['hs_last_layer'][0, pad_size:pad_size+det_query_num+track_query_num]
        else:
            track_instances.output_embedding = frame_res['hs_last_layer'][0, :pad_size+det_query_num+track_query_num]

        track_instances.scores = track_scores[:det_query_num+track_query_num]
        track_instances.pred_logits = frame_res['pred_logits'][0, :det_query_num+track_query_num]
        track_instances.pred_boxes = frame_res['pred_boxes'][0, :det_query_num+track_query_num]

        if track_scores.shape[0] > det_query_num+track_query_num:
            frame_res['track_dec_out']['decoder_out']['time_tracker_pred'] = {}

            frame_res['track_dec_out']['decoder_out']['time_tracker_pred'][5] = {
                'pred_logits': frame_res['pred_logits'][:, det_query_num+track_query_num:],
                'pred_boxes': frame_res['pred_boxes'][:, det_query_num+track_query_num:]
                }
            new_aux_outputs = []
            for layer_i, layer_aux_output in enumerate(frame_res['aux_outputs']):
                frame_res['track_dec_out']['decoder_out']['time_tracker_pred'][layer_i] = {
                'pred_logits': layer_aux_output['pred_logits'][:, det_query_num+track_query_num:],
                'pred_boxes': layer_aux_output['pred_boxes'][:, det_query_num+track_query_num:]
                }
                new_aux_outputs.append({
                'pred_logits': layer_aux_output['pred_logits'][:, :det_query_num+track_query_num],
                'pred_boxes': layer_aux_output['pred_boxes'][:, :det_query_num+track_query_num]
                })

            frame_res['pred_logits'] = frame_res['pred_logits'][:, :det_query_num+track_query_num]
            frame_res['pred_boxes'] = frame_res['pred_boxes'][:, :det_query_num+track_query_num]
            frame_res['aux_outputs'] = new_aux_outputs

        if self.training:
            # the track id will be assigned by the mather.
            frame_res['track_instances'] = track_instances
            track_instances = self.criterion.match_for_single_frame(frame_res)
        else:
            # each track will be assigned an unique global id by the track base.
            self.track_base.update(track_instances)

        if not is_last:
            out_track_instances = self.track_embed(track_instances, frame_res)
            frame_res['track_instances'] = out_track_instances
        else:
            frame_res['track_instances'] = None

        return frame_res