import moviepy.editor as mpy
import tensorflow as tf
from tensorflow.keras import backend as K
from tensorflow.keras.applications.imagenet_utils import preprocess_input
import cv2
from matplotlib import pyplot as plt

try:
    from tf_keras_vis import utils
    from tf_keras_vis.activation_maximization import ActivationMaximization
    from tf_keras_vis.utils.scores import CategoricalScore
    from tf_keras_vis.saliency import Saliency
    from tf_keras_vis.gradcam import GradcamPlusPlus
    from tf_keras_vis.scorecam import ScoreCAM
    from tf_keras_vis.utils import normalize
    from tf_keras_vis.utils.callbacks import Print
except:
    raise Exception("Please install tf-keras-vis: pip install tf-keras-vis")

import donkeycar as dk
from donkeycar.parts.tub_v2 import Tub
from donkeycar.utils import *
import numpy as np

DEG_TO_RAD = math.pi / 180.0

cfg = None

class MakeMovie(object):

    def run(self, args, parser):
        '''
        Load the images from a tub and create a movie from them.
        Movie
        '''
        global cfg

        if args.tub is None:
            print("ERR>> --tub argument missing.")
            parser.print_help()
            return

        conf = os.path.expanduser(args.config)
        if not os.path.exists(conf):
            print("No config file at location: %s. Add --config to specify\
                 location or run from dir containing config.py." % conf)
            return

        cfg = dk.load_config(conf)

        if args.type is None and args.model is not None:
            args.type = cfg.DEFAULT_MODEL_TYPE
            print("Model type not provided. Using default model type from config file")

        if args.salient:
            if args.model is None:
                print("ERR>> salient visualization requires a model. Pass with the --model arg.")
                parser.print_help()

            #if args.type not in ['linear', 'categorical']:
            #    print("Model type {} is not supported. Only linear or categorical is supported for salient visualization".format(args.type))
            #    parser.print_help()
            #    return

        self.tub = Tub(args.tub)

        start = args.start
        self.end_index = args.end if args.end != -1 else len(self.tub)
        num_frames = self.end_index - start

        # Move to the correct offset
        self.current = 0
        self.iterator = self.tub.__iter__()
        while self.current < start:
            self.iterator.next()
            self.current += 1

        self.scale = args.scale
        self.keras_part = None
        self.do_salient = False
        self.user = args.draw_user_input
        self.pilot_angle = 0.0
        self.pilot_throttle = 0.0
        self.pilot_score = 1.0 # used for color intensity
        self.user_angle = 0.0
        self.user_throttle = 0.0
        self.control_score = 0.25 # used for control size
        self.flag_test_pilot_angle = 1
        self.flag_test_pilot_throttle = 1
        self.flag_test_user_angle = 1
        self.flag_test_user_throttle = 1
        self.throttle_circle_pilot_angle = 0
        self.throttle_circle_user_angle = 0
        self.is_test = False
        self.last_pilot_throttle = 0.0 # used for color transparency
        self.last_user_throttle = 0.0 # used for color transparency
        self.pilot_throttle_trans = 1.0 # used for color transparency
        self.user_throttle_trans = 1.0 # used for color transparency
        self.pilot_throttle_trans_rate = 0.25 # used for color transparency
        self.user_throttle_trans_rate = 0.25 # used for color transparency

        if args.model is not None:
            self.keras_part = get_model_by_type(args.type, cfg=cfg)
            self.keras_part.load(args.model)
            if args.salient:
                self.do_salient = self.init_salient(self.keras_part.interpreter.model)

        print('making movie', args.out, 'from', num_frames, 'images')
        clip = mpy.VideoClip(self.make_frame, duration=((num_frames - 1) / cfg.DRIVE_LOOP_HZ))
        clip.write_videofile(args.out, fps=cfg.DRIVE_LOOP_HZ)

    @staticmethod
    def draw_line_into_image(angle, throttle, is_pred, img, color):
        """
        is_pred:
            True: from draw_model_prediction()
            False: from draw_user_input()
        """

        height, width, _ = img.shape
        mid_h = height//2
        length = height//4
        a1 = angle * 45.0
        l1 = throttle * length
        mid_w = width // 2 + (- 1 if is_pred else +1)

        p1 = tuple((mid_w - 2, mid_h - 1))
        p11 = tuple((int(p1[0] + l1 * math.cos((a1 + 270.0) * DEG_TO_RAD)),
                     int(p1[1] + l1 * math.sin((a1 + 270.0) * DEG_TO_RAD))))

        cv2.line(img, p1, p11, color, 2)

    def get_user_input(self, record, img):
        """
        Get the user input from record
        """
        if self.is_test:
            self.user_angle_a = float(record["user/angle"])
            self.user_throttle_a = float(record["user/throttle"])
        else:
            self.user_angle = float(record["user/angle"])
            self.user_throttle = float(record["user/throttle"])

    def get_model_prediction(self, img, salient_image):
        """
        Get the pilot input from model prediction
        """
        if self.keras_part is None:
            return

        #expected = tuple(self.keras_part.get_input_shape()[1:])
        expected = self.keras_part.interpreter.model.inputs[0].shape[1:]
        actual = img.shape

        # if model expects grey-scale but got rgb, covert
        if expected[2] == 1 and actual[2] == 3:
            # normalize image before grey conversion
            grey_img = rgb2gray(img)
            actual = grey_img.shape
            img = grey_img.reshape(grey_img.shape + (1,))

        if expected != actual:
            print(f"expected input dim {expected} didn't match actual dim "
                  f"{actual}")
            return

        if self.is_test:
            self.pilot_angle_a, self.pilot_throttle_a = self.keras_part.run(img)
        else:
            self.pilot_angle, self.pilot_throttle = self.keras_part.run(img)

    def draw_steering_distribution(self, img, salient_image):
        """
        query the model for it's prediction, draw the distribution of
        steering choices
        """
        from donkeycar.parts.keras import KerasCategorical

        if self.keras_part is None or type(self.keras_part) is not KerasCategorical:
            return

        pred_img = normalize_image(img)
        pred_img = pred_img.reshape((1,) + pred_img.shape)
        angle_binned, _ = self.keras_part.interpreter.predict(pred_img)

        x = 4
        dx = 4
        y = cfg.IMAGE_H - 4
        iArgMax = np.argmax(angle_binned)
        for i in range(15):
            p1 = (x, y)
            p2 = (x, y - int(angle_binned[0][i] * 100.0))
            if i == iArgMax:
                cv2.line(salient_image, p1, p2, (255, 0, 0), 2)
            else:
                cv2.line(salient_image, p1, p2, (200, 200, 200), 2)
            x += dx

    def init_salient(self, model):
        # Utility to search for layer index by name.
        # Alternatively we can specify this as -1 since it corresponds to the last layer.
        model.summary()
        self.output_names = []

        for i, layer in enumerate(model.layers):
            if "dropout" not in layer.name.lower() and "out" in layer.name.lower():
                self.output_names += [layer.name]

        if len(self.output_names) == 0:
            print("Failed to find the model layer named with 'out'. Skipping salient.")
            return False

        print("####################")
        print("Visualizing activations on layer:", *self.output_names)
        print("####################")

        # Create Saliency object.
        # If `clone` is True(default), the `model` will be cloned,
        # so the `model` instance will be NOT modified, but it takes a machine resources.
        self.saliency = Saliency(model,
                                 model_modifier=self.model_modifier,
                                 clone=False)
        # Create GradCAM++ object, Just only repalce class name to "GradcamPlusPlus"
        self.gradcampp = GradcamPlusPlus(model,
                                         model_modifier=self.model_modifier,
                                         clone=False)

        return True

    def draw_gradcam_pp(self, img):

        x = preprocess_input(img, mode='tf')

        # Generate heatmap with GradCAM++
        salient_map = self.gradcampp(self.loss,
                             x,
                             penultimate_layer=-1, # model.layers number
                             )

        return self.draw_mask(img, salient_map)

    def draw_salient(self, img):
        # https://github.com/keras-team/keras-applications/blob/master/keras_applications/imagenet_utils.py
        x = preprocess_input(img, mode='tf')

        # Generate saliency map with smoothing that reduce noise by adding noise
        salient_map = self.saliency(self.loss, x)
        return self.draw_mask(img, salient_map)

    def draw_mask(self, img, salient_map):
        if salient_map[0].size != cfg.IMAGE_W * cfg.IMAGE_H:
            print("salient size failed.")
            return

        salient_map = salient_map[0]
        salient_map = salient_map * 255
        salient_mask = cv2.cvtColor(salient_map, cv2.COLOR_GRAY2RGB)
        salient = cv2.applyColorMap(salient_map.astype('uint8'), cv2.COLORMAP_JET)
        salient = cv2.applyColorMap(salient, cv2.COLOR_BGR2RGB)
        salient = cv2.bitwise_and(salient, salient_mask.astype('uint8'))

        blend = cv2.addWeighted(img, 1.0, salient, 1.0, 0.0)

        return blend

    def make_frame(self, t):
        '''
        Callback to return an image from from our tub records.
        This is called from the VideoClip as it references a time.
        We don't use t to reference the frame, but instead increment
        a frame counter. This assumes sequential access.
        '''

        if self.current >= self.end_index:
            return None

        rec = self.iterator.next()
        img_path = os.path.join(self.tub.images_base_path, rec['cam/image_array'])
        camera_image = img_to_arr(Image.open(img_path))

        salient_image = None
        if self.do_salient:
            salient_image = self.draw_salient(camera_image)
            #salient_image = self.draw_gradcam_pp(camera_image)
            salient_image = salient_image.astype('uint8')
        if salient_image is None:
            salient_image = camera_image
        # resize
        if self.scale != 1:
            h, w, d = salient_image.shape
            dsize = (w * self.scale, h * self.scale)
            salient_image = cv2.resize(salient_image, dsize=dsize, interpolation=cv2.INTER_CUBIC)

        # draw control
        if self.keras_part is not None:
            self.get_model_prediction(camera_image, salient_image)
            self.draw_steering_distribution(camera_image, salient_image)
        if self.user: self.get_user_input(rec, salient_image)
        salient_image = self.draw_control_into_image(salient_image)

        """
        # left upper text
        display_str = []
        display_str.append(f"pilot_angle: {self.pilot_angle:.2f}")
        display_str.append(f"pilot_throttle: {self.pilot_throttle:.2f}")
        self.draw_text(salient_image, display_str)
        """

        self.current += 1
        # returns a 8-bit RGB array
        return salient_image

    def model_modifier(self, m):
        m.layers[-1].activation = tf.keras.activations.linear

    def loss(self, output):
        return (output[0])

    def draw_control_into_image(self, img):
        """
        test
        self.pilot_score = np.random.default_rng().uniform(low=0.2, high=1.0)
        self.pilot_throttle = np.random.default_rng().uniform(low=-1.0, high=1.0)
        self.pilot_angle = np.random.default_rng().uniform(low=-1.0, high=1.0)
        """
        if self.is_test:
            if self.pilot_angle >= 1.0:
                self.flag_test_pilot_angle = -0.1
            elif self.pilot_angle <= -1.0:
                self.flag_test_pilot_angle = +0.1
            self.pilot_angle += self.flag_test_pilot_angle

            if self.pilot_throttle >= 1.0:
                self.flag_test_pilot_throttle = -0.01
            elif self.pilot_throttle <= -1.0:
                self.flag_test_pilot_throttle = +0.01
            self.pilot_throttle += self.flag_test_pilot_throttle

            if self.user_angle >= 1.0:
                self.flag_test_user_angle = -0.05
            elif self.user_angle <= -1.0:
                self.flag_test_user_angle = +0.05
            self.user_angle += self.flag_test_user_angle

            if self.user_throttle >= 1.0:
                self.flag_test_user_throttle = -0.02
            elif self.user_throttle <= -1.0:
                self.flag_test_user_throttle = +0.02
            self.user_throttle += self.flag_test_user_throttle

            #self.control_score = abs(self.pilot_throttle/4*3) + 0.25

        height, width, _ = img.shape
        y = height//2
        x = width//2

        # prepare ellipse mask
        r_base = 6.0
        r_mask = int(height//r_base)
        white = np.ones_like(img)*255
        ellipse = np.zeros_like(img)

        r_pilot = int(height//r_base+(height//3.1)*self.control_score)
        r_user = int(height//r_base+(height//6.2)*self.control_score)

        pilot_trans, user_trans = self.trans_make()

        # draw pilot control
        green = (0,int(255*pilot_trans),0) # green for reverse
        blue = (0,0,int(255*pilot_trans)) # blue for reverse
        self.draw_ellipse(self.pilot_angle, self.pilot_throttle, x,y,r_pilot, ellipse, green, blue)

        pilot_circle_mask = cv2.circle(white, (int(x),int(y)), int(r_user+(height//10)*self.control_score), (0,0,0), -1).astype('uint8')
        # pilot mask
        ellipse = cv2.bitwise_and(ellipse, pilot_circle_mask)

        # draw user control
        green = (0,int(255*user_trans),0) # green for reverse
        blue = (0,0,int(255*user_trans)) # blue for reverse
        orange = (255,69,0) # orange for reverse
        self.draw_ellipse(self.user_angle, self.user_throttle, x,y,r_user, ellipse, green, blue)
        white = np.ones_like(img)*255

        user_circle_mask = cv2.circle(white, (int(x),int(y)), int(r_mask+(height//10)*self.control_score), (0,0,0), -1).astype('uint8')
        # user mask
        ellipse = cv2.bitwise_and(ellipse, user_circle_mask)

        # draw circle
        color1= (0,0,215)
        color2= (0,25,78)
        color3= (0,255,215)
        cv2.circle(ellipse, (int(x),int(y)), int(r_pilot), color1, 1)
        cv2.circle(ellipse, (int(x),int(y)), int(r_user), color2, 1)
        cv2.circle(ellipse, (int(x),int(y)), int(r_mask), color3, 1)

        # draw dot circle
        ellipse = self.draw_dot_circle2(ellipse, x,y,r_pilot+1+self.scale//4, 1+self.scale//4, (0,255,218), -1, 4*self.scale, self.pilot_throttle, True)
        ellipse = self.draw_dot_circle2(ellipse, x,y,r_user+1+self.scale//4, 1+self.scale//4, (0,255,218), -1, 4*self.scale, self.user_throttle, False)

        # draw speed meter
        #self.draw_analog_meter(ellipse, self.pilot_speed)
        self.draw_analog_direction_meter(ellipse, self.pilot_angle, self.pilot_throttle)
        ellipse = self.draw_digital_meter(ellipse, r_mask-(2+self.scale//4), self.pilot_throttle, 18*self.scale, pilot_trans)
        ellipse = self.draw_digital_meter(ellipse, r_mask-(2+self.scale//4)*2-(width//10)*self.control_score, self.user_throttle, 12*self.scale, user_trans)

        #self.draw_analog_meter(ellipse, -0.75)
        #print(f"r_pilot: {r_pilot}")
        # draw stripe circle
        if self.pilot_throttle >= 0:
            flag_pilot_throttle = 1
            add_pilot_deg = 180
        else:
            flag_pilot_throttle = -1
            add_pilot_deg = 0
        if self.user_throttle >= 0:
            flag_user_throttle = 1
            add_user_deg = 180
        else:
            flag_user_throttle = -1
            #add_user_deg = 0
            add_user_deg = 180
        """
        p1 = (x,y-flag_pilot_throttle*r_pilot)
        p2 = (x,y+flag_pilot_throttle*y//10)
        red = (255,0,0)
        pts = self.points_rotation([p1,p2], center=(x,y), degrees=flag_pilot_throttle*self.pilot_angle*90)
        cv2.line(ellipse, pts[0], pts[1], red, 2)
        cv2.line(ellipse, p1, p2, red, 2)
        """
        red=(255,0,0)
        #self.draw_arc_center_line(ellipse, x,y,r_pilot+1+(height//40),r_user+(height//20)*self.control_score, red, 2, degrees=flag_pilot_throttle*self.pilot_angle*90+add_pilot_deg)
        self.draw_arc_center_line(ellipse, x,y,r_user+1+(height//40),r_mask+(height//20)*self.control_score, red, 2, degrees=flag_user_throttle*self.user_angle*90+add_user_deg)

        # left upper text
        """
        display_str = []
        display_str.append(f"x: {x}")
        display_str.append(f"y: {y}")
        display_str.append(f"r_pilot: {r_pilot}")
        display_str.append(f"angle: {self.pilot_angle:.2f}")
        display_str.append(f"throttle: {self.pilot_throttle:.2f}")
        self.draw_text(ellipse, display_str)
        """

        # blur
        if self.scale <= 3:
            ellipse = cv2.GaussianBlur(ellipse,(3,3),0)
        else:
            ellipse = cv2.GaussianBlur(ellipse,(5,5),0)

        self.last_pilot_throttle = self.pilot_throttle
        self.last_user_throttle = self.user_throttle

        return cv2.addWeighted(img, 1.0, ellipse, 1.0, 0.0)

    def trans_make(self):
        if self.pilot_throttle != self.last_pilot_throttle:
            if self.pilot_throttle_trans >= 0.8:
                self.pilot_throttle_trans_rate = -0.25
            elif self.pilot_throttle_trans <= 0.3:
                self.pilot_throttle_trans_rate = +0.25
            self.pilot_throttle_trans += self.pilot_throttle_trans_rate
        else:
            self.pilot_throttle_trans = 1.0
        if self.user_throttle != self.last_user_throttle:
            if self.user_throttle_trans >= 0.8:
                self.user_throttle_trans_rate = -0.25
            elif self.user_throttle_trans <= 0.3:
                self.user_throttle_trans_rate = +0.25
            self.user_throttle_trans += self.user_throttle_trans_rate
        else:
            self.user_throttle_trans = 1.0
        return self.pilot_throttle_trans, self.user_throttle_trans

    def draw_analog_meter(self, img, value):
        height, width, _ = img.shape
        y = height//2
        x = width//2

        r_base = 6.0
        r_pilot = int(height//r_base+(height//3.1)*self.control_score)
        red = (255,0,0)

        p1 = (x,y-r_pilot)
        p2 = (x,y+y//10)
        pts = self.points_rotation([p1,p2], center=(x,y), degrees=value*180-90)
        cv2.line(img, tuple(pts[0]), tuple(pts[1]), red, 2)

    def draw_analog_direction_meter(self, img, angle, throttle):
        height, width, _ = img.shape
        y = height//2
        x = width//2

        r_base = 6.0
        r_pilot = int(height//r_base+(height//3.1)*self.control_score)
        red = (255,0,0)

        p1 = (x,y-r_pilot)
        p2 = (x,y+y//10)
        if throttle >= 0:
            flag_throttle = 1
            add_deg = 0
        else:
            flag_throttle = -1
            #add_deg = 180
            add_deg = 0

        pts = self.points_rotation([p1,p2], center=(x,y), degrees=flag_throttle*angle*90+add_deg)
        cv2.line(img, tuple(pts[0]), tuple(pts[1]), red, 2)

    def draw_arc_center_line(self, img, x,y,r1,r2, color, thickness, degrees):
        base_point1 = np.array([0,r1])
        base_point2 = np.array([0,r2])
        rot = self.rot(degrees)
        p1 = np.dot(rot, base_point1) # NEVER DO .astype('uint8'). x,y is more than 255.
        p2 = np.dot(rot, base_point2)
        cv2.line(img, (int(p1[0]+x),int(p1[1]+y)), (int(p2[0]+x),int(p2[1]+y)), color, thickness)

    def color_make(self,value):
        """
        Rainbow color maker.
        value: -1.0 to 1.0

        abs(value) 0.0: blue
        abs(value) 0.5: green
        abs(value) 1.0: red
        """
        value = abs(value)
        if value > 1:
            value = 1

        c = int(255*value)
        c1 = int(255*(value*2-0.5))
        c05 = int(255*value*2)
        if c > 255:
            c = 255
        elif c < 0:
            c = 0
        if c1 > 255:
            c1 = 255
        elif c1 < 0:
            c1 = 0
        if c05 > 255:
            c05 = 255
        elif c05 < 0:
            c05 = 0

        if 0 <= value and value < 0.5:
            color = (0,c05,255-c05) # blue -> green
        elif 0.5 <= value and value <= 1.0:
            color = (c1,c05-c1,0) # green -> red
        elif 1.0 < value:
            color = (255,0,0) # red

        return color

    def draw_digital_meter(self, img, r_mask, value, num, trans):
        """
        Rainbow digital throttle meter.
        img: image to draw
        r_mask: circumferential radius
        value: -1.0 to 1.0
        num: number of boxes
        trans: when the value is changed, boxes transparency is changed with this value
        """
        if num > 36:
            num = 36
        height, width, _ = img.shape
        y = height//2
        x = width//2
        h = (height//20)*self.control_score
        w = (width//10)*self.control_score
        base_points = np.array([[r_mask,h/2],[r_mask,-h/2],[-w+r_mask,-h/2],[-w+r_mask,h/2]])
        dot_angle = 360.0/num
        center = np.atleast_2d((x,y))
        start_angle = 145

        mask_img = np.zeros_like(img)
        meter_img = np.zeros_like(img)
        self.draw_digital_mask(value,x,y,r_mask,mask_img)

        for i in range(0, num):
            deg = i * dot_angle + start_angle
            box = self.points_rotation(base_points,center=(0,0),degrees=deg)
            box = ((box.T + center.T).T)
            if i == 0:
                color = (0,0,255*trans)
            else:
                if value >= 0:
                    color = self.color_make(i*dot_angle/270)
                else:
                    color = self.color_make((360-i*dot_angle)/270)
                color=(color[0]*trans,color[1]*trans,color[2]*trans)
            cv2.fillPoly(meter_img, np.array([box], dtype=np.int32), color)

        meter_img = cv2.bitwise_and(meter_img.astype('uint8'), mask_img.astype('uint8'))
        img = cv2.bitwise_or(img, meter_img.astype('uint8'))
        return img

    def draw_dot_circle(self, img, x,y,r, cr, color, thickness, num):
        """
        Draw outer dot circle.
        """
        if num > 36:
            num = 36
        base_point = np.array([r,0])
        dot_angle = 360.0/num
        for i in range(0, num):
            deg = i * dot_angle
            rot = self.rot(deg)
            rotated = np.dot(rot, base_point)
            img = cv2.circle(img, (int(rotated[0]+x),int(rotated[1]+y)), cr, color, thickness).astype('uint8')
        return img

    def draw_dot_circle2(self, img, x,y,r, cr, color, thickness, num, throttle, is_pilot):
        """
        Draw outer dot circle.
        This circle rotates according to the throttle.
        """
        if num > 36:
            num = 36
        dot_angle = 360.0/num
        rot_spd = throttle*dot_angle
        if rot_spd > dot_angle/2.0: # rotation speed limit
            rot_spd = dot_angle/2.0
        elif rot_spd < -dot_angle/2.0:
            rot_spd = -dot_angle/2.0
        if is_pilot:
            self.throttle_circle_pilot_angle += rot_spd
            start_angle = self.throttle_circle_pilot_angle
        else:
            self.throttle_circle_user_angle += rot_spd
            start_angle = self.throttle_circle_user_angle
        rot = self.rot(start_angle)
        base_point = np.dot(rot,np.array([r,0]))
        throttle = abs(throttle)
        for i in range(0, num):
            deg = i * dot_angle
            rot = self.rot(deg)
            rotated = np.dot(rot, base_point)

            color = self.color_make(throttle)
            img = cv2.circle(img, (int(rotated[0]+x),int(rotated[1]+y)), cr, color, thickness).astype('uint8')
            #print(f'i: {i}, deg: {deg}, (x,y): ({int(rotated[0]+x),int(rotated[1]+y)})')
        return img

    def rot(self, degrees):
        rad = np.deg2rad(degrees)
        return np.array([[np.cos(rad), -np.sin(rad)],
                         [np.sin(rad), np.cos(rad)]])

    def points_rotation_test(self, pts, center=(0,0), degrees=0):
        """
        the same as points_rotation()
        """
        res = None
        for pt in pts:
            base_point = np.array([pt[0]-center[0], pt[1]-center[1]])
            rot = self.rot(degrees)
            p = np.dot(rot, base_point)
            p = (int(p[0]+center[0]), int(p[1]+center[1]))
            if res is None:
                res = p
            else:
                res = np.vstack((res,p))
        return res

    def points_rotation(self, pts, center=(0,0), degrees=0):
        """
        https://stackoverflow.com/questions/34372480/rotate-point-about-another-point-in-degrees-python
        This function rotates the polygon coordinates.
        pts: 2D polygon coordinates
        center: Center coordinates of rotation
        degrees: angle
        """
        R = self.rot(degrees)
        o = np.atleast_2d(center)
        pts = np.atleast_2d(pts)
        return np.squeeze((R @ (pts.T-o.T) + o.T).T).astype('int32') # 'int16' or more. x,y is larger than 255. cv2.fillPoly required int32

    def draw_ellipse(self, angle, throttle, x,y,r, img, color, reverse_color):
        """
        Draw the outer arc.
        """
        angle_deg = int(throttle * 270)
        if throttle < 0:
            center_deg = +90 + abs(angle_deg/2)
            color = reverse_color
        else:
            center_deg = -90 - abs(angle_deg/2)

        if throttle < 0:
            angle = angle * -1.0
        rotate_deg = int(angle * 90) + center_deg

        cv2.ellipse(img,(int(x),int(y)),(int(r),int(r)),rotate_deg,0,angle_deg,color,-1)

    def draw_digital_mask(self, throttle, x,y,r, img):
        """
        throttle: -1.0 to 1.0
        x: x coordinate in the center of the image
        y: y coordinate in the center of the image
        r: radius of arc
        img: mask image
        """
        angle_deg = int(throttle * 270) # ellipse angle
        rotate_deg = 145 # ellipse start position
        color=(255,255,255)
        cv2.ellipse(img,(int(x),int(y)),(int(r),int(r)),rotate_deg,0,angle_deg,color,-1)

    def draw_text(self, img, display_str=[]):
        """ STATISTICS FONT """
        fontScale = img.shape[0]/1000.0
        if fontScale < 0.4:
            fontScale = 0.4
        fontThickness = 1 + int(fontScale)
        fontFace = cv2.FONT_HERSHEY_SIMPLEX

        max_text_width = 0
        max_text_height = 0

        """ DRAW BLACK BOX AND TEXT """
        [(text_width, text_height), baseLine] = cv2.getTextSize(text=display_str[0], fontFace=fontFace, fontScale=fontScale, thickness=fontThickness)
        x_left = int(baseLine)
        y_top = int(baseLine)
        for i in range(len(display_str)):
            [(text_width, text_height), baseLine] = cv2.getTextSize(text=display_str[i], fontFace=fontFace, fontScale=fontScale, thickness=fontThickness)
            if max_text_width < text_width:
                max_text_width = text_width
            if max_text_height < text_height:
                max_text_height = text_height
        """ DRAW BLACK BOX """
        cv2.rectangle(img, (x_left - 2, int(y_top)), (int(x_left + max_text_width + 2), int(y_top + len(display_str)*max_text_height*1.2+baseLine)), color=(0, 0, 0), thickness=-1)
        """ DRAW FPS, TEXT """
        for i in range(len(display_str)):
            cv2.putText(img, display_str[i], org=(x_left, y_top + int(max_text_height*1.2 + (max_text_height*1.2 * i))), fontFace=fontFace, fontScale=fontScale, thickness=fontThickness, color=(77, 255, 9))


import logging
import numpy as np
import requests
import json
import pickle

try:
    import superprzm  # Import superprzm.dll / .so
    _dll_loaded = True
except ImportError as e:
    logging.exception(e)
    _dll_loaded = False

def run(jid, sam_bin_path, name_temp, section, array_size):
    # np_array_out = np.random.rand(50,3)  # Dummy NumPy data

    # Run SuperPRZM as DLL
    print("sam_bin_path: ", sam_bin_path)
    np_array_huc_ids, np_array_out = superprzm.runmain.run(sam_bin_path, name_temp, section, array_size)

    huc_ids = create_huc_ids_list(np_array_huc_ids)

    # Send Numpy array to MongoDB/Motor server
    mongo_motor_insert(jid, huc_ids, np_array_out, name_temp, section)  # Motor only works on Unix-based OS

    return True

def mongo_motor_insert(jid, huc_ids, np_array, name_temp, section):
    # TODO: Remove the next line (was previously used for testing, production 'jid' is the 'jid' param)
    # jid = name_temp + "_" +section
    url = 'http://localhost:8787/sam/daily/' + jid
    # http_headers = {'Content-Type': 'application/json'}
    http_headers = {'Content-Type': 'application/octet-stream'}
    # data = json.dumps(create_mongo_document(np_array, name_temp, section))
    data = serialize(jid, huc_ids, np_array, name_temp, section)

    # Send data to Mongo server
    requests.post(url, data=data, headers=http_headers, timeout=30)

def create_huc_ids_list(np_array_huc_ids):
    """
    Create a Python list of HUC_IDs from the SuperPRZM output numpy array of HUC_IDs
    :param np_array_huc_ids: numpy array, [['0', '5', '0', '0', '1', '1', '0', '5', '4', '2', '1', '1'], [...]]
    :return: list
    """
    out_list = []
    for char_array in np_array_huc_ids:
        out_list.append("".join(list(char_array)))

    return out_list

def serialize(jid, huc_ids, np_array, name_temp, section):
    """
    Returns pickle to be sent to Mongo server.
    :param np_array:
    :param name_temp:
    :param section:
    :return:
    """
    return pickle.dumps({
        "jid": jid,
        'output': np_array,
        'huc_ids': huc_ids
    }, protocol=2)


# First import library
import pyrealsense2 as rs
# Import Numpy for easy array manipulation
import numpy as np
# Import OpenCV for easy image rendering
import cv2
# Import argparse for command-line options
import argparse
# Import os.path for file path manipulation
import os.path

# Create object for parsing command-line options
parser = argparse.ArgumentParser(description="Read recorded bag file and display depth stream in jet colormap.\
                                Remember to change the stream fps and format to match the recorded.")
# Add argument which takes path to a bag file as an input
parser.add_argument("-i", "--input", type=str, help="Path to the bag file", default="guilherme-rgbd.bag")
# Parse the command line arguments to an object
args = parser.parse_args()
# Safety if no parameter have been given
if not args.input:
    print("No input paramater have been given.")
    print("For help type --help")
    exit()
# Check if the given file have bag extension
if os.path.splitext(args.input)[1] != ".bag":
    print("The given file is not of correct file format.")
    print("Only .bag files are accepted")
    exit()

# Declare pointcloud object, for calculating pointclouds and texture mappings
pc = rs.pointcloud()

# We want the points object to be persistent so we can display the last cloud when a frame drops
points = rs.points()

# Create a pipeline
pipeline = rs.pipeline()

# Create a config and configure the pipeline to stream
#  different resolutions of color and depth streams
config = rs.config()

# Tell config that we will use a recorded device from file to be used by the pipeline through playback.
rs.config.enable_device_from_file(config, args.input)

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

# Start streaming
profile = pipeline.start(config)

# We'll use the colorizer to generate texture for our PLY
# (alternatively, texture can be obtained from color or infrared stream)
colorizer = rs.colorizer()

# Getting the depth sensor's depth scale (see rs-align example for explanation)
depth_sensor = profile.get_device().first_depth_sensor()
depth_scale = depth_sensor.get_depth_scale()
print("Depth Scale is: " , depth_scale)

# We will be removing the background of objects more than
#  clipping_distance_in_meters meters away
clipping_distance_in_meters = 1 #1 meter
clipping_distance = clipping_distance_in_meters / depth_scale

# Create an align object
# rs.align allows us to perform alignment of depth frames to others frames
# The "align_to" is the stream type to which we plan to align depth frames.
align_to = rs.stream.color
align = rs.align(align_to)

# Streaming loop
try:
    # Get frameset of color and depth
    frames = pipeline.wait_for_frames()

    # Align the depth frame to color frame
    aligned_frames = align.process(frames)

    # Get aligned frames
    aligned_depth_frame = aligned_frames.get_depth_frame() # aligned_depth_frame is a 640x480 depth image
    color_frame = aligned_frames.get_color_frame()

    colorized = colorizer.process(aligned_frames)

    # Create save_to_ply object
    ply = rs.save_to_ply("test.ply")

    # Set options to the desired values
    # In this example we'll generate a textual PLY with normals (mesh is already created by default)
    ply.set_option(rs.save_to_ply.option_ply_binary, False)
    ply.set_option(rs.save_to_ply.option_ply_normals, True)

    print("Saving to test.ply...")
    # Apply the processing block to the frameset which contains the depth frame and the texture
    ply.process(colorized)
    print("Done")

finally:
    pipeline.stop()


#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import numpy as np
import math

import neurogym as ngym
from neurogym import spaces

# TODO: Move to collection
class Pneumostomeopening(ngym.BaseEnv):
    metadata = {
        'paper_link': 'https://jeb.biologists.org/content/199/3/683.long',
        'paper_name': '''Operant conditioning of aerial respiratory behaviour
        in Lymnaea stagnalis''',
        'tags': ['operant conditioning', 'learning', 'Lymnaea Stagnalis',
                 'associative learning', 'aerial respiratory behaviour',
                 'hypoxia', 'molluscan model', 'system', 'snail']
    }

    def __init__(self, dt=100, rewards=None, timing=None, sigma=1.0,
                 extra_input_param=None):
        super().__init__(dt=dt)

        self.t = 0

        ## action and observation spaces:
        ## action space
        ## breathing actions: [breathing pneumostome, breathing skin]
        ## swim actions: [up, down, left, right]
        ## make it 3D (forward/backward) or 2D is enough ?
        ## observation space
        ## Discretizing beaker into level of depth
        name = {
            'depth1': 0,
            'depth2': 1,
            'depth2': 2,
            'depth2': 3}
        self.observation_space = spaces.Discrete(4, name=name)
        name = {
            'breathingpneumostome': 0,
            'breathingskin': 1,
            'up': 2,
            'down': 3}
        self.action_space = spaces.Discrete(4, name=name)
        # self.action_space = spaces.MultiDiscrete([(0, 2), (0, 2)])

        ## TODO: can do breathingpneumostome only when at the surface ==depth 1

        ## task specific variables
        self.oxygen_level = 10

    def _get_new_oxygen_level(self, action):
        """
        Update oxygen level of the agent based on action. Follow an
        exponential decay.
        """
        decay_constant = 1

        if action == 0: #breathingpneumostome
            self.oxygen_level += 5
        elif action == 1: #breathingskin
            self.oxygen_level +=1
        else : #swimming actions
            self.oxygen_level = self.oxygen_level * math.exp(-decay_constant * self.t)
        self.oxygen_level = int(round(self.oxygen_level))
        print('oxygen_level: {}'.format(self.oxygen_level))
        return self.oxygen_level

    def phase(self, phase, **kwargs):
        """
        different phase of training.
        set_default_behavior trains the network so that breathing behavior corresponds
        to the default states of the animal where breathing through the peumostome
        is a common occurence.
        Training_session train the network for operant-conditioning of respiratory
        behaviour pattern. Simulating poking the pneumostome as a negative reward.
        Testing_session test the respiratory behaviour change after training_session
        has been run.

        """
        ## Rewards
        if phase == 'set_default_behavior':
            rewards = {
                'swim': +0.1,  # reward given for anything else
                'skin': +0.5,  # reward given when breathing through the skin
                'pneumostome': +1.  # reward given when breathing through the pneumostome
                #'miss': # to force some type of trial duration limit if implementing trials in this task ?
            }
        elif phase == 'training_session':
            rewards = {
                'swim': +0.1,
                'skin': +1.,
                'pneumostome': -1.
            }
        elif phase == 'testing_session':
            rewards = {
                'swim': 0.,
                'skin': 0.,
                'pneumostome': 0.
            }
        return rewards

    def reset(self):
        self.agent_pos = np.random.randint(7)
        self.done = False
        return self.agent_pos

    def step(self, action):
        new_trial = False

        # define the phase of the behavior
        self.rewards = self.phase('set_default_behavior')

        # update reward depending on action and phase of behavioral training
        # update agent position
        if action == 0: #breathingpneumostome
            self.reward = self.rewards['pneumostome']
            self.agent_pos = self.agent_pos
        elif action == 1: #breathingskin
            self.reward = self.rewards['skin']
            self.agent_pos = self.agent_pos
        else: #up #down
            self.reward = self.rewards['swim']
            if action == 2:
                self.agent_pos = self.agent_pos + 1
            if action == 3:
                self.agent_pos = self.agent_pos - 1

        self.agent_pos = np.clip(self.agent_pos, 0, len(self.observation_space.name))

        # updating oxygen level according to action taken
        self.oxygen_level = self._get_new_oxygen_level(action)

        # if oxygen level drop at 0 or below then end experiment
        if self.oxygen_level <= 0:
            self.done = True

        self.t = self.t + 1

        ## potential problems:
        ## - limit number of consecutive pneumostome opening using t
        ## - limit number of consecutive breathing event so that not only do that ?
                ## by using breathing even only when below threshold of O2 ?
                ## or use refractory period ?

        print([self.agent_pos], self.reward, self.done, {'new_trial': new_trial})
        return np.array([self.agent_pos]), self.reward, self.done, {'new_trial': new_trial}

    def render(self,mode='human'):
        pass

    def close(self):
        pass

if __name__ == '__main__':
    import matplotlib.pyplot as plt
    from neurogym.tests import test_run, test_speed
    env = Pneumostomeopening()
    # test_run(env)
    # test_speed(env)
    from neurogym.utils.plotting import plot_env
    from neurogym.utils.test_plotting import test_plot
    # ngym.utils.plotting.run_env(env, num_steps=200)
    plot = plot_env(env, num_steps=100)
    plt.show(plot)
    # test_plot(env)

    # ob = data['ob']
    # ob.shape
    # actions = data['actions']
    #
    # ngym.utils.plotting.fig_(ob, actions, env=env)
    #
    # ob = np.array(ob)
    # actions = np.array(actions)


# -*- coding: utf-8 -*-
# A Survey on Negative Transfer
# https://github.com/chamwen/NT-Benchmark
import argparse
import os
import random
import os.path as osp
import torch as tr
import numpy as np
import utils.network as network

def create_folder(output_dir):

    if not osp.exists(output_dir):
        os.system('mkdir -p ' + output_dir)
    if not osp.exists(output_dir):
        os.mkdir(output_dir)

if __name__ == '__main__':
    seed = 2022
    tr.manual_seed(seed)
    tr.cuda.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    tr.cuda.manual_seed_all(seed)
    tr.backends.cudnn.deterministic = True
    mdl_init_dir = 'outputs/mdl_init/'
    dset_list = ['DomainNet', 'SEED', 'moon']

    ###################################################################################
    # Img data
    args = argparse.Namespace(bottleneck=1024, net='resnet50', layer='wn', classifier='bn')
    args.class_num = 40
    output_dir = osp.join(mdl_init_dir, dset_list[0])
    create_folder(output_dir)

    if args.net[0:3] == 'res':
        netF = network.ResBase(res_name=args.net).cuda()
    elif args.net[0:3] == 'vgg':
        netF = network.VGGBase(vgg_name=args.net).cuda()
    netB = network.feat_bottleneck(type=args.classifier, feature_dim=netF.in_features,
                                   bottleneck_dim=args.bottleneck).cuda()
    netC = network.feat_classifier(type=args.layer, class_num=args.class_num, bottleneck_dim=args.bottleneck).cuda()
    netD_clf = network.feat_classifier(type=args.layer, class_num=2, bottleneck_dim=args.bottleneck).cuda()
    netD_full = network.AdversarialNetwork(args.bottleneck, 2048).cuda()

    tr.save(netF.state_dict(), osp.join(output_dir, "netF.pt"))
    tr.save(netB.state_dict(), osp.join(output_dir, "netB.pt"))
    tr.save(netC.state_dict(), osp.join(output_dir, "netC.pt"))
    tr.save(netD_clf.state_dict(), osp.join(output_dir, "netD_clf.pt"))
    tr.save(netD_full.state_dict(), osp.join(output_dir, "netD_full.pt"))
    # netF.load_state_dict(tr.load(osp.join(output_dir, "netF.pt")))
    print('\nfinished init of DomainNet data...')

    ###################################################################################
    # SEED data
    args = argparse.Namespace(bottleneck=64, backbone='ShallowNet', layer='wn')
    args.input_dim = 310
    args.class_num = 3
    output_dir = osp.join(mdl_init_dir, dset_list[1])
    create_folder(output_dir)

    netF, netC = network.backbone_net(args, args.bottleneck)
    netD_full = network.AdversarialNetwork(args.bottleneck, 20).cuda()
    netD_clf = network.feat_classifier(type=args.layer, class_num=2, bottleneck_dim=args.bottleneck).cuda()

    tr.save(netF.state_dict(), osp.join(output_dir, "netF.pt"))
    tr.save(netC.state_dict(), osp.join(output_dir, "netC.pt"))
    tr.save(netD_full.state_dict(), osp.join(output_dir, "netD_full.pt"))
    tr.save(netD_clf.state_dict(), osp.join(output_dir, "netD_clf.pt"))
    print('\nfinished init of seed data...')

    ###################################################################################
    # Synth data
    args = argparse.Namespace(bottleneck=64, backbone='ShallowNet', layer='wn')
    args.input_dim = 2
    args.class_num = 2
    output_dir = osp.join(mdl_init_dir, dset_list[2])
    create_folder(output_dir)

    netF, netC = network.backbone_net(args, args.bottleneck)
    netD_full = network.AdversarialNetwork(args.bottleneck, 20).cuda()
    netD_clf = network.feat_classifier(type=args.layer, class_num=2, bottleneck_dim=args.bottleneck).cuda()

    tr.save(netF.state_dict(), osp.join(output_dir, "netF.pt"))
    tr.save(netC.state_dict(), osp.join(output_dir, "netC.pt"))
    tr.save(netD_full.state_dict(), osp.join(output_dir, "netD_full.pt"))
    tr.save(netD_clf.state_dict(), osp.join(output_dir, "netD_clf.pt"))
    print('\nfinished init of moon data...')


# -*- coding:utf-8 -*-
# author:平手友梨奈ii
# e-mail:1353593259@qq.com
# datetime:1993/12/01
# filename:cosine_anneal.py
# software: PyCharm

from yolo import YOLO
from PIL import Image
from keras.layers import Input
from keras import backend as K
from utils.utils import letterbox_image
from nets.yolo4 import yolo_body, yolo_eval_weighted_nms
import colorsys
import numpy as np
import os
from keras.models import load_model

from config.configs import CONFIG
from utils.boxes_wbf import weighted_boxes_fusion

import cv2

class MapYolo(YOLO):
    """
    detect test data and log results in txt file
    """

    def __init__(self, **kwargs):
        self.diou_threshold = CONFIG.DETECT.DIOU
        self.nms_method = CONFIG.DETECT.NMS_METHOD
        super(MapYolo, self).__init__(**kwargs)
        self.resolution = CONFIG.DETECT.RESOLUTION

    def generate(self):
        score = CONFIG.DETECT.SCORE

        model_path = os.path.expanduser(self.model_path)
        assert model_path.endswith('.h5'), 'Keras model or weights must be a .h5 file.'

        num_anchors = len(self.anchors)
        num_classes = len(self.class_names)

        try:
            self.yolo_model = load_model(model_path, compile=False)
        except:
            self.yolo_model = yolo_body(Input(shape=(None, None, 3)), num_anchors // 3, num_classes)
            self.yolo_model.load_weights(self.model_path)
        else:
            assert self.yolo_model.layers[-1].output_shape[-1] == \
                   num_anchors / len(self.yolo_model.output) * (num_classes + 5), \
                   'Mismatch between model and given anchor and class sizes'

        print('{} model, anchors, and classes loaded.'.format(model_path))

        # draw bounding boxes
        hsv_tuples = [(x / len(self.class_names), 1., 1.)
                      for x in range(len(self.class_names))]
        self.colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
        self.colors = list(
            map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)),
                self.colors))

        # random color
        np.random.seed(10101)
        np.random.shuffle(self.colors)
        np.random.seed(None)

        self.input_image_shape = K.placeholder(shape=(2,))

        boxes, scores, classes = yolo_eval_weighted_nms(self.yolo_model.output, self.anchors,
                                                        num_classes, self.input_image_shape,
                                                        score_threshold=score)
        return boxes, scores, classes

    def detect_image(self, image_id, image):

        f = open("./mAP/detection-results/" + image_id + ".txt", "w")

        # use letter box to resize the original img
        new_image_size = self.resolution
        boxed_image = letterbox_image(image, new_image_size)
        image_data = np.array(boxed_image, dtype='float32')
        image_data /= 255.
        image_data = np.expand_dims(image_data, 0)  # Add batch dimension.

        # detect img
        out_boxes, out_scores, out_classes = self.sess.run(
            [self.boxes, self.scores, self.classes],
            feed_dict={
                self.yolo_model.input: image_data,
                self.input_image_shape: [image.size[1], image.size[0]],
                # learning_phase = 0
                # so, BN uses historical mean and std
                K.learning_phase(): 0
            })

        res_boxes = []
        res_scores = []
        res_classes = []

        for i in range(len(out_boxes)):
            boxes_class = out_boxes[i]
            scores_class = out_scores[i]
            fusion_boxes = np.concatenate([boxes_class, np.expand_dims(scores_class, axis=1)], axis=-1)
            # because i have filter boxes by score threshold
            # s0, i set score_thres = 0.0 here
            results = weighted_boxes_fusion(boxes=fusion_boxes,
                                            iou_thres=CONFIG.DETECT.IOU,
                                            score_thres=0.0)
            if len(results) > 0:
                res_boxes.append(results[..., :4])
                res_scores.append(results[..., 4])
                res_classes.append(np.ones_like(results[..., 4], dtype=np.int) * i)
            else:
                continue
        if len(res_boxes) > 0:
            res_boxes = np.concatenate(res_boxes, axis=0)
            res_scores = np.concatenate(res_scores, axis=0)
            res_classes = np.concatenate(res_classes, axis=0)

        for i, c in enumerate(res_classes):
            predicted_class = self.class_names[int(c)]
            predicted_class = predicted_class.replace(' ', '_')
            score = str(res_scores[i])

            top, left, bottom, right = res_boxes[i]
            f.write("%s %s %s %s %s %s\n" % (predicted_class, score[:6],
                                             str(int(left)), str(int(top)), str(int(right)), str(int(bottom))))

        f.close()

        img_array = np.asarray(image)
        for i in range(len(res_boxes)):
            boxes = res_boxes[i]
            top, left, bottom, right = boxes
            cv2.rectangle(img_array, (int(left), int(top)), (int(right), int(bottom)), (0, 255, 0), 2)
        cv2.imshow('img', img_array)
        cv2.waitKey(0)
        return

yolo = MapYolo()

image_ids = open('VOCdevkit/VOC2088/ImageSets/Main/test.txt', encoding='utf-8').read().strip().split()

if not os.path.exists("./mAP"):
    os.makedirs("./mAP")
if not os.path.exists("./mAP/detection-results"):
    os.makedirs("./mAP/detection-results")
if not os.path.exists("./mAP/images-optional"):
    os.makedirs("./mAP/images-optional")

for image_id in image_ids:
    image_path = "./VOCdevkit/VOC2088/JPEGImages/" + image_id + ".jpg"
    image = Image.open(image_path)
    image_id = image_id.split('/')[1]
    yolo.detect_image(image_id, image)
    print(image_id, " done!")

print("Conversion completed!")


import json
import os
import os.path as osp
#from collections import OrderedDict
import tempfile

import numpy as np

from mmpose.core.evaluation.top_down_eval import (keypoint_nme,
                                                  keypoint_pck_accuracy)
from mmpose.datasets.builder import DATASETS
from mmpose.datasets.datasets.base import Kpt2dSviewRgbImgTopDownDataset

######

from collections import OrderedDict, defaultdict

import json_tricks as json
import numpy as np
from mmcv import Config, deprecated_api_warning
from xtcocotools.cocoeval import COCOeval

from ....core.post_processing import oks_nms, soft_oks_nms
from ...builder import DATASETS
from ..base import Kpt2dSviewRgbImgTopDownDataset

#####

@DATASETS.register_module()
class TopDownCOCOTinyDataset(Kpt2dSviewRgbImgTopDownDataset):

    def __init__(self,
                 ann_file,
                 img_prefix,
                 data_cfg,
                 pipeline,
                 dataset_info=None,
                 test_mode=False):
        super().__init__(
            ann_file, img_prefix, data_cfg, pipeline, dataset_info, coco_style=False, test_mode=test_mode)

        # flip_pairs, upper_body_ids and lower_body_ids will be used
        # in some data augmentations like random flip
        self.ann_info['flip_pairs'] = [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10],
                                       [11, 12], [13, 14], [15, 16]]
        self.ann_info['upper_body_ids'] = (0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10)
        self.ann_info['lower_body_ids'] = (11, 12, 13, 14, 15, 16)

        self.ann_info['joint_weights'] = None
        self.ann_info['use_different_joint_weights'] = False

        self.dataset_name = 'coco_tiny'
        self.db = self._get_db()

    def _get_db(self):
        with open(self.ann_file) as f:
            anns = json.load(f)

        db = []
        for idx, ann in enumerate(anns):
            # get image path
            image_file = osp.join(self.img_prefix, ann['image_file'])
            # get bbox
            bbox = ann['bbox']
            center, scale = self._xywh2cs(*bbox)
            # get keypoints
            keypoints = np.array(
                ann['keypoints'], dtype=np.float32).reshape(-1, 3)
            num_joints = keypoints.shape[0]
            joints_3d = np.zeros((num_joints, 3), dtype=np.float32)
            joints_3d[:, :2] = keypoints[:, :2]
            joints_3d_visible = np.zeros((num_joints, 3), dtype=np.float32)
            joints_3d_visible[:, :2] = np.minimum(1, keypoints[:, 2:3])

            sample = {
                'image_file': image_file,
                'center': center,
                'scale': scale,
                'bbox': bbox,
                'rotation': 0,
                'joints_3d': joints_3d,
                'joints_3d_visible': joints_3d_visible,
                'bbox_score': 1,
                'bbox_id': idx,
            }
            db.append(sample)

        return db

    def _xywh2cs(self, x, y, w, h):
        """This encodes bbox(x, y, w, h) into (center, scale)
        Args:
            x, y, w, h
        Returns:
            tuple: A tuple containing center and scale.
            - center (np.ndarray[float32](2,)): center of the bbox (x, y).
            - scale (np.ndarray[float32](2,)): scale of the bbox w & h.
        """
        aspect_ratio = self.ann_info['image_size'][0] / self.ann_info[
            'image_size'][1]
        center = np.array([x + w * 0.5, y + h * 0.5], dtype=np.float32)
        if w > aspect_ratio * h:
            h = w * 1.0 / aspect_ratio
        elif w < aspect_ratio * h:
            w = h * aspect_ratio

        # pixel std is 200.0
        scale = np.array([w / 200.0, h / 200.0], dtype=np.float32)
        # padding to include proper amount of context
        scale = scale * 1.25
        return center, scale

    @deprecated_api_warning(name_dict=dict(outputs='results'))
    def evaluate(self, results, res_folder=None, metric='mAP', **kwargs):
        """Evaluate coco keypoint results. The pose prediction results will be
        saved in ``${res_folder}/result_keypoints.json``.

        Note:
            - batch_size: N
            - num_keypoints: K
            - heatmap height: H
            - heatmap width: W

        Args:
            results (list[dict]): Testing results containing the following
                items:

                - preds (np.ndarray[N,K,3]): The first two dimensions are \
                    coordinates, score is the third dimension of the array.
                - boxes (np.ndarray[N,6]): [center[0], center[1], scale[0], \
                    scale[1],area, score]
                - image_paths (list[str]): For example, ['data/coco/val2017\
                    /000000393226.jpg']
                - heatmap (np.ndarray[N, K, H, W]): model output heatmap
                - bbox_id (list(int)).
            res_folder (str, optional): The folder to save the testing
                results. If not specified, a temp folder will be created.
                Default: None.
            metric (str | list[str]): Metric to be performed. Defaults: 'mAP'.

        Returns:
            dict: Evaluation results for evaluation metric.
        """
        metrics = metric if isinstance(metric, list) else [metric]
        allowed_metrics = ['mAP']
        for metric in metrics:
            if metric not in allowed_metrics:
                raise KeyError(f'metric {metric} is not supported')

        if res_folder is not None:
            tmp_folder = None
            res_file = osp.join(res_folder, 'result_keypoints.json')
        else:
            tmp_folder = tempfile.TemporaryDirectory()
            res_file = osp.join(tmp_folder.name, 'result_keypoints.json')

        kpts = defaultdict(list)

        for result in results:
            preds = result['preds']
            boxes = result['boxes']
            image_paths = result['image_paths']
            bbox_ids = result['bbox_ids']

            batch_size = len(image_paths)
            for i in range(batch_size):
                image_id = self.name2id[image_paths[i][len(self.img_prefix):]]
                kpts[image_id].append({
                    'keypoints': preds[i],
                    'center': boxes[i][0:2],
                    'scale': boxes[i][2:4],
                    'area': boxes[i][4],
                    'score': boxes[i][5],
                    'image_id': image_id,
                    'bbox_id': bbox_ids[i]
                })
        kpts = self._sort_and_unique_bboxes(kpts)

        # rescoring and oks nms
        num_joints = self.ann_info['num_joints']
        vis_thr = self.vis_thr
        oks_thr = self.oks_thr
        valid_kpts = []
        for image_id in kpts.keys():
            img_kpts = kpts[image_id]
            for n_p in img_kpts:
                box_score = n_p['score']
                kpt_score = 0
                valid_num = 0
                for n_jt in range(0, num_joints):
                    t_s = n_p['keypoints'][n_jt][2]
                    if t_s > vis_thr:
                        kpt_score = kpt_score + t_s
                        valid_num = valid_num + 1
                if valid_num != 0:
                    kpt_score = kpt_score / valid_num
                # rescoring
                n_p['score'] = kpt_score * box_score

            if self.use_nms:
                nms = soft_oks_nms if self.soft_nms else oks_nms
                keep = nms(img_kpts, oks_thr, sigmas=self.sigmas)
                valid_kpts.append([img_kpts[_keep] for _keep in keep])
            else:
                valid_kpts.append(img_kpts)

        self._write_coco_keypoint_results(valid_kpts, res_file)

        info_str = self._do_python_keypoint_eval(res_file)
        name_value = OrderedDict(info_str)

        if tmp_folder is not None:
            tmp_folder.cleanup()

        return name_value

    def _write_coco_keypoint_results(self, keypoints, res_file):
        """Write results into a json file."""
        data_pack = [{
            'cat_id': self._class_to_coco_ind[cls],
            'cls_ind': cls_ind,
            'cls': cls,
            'ann_type': 'keypoints',
            'keypoints': keypoints
        } for cls_ind, cls in enumerate(self.classes)
                     if not cls == '__background__']

        results = self._coco_keypoint_results_one_category_kernel(data_pack[0])

        with open(res_file, 'w') as f:
            json.dump(results, f, sort_keys=True, indent=4)

    def _coco_keypoint_results_one_category_kernel(self, data_pack):
        """Get coco keypoint results."""
        cat_id = data_pack['cat_id']
        keypoints = data_pack['keypoints']
        cat_results = []

        for img_kpts in keypoints:
            if len(img_kpts) == 0:
                continue

            _key_points = np.array(
                [img_kpt['keypoints'] for img_kpt in img_kpts])
            key_points = _key_points.reshape(-1,
                                             self.ann_info['num_joints'] * 3)

            result = [{
                'image_id': img_kpt['image_id'],
                'category_id': cat_id,
                'keypoints': key_point.tolist(),
                'score': float(img_kpt['score']),
                'center': img_kpt['center'].tolist(),
                'scale': img_kpt['scale'].tolist()
            } for img_kpt, key_point in zip(img_kpts, key_points)]

            cat_results.extend(result)

        return cat_results

    def _do_python_keypoint_eval(self, res_file):
        """Keypoint evaluation using COCOAPI."""
        coco_det = self.coco.loadRes(res_file)
        coco_eval = COCOeval(self.coco, coco_det, 'keypoints', self.sigmas)
        coco_eval.params.useSegm = None
        coco_eval.evaluate()
        coco_eval.accumulate()
        coco_eval.summarize()

        stats_names = [
            'AP', 'AP .5', 'AP .75', 'AP (M)', 'AP (L)', 'AR', 'AR .5',
            'AR .75', 'AR (M)', 'AR (L)'
        ]

        info_str = list(zip(stats_names, coco_eval.stats))

        return info_str

    def _sort_and_unique_bboxes(self, kpts, key='bbox_id'):
        """sort kpts and remove the repeated ones."""
        for img_id, persons in kpts.items():
            num = len(persons)
            kpts[img_id] = sorted(kpts[img_id], key=lambda x: x[key])
            for i in range(num - 1, 0, -1):
                if kpts[img_id][i][key] == kpts[img_id][i - 1][key]:
                    del kpts[img_id][i]

        return kpts

    def evaluate1111(self, results, res_folder=None, metric='PCK', **kwargs):
#    def evaluate(self, results, res_folder=None, metric='PCK', **kwargs):
        """Evaluate keypoint detection results. The pose prediction results will
        be saved in `${res_folder}/result_keypoints.json`.

        Note:
        batch_size: N
        num_keypoints: K
        heatmap height: H
        heatmap width: W

        Args:
        results (list(preds, boxes, image_path, output_heatmap))
            :preds (np.ndarray[N,K,3]): The first two dimensions are
                coordinates, score is the third dimension of the array.
            :boxes (np.ndarray[N,6]): [center[0], center[1], scale[0]
                , scale[1],area, score]
            :image_paths (list[str]): For example, ['Test/source/0.jpg']
            :output_heatmap (np.ndarray[N, K, H, W]): model outputs.

        res_folder (str, optional): The folder to save the testing
                results. If not specified, a temp folder will be created.
                Default: None.
        metric (str | list[str]): Metric to be performed.
            Options: 'PCK', 'NME'.

        Returns:
            dict: Evaluation results for evaluation metric.
        """
        metrics = metric if isinstance(metric, list) else [metric]
        allowed_metrics = ['PCK', 'NME']
        for metric in metrics:
            if metric not in allowed_metrics:
                raise KeyError(f'metric {metric} is not supported')

        if res_folder is not None:
            tmp_folder = None
            res_file = osp.join(res_folder, 'result_keypoints.json')
        else:
            tmp_folder = tempfile.TemporaryDirectory()
            res_file = osp.join(tmp_folder.name, 'result_keypoints.json')

        kpts = []
        for result in results:
            preds = result['preds']
            boxes = result['boxes']
            image_paths = result['image_paths']
            bbox_ids = result['bbox_ids']

            batch_size = len(image_paths)
            for i in range(batch_size):
                kpts.append({
                    'keypoints': preds[i].tolist(),
                    'center': boxes[i][0:2].tolist(),
                    'scale': boxes[i][2:4].tolist(),
                    'area': float(boxes[i][4]),
                    'score': float(boxes[i][5]),
                    'bbox_id': bbox_ids[i]
                })
        kpts = self._sort_and_unique_bboxes(kpts)

        self._write_keypoint_results(kpts, res_file)
        info_str = self._report_metric(res_file, metrics)
        name_value = OrderedDict(info_str)

        if tmp_folder is not None:
            tmp_folder.cleanup()

        return name_value

    def _report_metric(self, res_file, metrics, pck_thr=0.3):
        """Keypoint evaluation.

        Args:
        res_file (str): Json file stored prediction results.
        metrics (str | list[str]): Metric to be performed.
            Options: 'PCK', 'NME'.
        pck_thr (float): PCK threshold, default: 0.3.

        Returns:
        dict: Evaluation results for evaluation metric.
        """
        info_str = []

        with open(res_file, 'r') as fin:
            preds = json.load(fin)
        assert len(preds) == len(self.db)

        outputs = []
        gts = []
        masks = []

        for pred, item in zip(preds, self.db):
            outputs.append(np.array(pred['keypoints'])[:, :-1])
            gts.append(np.array(item['joints_3d'])[:, :-1])
            masks.append((np.array(item['joints_3d_visible'])[:, 0]) > 0)

        outputs = np.array(outputs)
        gts = np.array(gts)
        masks = np.array(masks)

        normalize_factor = self._get_normalize_factor(gts)

        if 'PCK' in metrics:
            _, pck, _ = keypoint_pck_accuracy(outputs, gts, masks, pck_thr,
                                              normalize_factor)
            info_str.append(('PCK', pck))

        if 'NME' in metrics:
            info_str.append(
                ('NME', keypoint_nme(outputs, gts, masks, normalize_factor)))

        return info_str

    @staticmethod
    def _write_keypoint_results(keypoints, res_file):
        """Write results into a json file."""

        with open(res_file, 'w') as f:
            json.dump(keypoints, f, sort_keys=True, indent=4)

    @staticmethod
    def _sort_and_unique_bboxes(kpts, key='bbox_id'):
        """sort kpts and remove the repeated ones."""
        kpts = sorted(kpts, key=lambda x: x[key])
        num = len(kpts)
        for i in range(num - 1, 0, -1):
            if kpts[i][key] == kpts[i - 1][key]:
                del kpts[i]

        return kpts

    @staticmethod
    def _get_normalize_factor(gts):
        """Get inter-ocular distance as the normalize factor, measured as the
        Euclidean distance between the outer corners of the eyes.

        Args:
            gts (np.ndarray[N, K, 2]): Groundtruth keypoint location.

        Return:
            np.ndarray[N, 2]: normalized factor
        """

        interocular = np.linalg.norm(
            gts[:, 0, :] - gts[:, 1, :], axis=1, keepdims=True)
        return np.tile(interocular, [1, 2])


import numpy as np
import torch
import argparse
from tqdm import tqdm
import os.path as osp
from data import Data
from torch.utils.data import DataLoader
from model import ModelDSR
import itertools

parser = argparse.ArgumentParser()

parser.add_argument('--resume', type=str, help='path to model')
parser.add_argument('--data_path', type=str, help='path to data')
parser.add_argument('--test_type', type=str, choices=['motion_visible', 'motion_full', 'mask_ordered', 'mask_unordered'])

parser.add_argument('--gpu', type=int, default=0, help='gpu id (single gpu)')
parser.add_argument('--object_num', type=int, default=5, help='number of objects')
parser.add_argument('--seq_len', type=int, default=10, help='sequence length')
parser.add_argument('--batch', type=int, default=12, help='batch size')
parser.add_argument('--workers', type=int, default=2, help='number of workers in data loader')

parser.add_argument('--model_type', type=str, default='dsr', choices=['dsr', 'single', 'nowarp', 'gtwarp', '3dflow'])
parser.add_argument('--transform_type', type=str, default='se3euler', choices=['affine', 'se3euler', 'se3aa', 'se3spquat', 'se3quat'])

def main():
    args = parser.parse_args()
    torch.cuda.set_device(args.gpu)

    data, loaders = {}, {}
    for split in ['test']:
        data[split] = Data(data_path=args.data_path, split=split, seq_len=args.seq_len)
        loaders[split] = DataLoader(dataset=data[split], batch_size=args.batch, num_workers=args.workers)
    print('==> dataset loaded: [size] = {0}'.format(len(data['test'])))

    model = ModelDSR(
        object_num=args.object_num,
        transform_type=args.transform_type,
        motion_type='se3' if args.model_type != '3dflow' else 'conv',
    )
    model.cuda()

    checkpoint = torch.load(args.resume, map_location=torch.device(f'cuda:{args.gpu}'))
    model.load_state_dict(checkpoint['state_dict'])
    print('==> resume: ' + args.resume)

    with torch.no_grad():
        if args.test_type == 'motion_visible':
            evaluation_motion_visible(args, model, loaders['test'])

        if args.test_type == 'motion_full':
            evaluation_motion_full(args, model, loaders['test'])

        if args.test_type == 'mask_ordered':
            evaluation_mask_ordered(args, model, loaders['test'])

        if args.test_type == 'mask_unordered':
            evaluation_mask_unordered(args, model, loaders['test'])

def evaluation_mask_unordered(args, model, loader):
    print(f'==> evaluation_mask (unordered)')
    iou_dict = [[] for _ in range(args.seq_len)]
    for batch in tqdm(loader):
        batch_size = batch['0-action'].size(0)
        last_s = model.get_init_repr(batch_size).cuda()
        logit_pred_list, mask_gt_list = [], []
        for step_id in range(args.seq_len):
            output = model(
                input_volume=batch['%d-tsdf' % step_id].cuda().unsqueeze(1),
                last_s=last_s,
                input_action=batch['%d-action' % step_id].cuda(),
                input_motion=batch['%d-scene_flow_3d' % step_id].cuda() if args.model_type=='gtwarp' else None,
                no_warp=args.model_type=='nowarp'
            )
            if not args.model_type == 'single':
                last_s = output['s'].data

            logit_pred = output['init_logit']
            mask_gt = batch['%d-mask_3d' % step_id].cuda()
            iou_unordered = calc_iou_unordered(logit_pred, mask_gt)
            iou_dict[step_id].append(iou_unordered)
    print('mask_unordered (IoU) = ', np.mean([np.mean(np.concatenate(iou_dict[i])) for i in range(args.seq_len)]))

def calc_iou_unordered(logit_pred, mask_gt_argmax):
    # logit_pred: [B, K, S1, S2, S3], softmax, the last channel is empty
    # mask_gt_argmax: [B, S1, S2, S3], 0 represents empty
    B, K, S1, S2, S3 = logit_pred.size()
    logit_pred_argmax = torch.argmax(logit_pred, dim=1, keepdim=True)
    mask_gt_argmax = torch.unsqueeze(mask_gt_argmax, 1)
    mask_pred_onehot = torch.zeros_like(logit_pred).scatter(1, logit_pred_argmax, 1)[:, :-1]
    mask_gt_onehot = torch.zeros_like(logit_pred).scatter(1, mask_gt_argmax, 1)[:, 1:]
    K -= 1
    info_dict = {'I': np.zeros([B, K, K]), 'U': np.zeros([B, K, K])}
    for b in range(B):
        for i in range(K):
            for j in range(K):
                mask_gt = mask_gt_onehot[b, i]
                mask_pred = mask_pred_onehot[b, j]
                I = torch.sum(mask_gt * mask_pred).item()
                U = torch.sum(mask_gt + mask_pred).item() - I
                info_dict['I'][b, i, j] = I
                info_dict['U'][b, i, j] = U
    batch_ious = []
    for b in range(B):
        best_iou, best_p = 0, None
        for p in list(itertools.permutations(range(K))):
            cur_I = [info_dict['I'][b, i, p[i]] for i in range(K)]
            cur_U = [info_dict['U'][b, i, p[i]] for i in range(K)]
            cur_iou = np.mean(np.array(cur_I) / np.maximum(np.array(cur_U), 1))
            if cur_iou > best_iou:
                best_iou = cur_iou
        batch_ious.append(best_iou)

    return np.array(batch_ious)

def evaluation_mask_ordered(args, model, loader):
    print(f'==> evaluation_mask (ordered)')
    iou_dict = []
    for batch in tqdm(loader):
        batch_size = batch['0-action'].size(0)
        last_s = model.get_init_repr(batch_size).cuda()
        logit_pred_list, mask_gt_list = [], []
        for step_id in range(args.seq_len):
            output = model(
                input_volume=batch['%d-tsdf' % step_id].cuda().unsqueeze(1),
                last_s=last_s,
                input_action=batch['%d-action' % step_id].cuda(),
                input_motion=batch['%d-scene_flow_3d' % step_id].cuda() if args.model_type=='gtwarp' else None,
                no_warp=args.model_type=='nowarp'
            )
            if not args.model_type == 'single':
                last_s = output['s'].data

            logit_pred = output['init_logit']
            mask_gt = batch['%d-mask_3d' % step_id].cuda()
            logit_pred_list.append(logit_pred)
            mask_gt_list.append(mask_gt)
        iou_ordered = calc_iou_ordered(logit_pred_list, mask_gt_list)
        iou_dict.append(iou_ordered)
    print('mask_ordered (IoU) = ', np.mean(np.concatenate(iou_dict)))

def calc_iou_ordered(logit_pred_list, mask_gt_argmax_list):
    # logit_pred_list: [L, B, K, S1, S2, S3], softmax, the last channel is empty
    # mask_gt_argmax_list: [L, B, S1, S2, S3], 0 represents empty
    L = len(logit_pred_list)
    B, K, S1, S2, S3 = logit_pred_list[0].size()
    K -= 1
    info_dict = {'I': np.zeros([L, B, K, K]), 'U': np.zeros([L, B, K, K])}
    for l in range(L):
        logit_pred = logit_pred_list[l]
        mask_gt_argmax = mask_gt_argmax_list[l]
        logit_pred_argmax = torch.argmax(logit_pred, dim=1, keepdim=True)
        mask_gt_argmax = torch.unsqueeze(mask_gt_argmax, 1)
        mask_pred_onehot = torch.zeros_like(logit_pred).scatter(1, logit_pred_argmax, 1)[:, :-1]
        mask_gt_onehot = torch.zeros_like(logit_pred).scatter(1, mask_gt_argmax, 1)[:, 1:]
        for b in range(B):
            for i in range(K):
                for j in range(K):
                    mask_gt = mask_gt_onehot[b, i]
                    mask_pred = mask_pred_onehot[b, j]
                    I = torch.sum(mask_gt * mask_pred).item()
                    U = torch.sum(mask_gt + mask_pred).item() - I
                    info_dict['I'][l, b, i, j] = I
                    info_dict['U'][l, b, i, j] = U
    batch_ious = []
    for b in range(B):
        best_iou, best_p = 0, None
        for p in list(itertools.permutations(range(K))):
            cur_I = [info_dict['I'][l, b, i, p[i]] for l in range(L) for i in range(K)]
            cur_U = [info_dict['U'][l, b, i, p[i]] for l in range(L) for i in range(K)]
            cur_iou = np.mean(np.array(cur_I) / np.maximum(np.array(cur_U), 1))
            if cur_iou > best_iou:
                best_iou = cur_iou
        batch_ious.append(best_iou)

    return np.array(batch_ious)

def evaluation_motion_visible(args, model, loader):
    print('==> evaluation_motion (visible surface)')
    mse_dict = [0 for _ in range(args.seq_len)]
    data_num = 0
    for batch in tqdm(loader):
        batch_size = batch['0-action'].size(0)
        data_num += batch_size
        last_s = model.get_init_repr(batch_size).cuda()
        for step_id in range(args.seq_len):
            output = model(
                input_volume=batch['%d-tsdf' % step_id].cuda().unsqueeze(1),
                last_s=last_s,
                input_action=batch['%d-action' % step_id].cuda(),
                input_motion=batch['%d-scene_flow_3d' % step_id].cuda() if args.model_type=='gtwarp' else None,
                no_warp=args.model_type=='nowarp'
            )
            if not args.model_type in ['single', '3dflow'] :
                last_s = output['s'].data

            tsdf = batch['%d-tsdf' % step_id].cuda().unsqueeze(1)
            mask = batch['%d-mask_3d' % step_id].cuda().unsqueeze(1)
            surface_mask = ((tsdf > -0.99).float()) * ((tsdf < 0).float()) * ((mask > 0).float())
            surface_mask[..., 0] = 0

            target = batch['%d-scene_flow_3d' % step_id].cuda()
            pred = output['motion']

            mse = torch.sum((target - pred) ** 2 * surface_mask, dim=[1, 2, 3, 4]) / torch.sum(surface_mask, dim=[1, 2, 3, 4])
            mse_dict[step_id] += torch.sum(mse).item() * 0.16
            # 0.16(0.4^2) is the scale to convert the unit from "voxel" to "cm".
            # The voxel size is 0.4cm. Here we use seuqre error.
    print('motion_visible (MSE in cm) = ', np.mean([np.mean(mse_dict[i]) / data_num for i in range(args.seq_len)]))

def evaluation_motion_full(args, model, loader):
    print('==> evaluation_motion (full volume)')
    mse_dict = [0 for _ in range(args.seq_len)]
    data_num = 0
    for batch in tqdm(loader):
        batch_size = batch['0-action'].size(0)
        data_num += batch_size
        last_s = model.get_init_repr(batch_size).cuda()
        for step_id in range(args.seq_len):
            output = model(
                input_volume=batch['%d-tsdf' % step_id].cuda().unsqueeze(1),
                last_s=last_s,
                input_action=batch['%d-action' % step_id].cuda(),
                input_motion=batch['%d-scene_flow_3d' % step_id].cuda() if args.model_type=='gtwarp' else None,
                no_warp=args.model_type=='nowarp'
            )
            if not args.model_type in ['single', '3dflow'] :
                last_s = output['s'].data

            target = batch['%d-scene_flow_3d' % step_id].cuda()
            pred = output['motion']

            mse = torch.mean((target - pred) ** 2, dim=[1, 2, 3, 4])
            mse_dict[step_id] += torch.sum(mse).item() * 0.16
            # 0.16(0.4^2) is the scale to convert the unit from "voxel" to "cm".
            # The voxel size is 0.4cm. Here we use seuqre error.
    print('motion_full (MSE in cm) = ', np.mean([np.mean(mse_dict[i]) / data_num for i in range(args.seq_len)]))

if __name__ == '__main__':
    main()


import unittest
import numpy as np

try:
    import SparseSC as SC
except ImportError:
    raise RuntimeError("SparseSC is not installed. Use 'pip install -e .' or 'conda develop .' from repo root to install in dev mode")
from SparseSC.fit import fit

#For the synth comparisons you will need rpy2 and the Synth package loaded in there.
# to do the latter:
# python
# from test.test_simulation import installRSynth
# installRSynth()
# installRCausalImpact()

def installRSynth():
    from rpy2.robjects.packages import importr
    utils = importr('utils')
    utils.install_packages('Synth')

def installRCausalImpact():
    from rpy2.robjects.packages import importr
    utils = importr('utils')
    utils.install_packages('CausalImpact')

class Simulation(unittest.TestCase):
    @staticmethod
    def fitRCausalImpact(Y_pre, Y_post, treated_units):
        import rpy2.robjects as ro #ro.r is the R instace
        #Automatically convert numpy arrays to R vector/matrix
        from rpy2.robjects import numpy2ri
        numpy2ri.activate()

        control_units = [u for u in range(Y_pre.shape[0]) if u not in treated_units]

        #np.matrix is not automatically converted so use ndarray
        if type(Y_pre).__name__=="matrix":
            Y_pre = Y_pre.A
        if type(Y_post).__name__=="matrix":
            Y_post = Y_post.A

        Y = np.hstack((Y_pre, Y_post))
        Y_sc = np.full(Y.shape, np.nan)
        Y_c = Y[control_units,:]
        T0 = Y_pre.shape[1]
        N,T = Y.shape

        try:
            CausalImpact = importr('CausalImpact')
        except:
            raise RuntimeError("Need the 'CausalImpact' package loaded in the rpy2 R environment. Use test.test_simulation.installRCausalImpact")

        for unit in range(N): #
            if unit in treated_units:
                data = np.hstack((Y[unit,:].T, Y_c.T))
            else:
                data = np.hstack((Y[unit,:].T, np.delete(Y_c, unit, 0).T))
            r_data = ro.r.matrix(data, nrow=df.shape[0], ncol=df.shape[1])

            r_casaulimpact_out = CausalImpact.CausalImpact(data=r_data, pre_period=ro.IntVector([1, T0]), post_period=ro.IntVector([T0+1, T]))
            #can't seem to get weights from impact$model$bsts.model
            #r_summary = r_casaulimpact_out[r_casaulimpact_out.names.index('summary')]
            #te = r_summary[r_summary.names.index('AbsEffect')]
            #p = r_summary[r_summary.names.index('p')]
            #te_ci_l = r_summary[r_summary.names.index('AbsEffect.lower')]
            #te_ci_u = r_summary[r_summary.names.index('AbsEffect.upper')]
            point_pred = np.array(r_casaulimpact_out[r_casaulimpact_out.names.index('series')])[:,2] #'point.pred' from 'zoo' object
            Y_sc[unit,:] = point_pred

        class RCausalImpact:
            def __self__(self, Y_sc):
                self.Y_sc = Y_sc

        return RCausalImpact(Y_sc)

    @staticmethod
    def fitRSynth(Y_pre, treated_units, verbose=False):
        #See the Synth.PDF from the CRAN package
        import rpy2.robjects as ro #ro.r is the R instace
        assert len(treated_units)==1 #for now
        control_units = [u for u in range(Y_pre.shape[0]) if u not in treated_units]

        #np.matrix is not automatically converted so use ndarray
        if type(Y_pre).__name__=="matrix":
            Y_pre = Y_pre.A

        N1, N0 = len(treated_units), len(control_units)
        N = N1 + N0
        T0 = Y_pre.shape[1]

        class RSynth:
            def __init__(self, sc_weights, Vs, control_units):
                self.sc_weights = sc_weights #NxN0
                self.Vs = Vs #NxT0
                self.control_units = control_units

            def predict(self, mat_c):
                #mat_c: N0xk
                return np.dot(sc_weights, mat_c[control_units,:])

        #Automatically convert numpy arrays to R vector/matrix
        from rpy2.robjects import numpy2ri
        numpy2ri.activate()

        from rpy2.robjects.packages import importr
        try:
            Synth = importr('Synth')
        except:
            raise RuntimeError("Need the 'Synth' package loaded in the rpy2 R environment. Use test.test_simulation.installRSynth")

        sc_weights = np.zeros((N,N0))
        Vs = np.zeros((N,T0))

        #Difficult to suppress Synth output. The following doesn't work
        #with capture() as out: #https://stackoverflow.com/questions/5136611/ #didn't suppress
        #r_synth_out = utils.capture_output(Synth.synth(X1=r_X1, X0=r_X0, Z1=r_X1, Z0=r_X0), file=tc) #API error
        if not verbose:
            r_tc = ro.r.textConnection("messages","w")
            ro.r.sink(r_tc)

        X0=Y_pre[control_units, :].T
        r_X0 = ro.r.matrix(X0, nrow=X0.shape[0], ncol=X0.shape[1])
        for treated_unit in treated_units:
            X1=Y_pre[[treated_unit], :].T
            r_X1 = ro.r.matrix(X1, nrow=X1.shape[0], ncol=X1.shape[1])

            r_synth_out = Synth.synth(X1=r_X1, X0=r_X0, Z1=r_X1, Z0=r_X0)
            sc_weights[treated_unit,:] = np.squeeze(np.array(r_synth_out[r_synth_out.names.index('solution.w')]))
            Vs[treated_unit,:] = np.squeeze(np.array(r_synth_out[r_synth_out.names.index('solution.v')]))

        from tqdm import tqdm
        for p_c, p_ct in enumerate(tqdm(control_units)):
            donors_ct = control_units.copy()
            donors_ct.remove(p_ct)
            donors_c = list(range(N0))
            donors_c.remove(p_c)
            X1=Y_pre[[p_ct], :].T
            X0=Y_pre[donors_ct, :].T
            r_X1 = ro.r.matrix(X1, nrow=X1.shape[0], ncol=X1.shape[1])
            r_X0 = ro.r.matrix(X0, nrow=X0.shape[0], ncol=X0.shape[1])

            with capture() as out: #doesn't have quiet option
                r_synth_out = Synth.synth(X1=r_X1, X0=r_X0, Z1=r_X1, Z0=r_X0)
            sc_weights[p_ct,donors_c] = np.squeeze(np.array(r_synth_out[r_synth_out.names.index('solution.w')]))
            Vs[p_ct,:] = np.squeeze(np.array(r_synth_out[r_synth_out.names.index('solution.v')]))

        if not verbose:
            ro.r.sink()
            ro.r.close(r_tc)

        return RSynth(sc_weights, Vs, control_units)

    @staticmethod
    def dataFactorDGP_AA(N0s=[20], T0s=[10], S=1):
        import itertools
        from dgp.factor_model import factor_dgp

        N1,T1 = 1,1
        data_dict = {}
        for counter, (N0,T0) in enumerate(itertools.product(N0s, T0s)):
            data_dict[(N0,T0)] = []
            for s in range(S):
                _, _, Y_pre_C, Y_pre_T, Y_post_C, Y_post_T, l_C, l_T = factor_dgp(N0, N1, T0, T1, K=0, R=0, F=1)
                Y_pre = np.vstack((Y_pre_T, Y_pre_C))
                Y_post = np.vstack((Y_post_T, Y_post_C))
                l = np.vstack((l_T, l_C))
                data_dict[(N0,T0)].append((Y_pre, Y_post, l))
        return(data_dict)

    @staticmethod
    def fitSparseSC_wrapper(Y_pre, Y_post, treated_units):
        return(fit(X=Y_pre, Y=Y_post, treated_units=treated_units,
                              model_type="retrospective", constrain="simplex",
                              print_path = False, progress = False, verbose=0))

    @staticmethod
    def fitRSynth_wrapper(Y_pre, Y_post, treated_units):
        return(Simulation.fitRSynth(Y_pre, treated_units))

    @staticmethod
    def SCFactor_AA_runner(data_dict, sc_method):
        treated_units = [0]
        data_keys = data_dict.keys()
        results = np.zeros((len(data_keys), 2))
        for counter, (N0,T0) in enumerate(data_keys):
            print(counter)
            datasets = data_dict[(N0,T0)]
            S = len(datasets)
            comb_results = np.empty((S,2))
            for s in range(S):
                Y_pre, Y_post, l = datasets[s]
                fit_res = sc_method(Y_pre, Y_post, treated_units)
                comb_results[s,0] = np.mean(np.square(Y_post-fit_res.predict(Y_post)))
                comb_results[s,1] = np.mean(np.square(l-fit_res.predict(l)))
            results[counter,:] = np.mean(comb_results, axis=1)
        return(results)

    def testFactorDGP_AA(self, N0s=[20], T0s=[10], S=1):
        import pickle
        import itertools

        treated_units = [0]

        N_combos = len(N0s)*len(T0s)
        data_dic = Simulation.dataFactorDGP_AA(N0s=[20], T0s=[10], S=1)

        data_pkl_file = 'sim_data.pkl'
        with open(data_pkl_file, 'wb') as output:
            pickle.dump(data_dic, output)
        #with open(data_pkl_file, 'rb') as input:
        #    data_dic = pickle.load(input)

        results_ssc = Simulation.SCFactor_AA_runner(data_dic, Simulation.fitSparseSC_wrapper)
        results = results_ssc
        #results_sc = Simulation.SCFactor_AA_runner(data_dic, Simulation.fitRSynth_wrapper)
        #results = np.hstack((results_sc, results_ssc))

        print(results)
        res_pkl_file = 'sim_results.pkl'
        with open(res_pkl_file, 'wb') as output:
            pickle.dump(results, output)
        #with open(pkl_file, 'rb') as input:
        #    results = pickle.load(input)

    # Come up with examples targetting Synth's weaknesses (show it needs lots of data to get to where SparseSC can with little)
    # - Synth over fits by fitting a V for each or not penalizing V
    # Come up with example where

    ##later) Can we match a longer set of factor loadings
    # Other Counterfactual prediction:
    ## a) Compare to SC (big N0, small T0, then SC; or many factors; should do bad) to basic time-series model

if __name__ == "__main__":
    import random

    random.seed(12345)
    np.random.seed(10101)

    s = Simulation()
    s.testFactorDGP_AA()
    # unittest.main()


"""
Code to compute and aggregate hashcodes over a large collection
of data, e.g. the whole million song dataset, using multiple
processes.

Many parameters are hard-coded!

For the help menu, simply launch the code:
    python compute_hashcodes_mprocess.py

Copyright 2011, Thierry Bertin-Mahieux <tb2332@columbia.edu>

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program. If not, see <http://www.gnu.org/licenses/>.
"""

import os
import sys
import glob
import time
import copy
import sqlite3
import datetime
import multiprocessing
import numpy as np

# hash codes params
WIN=3
DECAY=.995
NORMALIZE=True
WEIGHT_MARGIN=.60
MAX_PER_FRAME=3
LEVELS = [3]
COMPRESSION=2

def get_all_files(basedir,ext='.h5') :
    """
    From a root directory, go through all subdirectories
    and find all files with the given extension.
    Return all absolute paths in a list.
    """
    allfiles = set()
    for root, dirs, files in os.walk(basedir):
        files = glob.glob(os.path.join(root,'*'+ext))
        for f in files :
            allfiles.add( os.path.abspath(f) )
    return list(allfiles)

def aggregate_dbs(conn, tmpdbpath):
    """
    Aggregate the data from a filled, temporary database,
    to a main one already initialized.
    """
    try:
        CHT.__name__
    except NameError:
        #print 'We import cover hash table'
        import cover_hash_table as CHT
    t1 = time.time()
    conn_tmp = sqlite3.connect(tmpdbpath)
    # get number of hashes tables
    q = "SELECT cnt FROM num_hash_tables LIMIT 1"
    res = conn_tmp.execute(q)
    num_tables = res.fetchone()[0]
    # get all tids
    q = "SELECT ROWID, tid FROM tids"
    res = conn_tmp.execute(q)
    tidid_tid_pairs = res.fetchall()
    # iterate over all these tids
    for tidid, tid in tidid_tid_pairs:
        # iterate over all tables
        for tablecnt in range(1, num_tables + 1):
            table_name = 'hashes' + str(tablecnt)
            q = "SELECT jumpcode, weight FROM " + table_name
            q += " WHERE tidid=" + str(tidid)
            res = conn_tmp.execute(q)
            jcode_weight_pairs = res.fetchall()
            tid_is_in = False
            for jcode, weight in jcode_weight_pairs:
                CHT.add_hash(conn, tid, jcode, weight, tid_is_in=tid_is_in)
                tid_is_in = True
        # commit
        conn.commit()
    # done
    conn_tmp.close()
    # verbose
    timestr = str(datetime.timedelta(seconds=time.time()-t1))
    print 'Aggregated %s in %s.' % (tmpdbpath, timestr)

def my_vacuum(dbpathnew, dbpathold):
    """
    My own vacuum function on cover_hash_table.
    It works by copying and is slow!
    My main use, transform the page size to 4096
    Here because of its use of 'aggregate_dbs'
    """
    if os.path.exists(dbpathnew):
        print 'ERROR: %s already exists.' % dbpathnew
        return
    if not os.path.isfile(dbpathold):
        print 'ERROR: %s is not a file.' % dbpathold
        return
    # create new db
    import cover_hash_table as CHT
    conn = sqlite3.connect(dbpathnew)
    conn.execute('PRAGMA temp_store = MEMORY;')
    conn.execute('PRAGMA synchronous = OFF;')
    conn.execute('PRAGMA journal_mode = OFF;') # no ROLLBACK!
    conn.execute('PRAGMA page_size = 4096;')
    conn.execute('PRAGMA cache_size = 500000;') # page_size=4096, 500000->2GB
    CHT.init_db(conn)
    # copy
    aggregate_dbs(conn, dbpathold)
    # reindex
    CHT.reindex(conn)
    # done
    conn.commit()
    conn.close()

def create_fill_one_partial_db(filelist=None, outputdb=None):
    """
    This is the main function called by each process
    """
    # assert we have the params
    assert (not filelist is None) and (not outputdb is None), "internal arg passing error...!"
    # must be imported there... maybe... because of local num_hash_tables count
    import cover_hash_table as CHT
    # other imports
    import quick_query_test as QQT # should be replaced in the future
    import fingerprint_hash as FH
    # create output db, including PRAGMA
    conn = sqlite3.connect(outputdb)
    conn.execute('PRAGMA temp_store = MEMORY;')
    conn.execute('PRAGMA synchronous = OFF;')
    conn.execute('PRAGMA journal_mode = OFF;') # no ROLLBACK!
    conn.execute('PRAGMA cache_size = 1000000;') # default=2000, page_size=1024
    CHT.init_db(conn)
    # iterate over files
    cnt_tid_added = 0
    for filepath in filelist:
        # get bthcroma
        btchroma = QQT.get_cpressed_btchroma(filepath, compression=COMPRESSION)
        if btchroma is None:
            continue
        # get tid from filepath (faster than querying h5 file, less robust)
        tid = os.path.split(os.path.splitext(filepath)[0])[1]
        # get jumps
        landmarks = FH.get_landmarks(btchroma, decay=DECAY, max_per_frame=MAX_PER_FRAME)
        jumps = FH.get_jumps(landmarks, win=WIN)
        cjumps = FH.get_composed_jumps(jumps, levels=LEVELS, win=WIN)
        # add them
        jumpcodes = map(lambda cj: FH.get_jumpcode_from_composed_jump(cj, maxwin=WIN), cjumps)
        CHT.add_jumpcodes(conn, tid, jumpcodes, normalize=NORMALIZE, commit=False)
        cnt_tid_added += 1
        if cnt_tid_added % 1000 == 0:
            conn.commit()
        # debug
        if cnt_tid_added % 500 == 0:
            print 'We added %d tid in the hash table(s) of %s.' % (cnt_tid_added,
                                                                   outputdb)
    # we index
    CHT.reindex(conn)
    # close connection
    conn.close()
    # done
    return

# error passing problems
class KeyboardInterruptError(Exception):
    pass

# for multiprocessing
def create_fill_one_partial_db_wrapper(args):
    """ wrapper function for multiprocessor, calls run_steps """
    try:
        create_fill_one_partial_db(**args)
    except KeyboardInterrupt:
        raise KeyboardInterruptError()

def die_with_usage():
    """ HELP MENU """
    print 'compute_hashcodes_mprocess.py'
    print '   by T. Bertin-Mahieux (2011) Columbia University'
    print ''
    print 'Code to extract all hashcodes from a give folder of'
    print 'data, e.g. the whole millin song dataset.'
    print 'Creates as many db as process, then aggregate them into one'
    print 'USAGE'
    print '   python compute_hashcodes_mprocess.py <maindir> <output.db> <nthreads>'
    sys.exit(0)

if __name__ == '__main__':

    # help menu
    if len(sys.argv) < 4:
        die_with_usage()

    # params
    maindir = sys.argv[1]
    outputdb = sys.argv[2]
    nthreads = int(sys.argv[3])

    # sanity checks
    if not os.path.isdir(maindir):
        print 'ERROR: %s is not a directory.' % maindir
        sys.exit(0)
    if os.path.exists(outputdb):
        print 'ERROR: %s already exists.' % outputdb
        sys.exit(0)
    if nthreads < 1 or nthreads > 8:
        print 'REALLY? %d processes?' % nthreads
        sys.exit(0)

    # FIRST PASS, ITERATE OVER ALL FILES
    # temp db list
    if nthreads == 1:
        tmpdbs = [outputdb]
    else:
        tmpdbs = map(lambda n: outputdb + '_tmp' + str(n) + '.db', range(nthreads))
        for tdb in tmpdbs:
            if os.path.exists(tdb):
                print 'ERROR: tmp db %s already exists.' % tdb
                sys.exit(0)
    # list all files
    allfiles = get_all_files(maindir)
    print 'We got %d h5 files.' % len(allfiles)
    files_per_thread = int(len(allfiles) * 1. / nthreads + 1.)
    print 'That gives us ~%d files per thread.' % files_per_thread
    # params
    params_list = []
    for k in range(nthreads):
        # params for one specific thread
        p = {'outputdb': copy.deepcopy(tmpdbs[k]),
             'filelist': copy.deepcopy(allfiles[files_per_thread * k:
                                                files_per_thread * (k + 1)])}
        params_list.append(p)
    # create pool, launch using the list of params
    # we underuse multiprocessing, we will have as many processes
    # as jobs needed
    pool = multiprocessing.Pool(processes=nthreads)
    try:
        pool.map(create_fill_one_partial_db_wrapper, params_list)
        pool.close()
        pool.join()
    except KeyboardInterrupt:
        print 'MULTIPROCESSING'
        print 'stopping multiprocessing due to a keyboard interrupt'
        pool.terminate()
        pool.join()
    except Exception, e:
        print 'MULTIPROCESSING'
        print 'got exception: %r, terminating the pool' % (e,)
        pool.terminate()
        pool.join()

    # SECOND PASS, AGGREGATE (ONE THREAD)
    if nthreads == 1:
        print 'We are done (there was one thread, no aggregation!)'
        sys.exit(0)
    # create final output
    import cover_hash_table as CHT
    conn = sqlite3.connect(outputdb)
    conn.execute('PRAGMA temp_store = MEMORY;')
    conn.execute('PRAGMA synchronous = OFF;')
    conn.execute('PRAGMA journal_mode = OFF;') # no ROLLBACK!
    conn.execute('PRAGMA page_size = 4096;')
    conn.execute('PRAGMA cache_size = 500000;') # page_size=4096, 500000->2GB
    CHT.init_db(conn)
    print 'Final db initialized (including PRAGMA settings)'

    # iterate over temporary dbs
    for tdb in tmpdbs:
        aggregate_dbs(conn, tdb)
    # index the final db
    CHT.reindex(conn)
    # all done
    conn.commit()
    conn.close()
    print 'ALL DONE! you should delete the temporary databases...'
    print tmpdbs


import numpy as np
import reinforced_epos.helpers.config as cf
import reinforced_epos.helpers.dataset as ds
from tensorboard.main import main

#import json
from subprocess import call
import os

def execute_iepos(dataset=None):
    max_iterations = cf.IEPOS_ITER
    children_node = cf.IEPOS_CHILDREN
    lamda = cf.IEPOS_LAMBDA
    seed = cf.IEPOS_SEED
    if(dataset is None):
        dataset = ds.get_dataset(normalize=False)
    print(np.shape(dataset))
    experiment_epos_folder = cf.get_experiement_folder()
    path_dataset = experiment_epos_folder + "numpy_dataset.npy"
    jar = os.path.abspath("../../bin/iepos.jar")
    call(["java", "-Xmx10g", "-jar", jar,  experiment_epos_folder, str(cf.IEPOS_CHILDREN), str(cf.IEPOS_ITER), str(cf.IEPOS_LAMBDA), str(cf.IEPOS_SEED), str(cf.IEPOS_ALGO)], cwd="../../bin/")

execute_iepos()